Support multiple StableHLO Composite outputs

[chunnienc]

This PR adds the support for marking multiple tensors in StableHLOCompositeBuilder.mark_outputs.

Example:

class Model(torch.nn.Module):

      def __init__(self):
        super().__init__()

      def forward(self, x, y):
        builder = StableHLOCompositeBuilder("sample_composite")
        x, y = builder.mark_inputs(x, y)
        a = x + y
        b = x - y
        c = x + 1
        a, b, c = builder.mark_outputs(a, b, c)
        return a + b + c

SHLO:

module @IrToHlo.42 attributes {mhlo.cross_program_prefetches = [], mhlo.is_dynamic = false, mhlo.use_auto_spmd_partitioning = false} {
  func.func @main(%arg0: tensor<5x5xf32>, %arg1: tensor<5x5xf32>) -> tensor<5x5xf32> {
    %0:3 = stablehlo.custom_call @stablehlo.composite(%arg0, %arg1) {called_computations = [@sample_composite.impl], composite.backend_config = {attributes = {}, name = "sample_composite"}} : (tensor<5x5xf32>, tensor<5x5xf32>) -> (tensor<5x5xf32>, tensor<5x5xf32>, tensor<5x5xf32>)
    %1 = stablehlo.add %0#0, %0#1 : tensor<5x5xf32>
    %2 = stablehlo.add %1, %0#2 : tensor<5x5xf32>
    return %2 : tensor<5x5xf32>
  }
  func.func private @sample_composite.impl(%arg0: tensor<5x5xf32>, %arg1: tensor<5x5xf32>) -> (tensor<5x5xf32>, tensor<5x5xf32>, tensor<5x5xf32>) {
    %0 = stablehlo.constant dense<1.000000e+00> : tensor<5x5xf32>
    %1 = stablehlo.add %arg0, %arg1 : tensor<5x5xf32>
    %2 = stablehlo.subtract %arg0, %arg1 : tensor<5x5xf32>
    %3 = stablehlo.add %arg0, %0 : tensor<5x5xf32>
    return %1, %2, %3 : tensor<5x5xf32>, tensor<5x5xf32>, tensor<5x5xf32>
  }
}

[miladm]

Suggesting to include references to issues on public URLs on GH. Thanks very much.

[lsy323]

LGTM, left a few comments/questions. Thanks @chunnienc!