message: Register TPUPartitionedInput in TF MLIR ODS.
diff: @@ -2042,5 +2042,24 @@ argument "ref" in addition to returning the results.
  TF_DerivedOperandTypeAttr T = TF_DerivedOperandTypeAttr<0>;
  }
  +def TF_TPUPartitionedInputOp : TF_Op<"TPUPartitionedInput", [NoSideEffect]> {
  + let summary = [{
  +An op that groups a list of partitioned inputs together. This op
  + }];
  +
  + let arguments = (ins
  + Variadic<TF_Tensor>:$inputs,
  +
  + DefaultValuedAttr<I64Attr, "0">:$partition_dim,
  + OptionalAttr<StrAttr>:$_XlaSharding
  + );
  +
  + let results = (outs
  + TF_Tensor:$output
  + );
  +
  + TF_DerivedOperandTypeAttr T = TF_DerivedOperandTypeAttr<0>;
  + TF_DerivedOperandSizeAttr N = TF_DerivedOperandSizeAttr<0>;
  +}
  #endif // TF_OPS
  