// argument index.
llvm::SmallVector<llvm::SmallVector<int64_t, 4>, 4> GetMetadataArgumentMapping(
    const tpu::TPUCompileMetadataProto& metadata);

// Gets the proper tensor dimension from XLA OpSharding.
// "replicate_on_last_tile_dim" and "last_tile_dims" should be deducted from the
// real Tensor dimensions when tiled.
// For example:
// f32[8,512](sharding={devices=[1,1,2]0,1 last_tile_dims={REPLICATED})

  for (auto split_op : split_ops_for_tiled_input) {
    for (auto split_op_output : split_op.getResults()) {
      int64_t repeat_count =
          input_sharding.replicate_on_last_tile_dim()
              ? *input_sharding.tile_assignment_dimensions().rbegin()
              : 1;
      for (int64_t i = 0; i < repeat_count; ++i) {
        tiled_inputs->push_back(split_op_output);
      }
    }

// -----

// Tests TPIv2 with a "partially tiled" XLA annotation where:
//   type: OTHER
//   tile_assignment_dimensions: [4, 1, 1, 1, 2]
//   tile_assignment_devices: [0, 1, 2, 3, 4, 5, 6, 7]
//   replicate_on_last_tile_dim: true
// Serialized string:
//   "\08\03\1A\05\04\01\01\01\02\22\08\00\01\02\03\04\05\06\070\01"

// CHECK-LABEL: func @partial_tile_partitioned_variable

    // have one or more dimension(s) than its rank; so, we subtract them to
    // determine which rank the sharding is compatible with.
    tile_assignment_rank -= (int)sharding->replicate_on_last_tile_dim();
    tile_assignment_rank -= sharding->last_tile_dims_size();

    if (tensor_rank < tile_assignment_rank) {
      if (partitioned_op) {
        partitioned_op->emitError()

    func.return %4, %3 : tensor<*xi32>, tensor<*xi1>
  }
}

// -----

// Tests inputs to TPUComputation that are tiled in multiple dimensions with
// replicate_on_last_tile_dim set.

// The following OpSharding is used for TPU computation inputs in below test:
// Proto debug string:
//  input 0
//   type: OTHER
//   tile_assignment_dimensions: 2