break;
      }
    }
    TF::TPUReplicatedInputOp input =
        llvm::cast<TF::TPUReplicatedInputOp>(input_ptr);

    // Validate that all users of the TPUReplicatedInput are ReadVariable
    // or AssignVariable ops and check if any are outside the forwards pass.
    bool duplicate_needed = false;
    for (Operation* next_user : input.getOutput().getUsers()) {

  // CHECK-SAME: tensor<*xi32>
  func.return %1 : tensor<*xi32>
}

// If multiple types are used across reads and writes, check that the read uses
// the most refined type. The first ReadVariable should refine the type from
// *xi32 to ?xi32 and the assign should refine it further to 4xi32.
// CHECK-LABEL: func @type_refinement_use_refined_type
func.func @type_refinement_use_refined_type() -> tensor<4xi32> {