%outputs_2, %control_3 = tf_executor.island wraps "tf.XlaSpmdFullToShardShape"(%outputs_0) {dim = -1 : i64, manual_sharding = "\08\03\1A\02\02\01\22\02\00\01", unspecified_dims = []} : (tensor<2x2xf32>) -> tensor<1x2xf32>

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

    // CHECK: %[[MANUAL:.*]] = mhlo.custom_call @SPMDFullToShardShape(%[[SHARDING]]) {backend_config = "", mhlo.sharding = "{manual}"} : (tensor<2x2xi64>) -> tensor<1x2xi64>
    // CHECK: return %[[MANUAL]]
    %0 = "tf.XlaSpmdFullToShardShape"(%arg0) {dim = -1 : i64, manual_sharding = "\08\03\1A\02\02\01\22\02\00\01", unspecified_dims = []} : (tensor<2x2xi64>) -> tensor<1x2xi64>
    func.return %0 : tensor<1x2xi64>
  }

    // CHECK:         }, {
    // CHECK:           "tf_device.cluster"
    // CHECK:             %[[A:.+]] = "tf.OpA"
    // CHECK:             %[[A_SHARD:.+]] = "tf.XlaSpmdFullToShardShape"(%[[A]]) <{dim = -1 : i64, manual_sharding = "\08\03\1A\02\02\01\22\02\00\01", unspecified_dims = []}> : (tensor<2x2xi64>) -> tensor<1x2xi64>
    // CHECK:             %[[B:.+]] = "tf._XlaHostComputeMlir"(%[[A_SHARD]])

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

  let results = (outs
    Res<TF_Tensor, [{A `Tensor` of type T.}]>:$output
  );

  TF_DerivedOperandTypeAttr T = TF_DerivedOperandTypeAttr<0>;
}

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

  let description = [{