.Device(DEVICE),                                 \
                          XlaCompileOnDemandOp);                               \
  REGISTER_KERNEL_BUILDER(Name("XlaSpmdShardToFullShape").Device(DEVICE),      \
                          XlaCompileOnDemandOp);                               \
  REGISTER_KERNEL_BUILDER(Name("XlaSharding").Device(DEVICE),                  \

          %outputs_7, %control_8 = tf_executor.island wraps "tf.XlaSpmdShardToFullShape"(%outputs_5) {dim = -1 : i64, full_shape = #tf_type.shape<2x2>, manual_sharding = "\08\03\1A\02\02\01\22\02\00\01", unspecified_dims = []} : (tensor<1x2xf32>) -> tensor<2x2xf32>

                                         "XlaSharding",
                                         "XlaSort",
                                         "XlaSpmdFullToShardShape",
                                         "XlaSpmdShardToFullShape",
                                         "XlaSvd",
                                         "XlaVariadicReduceV2",
                                         "XlaVariadicSort",

    // CHECK: %[[FULL:.*]] = mhlo.custom_call @SPMDShardToFullShape(%[[SHARDING]]) {backend_config = "", mhlo.sharding = "{devices=[2,1]0,1}"} : (tensor<1x2xi64>) -> tensor<2x2xi64>
    // CHECK: return %[[FULL]]
    %0 = "tf.XlaSpmdShardToFullShape"(%arg0) {dim = -1 : i64, full_shape = #tf_type.shape<2x2>, manual_sharding = "\08\03\1A\02\02\01\22\02\00\01", unspecified_dims = []} : (tensor<1x2xi64>) -> tensor<2x2xi64>
    func.return %0 : tensor<2x2xi64>
  }

    // CHECK-SAME:          recv_key = "host_compute_channel_0_retvals"
    // CHECK-SAME:          send_key = "host_compute_channel_0_args"
    // CHECK:             %[[B_FULL:.+]] = "tf.XlaSpmdShardToFullShape"(%[[B]]) <{dim = -1 : i64, full_shape = #tf_type.shape<2x2>, manual_sharding = "\08\03\1A\02\02\01\22\02\00\01", unspecified_dims = []}> : (tensor<1x2xi64>) -> tensor<2x2xi64>
    // CHECK:             "tf.OpC"(%[[B_FULL]])

      "XlaSend",
      "XlaSetBound",
      "XlaSetDynamicDimensionSize",
      "XlaSharding",
      "XlaSort",
      "XlaSplitND",
      "XlaSpmdFullToShardShape",
      "XlaSpmdShardToFullShape",
      "XlaSvd",
      "XlaVariadicReduce",
      "XlaVariadicReduceV2",
      "XlaVariadicSort",
      "XlaWhile",
      "Zeta",
      "_Arg",
      "_ArrayToList",

  );

  let results = (outs
    TF_Tensor:$output
  );

  TF_DerivedOperandTypeAttr T = TF_DerivedOperandTypeAttr<0>;
}

def TF_XlaSpmdShardToFullShapeOp : TF_Op<"XlaSpmdShardToFullShape", [Pure]> {
  let summary = [{
An op used by XLA SPMD partitioner to switch from manual partitioning to
  }];

  let description = [{