// RUN: tf-opt %s -tf-replicated-clustering-bridge-v2 -tfrt-lower-cluster-to-runtime-ops-tpu 2>&1 | FileCheck %s

// TPUReshardVariables should be inserted even when While functions' shapes are
// different than While operand shapes. Test the whole tf-tpu-bridge because
// correct insertion of TPUReshardVariables depends on multiple passes including
// TPUVariableRuntimeReformatting, ShapeInference, WhileRegion canonicalization,
// and TPUMergeVariablesWithExecute.

            // CHECK: "tf_device.launch"
            // CHECK-SAME: device = "TPU_REPLICATED_CORE_0"
            // CHECK-NEXT: "tf.TPUReshardVariables"(%[[ID]], %[[R1]], %[[COMPILE]]#1, %[[R_STATE]])
            // CHECK-NEXT: tf_device.return
            // CHECK: "tf.TPUExecuteAndUpdateVariables"(%[[ID]], %[[R1]], %[[COMPILE]]#1)
            "tf_device.launch"() ({

  });
  if (reshard_was_inserted)
    VLOG(1) << "tf-tpu-variable-runtime-reformatting inserted at least one "
               "TPUReshardVariables";
  else
    VLOG(1) << "tf-tpu-variable-runtime-reformatting inserted no "
               "TPUReshardVariables";
}

}  // namespace

std::unique_ptr<OperationPass<ModuleOp>>
CreateTPUVariableRuntimeReformattingPass() {

            "Enables TPUVariableRuntimeReformatting pass for MLIR-Based "
            "TensorFlow Compiler Bridge. This enables weight update sharding "
            "and creates TPUReshardVariables ops.")});

  AppendMarkForCompilationPassFlagsInternal(flag_list);
  xla::ParseFlagsFromEnvAndDieIfUnknown("TF_XLA_FLAGS", *flag_list);

  mlir_flags = new MlirCommonFlags;

  );

  TF_DerivedOperandTypeAttr T = TF_DerivedOperandTypeAttr<0>;
  TF_DerivedResultSizeAttr num_replicas = TF_DerivedResultSizeAttr<0>;
}

def TF_TPUReshardVariablesOp : TF_Op<"TPUReshardVariables", []> {
  let summary = "Op that reshards on-device TPU variables to specified state.";

  let description = [{
Op that reshards on-device TPU variables to specified state. Internal use only.