input_index, tiled_input_size, num_cores_per_replica));
    };

    // If input is already partitioned using the `tf.TPUPartitionedInputV2` op,
    // only replicated sharding is supported where i-th operand to
    // `tf.TPUPartitionedInputV2` op is input to the i-th logical device.
    if (auto partitioned_input =
            llvm::dyn_cast_or_null<mlir::TF::TPUPartitionedInputV2Op>(

// attributes back to legacy attributes.
std::unique_ptr<OperationPass<func::FuncOp>>
CreateConvertToLegacyCompileAndReplicateAttributesPass();

// Creates a pass that converts all TPUPartitionedInput to TPUPartitionedInputV2
std::unique_ptr<OperationPass<func::FuncOp>>
CreateTPUPartitionedOpConversionPass();

std::unique_ptr<OperationPass<ModuleOp>> CreateTPUValidateInputsPass();

  let arguments = (ins
    I64Attr:$logical_core
  );

  let results = (outs
    TF_Int64Tensor:$device_ordinal
  );
}

// We must manually define TPUPartitionedInput, TPUPartitionedInputV2,
// TPUPartitionedOutput, and TPUPartitionedOutputV2 since they have an
// optional attribute, _XlaSharding, unlike the TensorFlow definition.
def TF_TPUPartitionedInputOp : TF_Op<"TPUPartitionedInput", [Pure]> {

  func.return %0 : tensor<1xi1>
}

// -----

// Test invalid tf.TPUPartitionedInputV2 with packing
func.func @testPackedTPUPartitionedInputV2(tensor<2x4xf32>, tensor<2x4xf32>) -> tensor<4x4xf32> {
^bb0(%arg0: tensor<2x4xf32>, %arg1: tensor<2x4xf32>):
  // expected-error @+1 {{expected 1 inputs, got 2}}

      tensorflow::GetTypeFromTFTensorShape({}, IntegerType::get(context, 1)));
  return success();
}

//===----------------------------------------------------------------------===//
// TPUPartitionedInputV2
//===----------------------------------------------------------------------===//

// This method mimics this op's core/TF-level shape inference logic
LogicalResult TPUPartitionedInputV2Op::verify() {