inline constexpr llvm::StringRef kReplicationInfoAttr = "_replication_info";
// Marks a node for XLA-TPU compilation. The attribute value indicates the
// associated compilation cluster and replication metadata op.
inline constexpr llvm::StringRef kTpuReplicateAttr = "_tpu_replicate";
// Device types.
inline constexpr llvm::StringRef kTpuDevice = "TPU";
// _xla_outside_compilation
inline constexpr llvm::StringRef kXlaOutsideCompilationAttr =

    });
    if (replicate)
      reshard_was_inserted |= HandleReplicateOp(while_op, replicate);
  });
  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>>

                        std::vector<int64_t>(),
                        "number of tasks from available TPU devices must be "
                        "'num_tasks' in 'topology' (1), got 2"),
        std::make_tuple(2, 1, TopologyWithMeshShapeAndTasks({1, 1, 1, 1}, 2, 2),
                        std::vector<int64_t>(),
                        "number of TPU devices available per task must be "

  void runOnFunction(
      func::FuncOp func,
      const TF::ResourceAliasAnalysis::Info& resource_alias_analysis);

  StringRef getArgument() const final { return "tf-tpu-dynamic-layout-pass"; }

  StringRef getDescription() const final {
    return "Inserts TPU layout ops to determine layout at run time.";
  }
};

// Checks if the input producer op is supported in this transform. Right now, we

    return absl::OkStatus();
  }

  // TFRT-TPU is used if device type is `DEVICE_TPU` and platform_info does not
  // have `xla_device_metadata`. This is used for TFRT-TPU when
  // BuildXlaDeviceCompiler() is called in GetCompilerIr(). Currently only
  // lowering to HLO is needed there and xla::LocalClient doesn't support
  // building the executable for TFRT-TPU and hence, is set to nullptr here.

  let summary = "Op that compiles a computation in MLIR into a TPU program, and loads and executes it on a TPU device.";

  let description = [{
For the internal use of the TPU compiler.

'static_shapes' are tensors specifying the maximum dimension sizes for the tensors specified in `dynamic_operands`.
'args' are inputs to the TPU computation.

  // CHECK: %[[IDENTIFY:.*]] = "tf.Identity"(%[[SUBGRAPH_0]]#1) {device = ""} : (tensor<1024x3xf32>) -> tensor<1024x3xf32>
  // CHECK: %[[SUBGRAPH_1:.*]] = "tf.XlaCallModule"() <{Sout = [#tf_type.shape<1024x3>], {{.*}} ["CPU", "TPU"], {{.*}}}> {_entry_function = @_stablehlo_main_1

                          WeightOnlyPtqComponent::kName, *function_aliases,
                          *ctx, *module));

  // Remove the `tpu` tag for exporting because the output quantized model is
  // essentially a CPU model.
  tags.erase("tpu");

  py_function_library.SaveExportedModel(
      dst_saved_model_path, post_calibrated_exported_model,
      src_saved_model_path, tags, signature_def_map);

  return builder->create<mlir::TF::ConcatOp>(
      location, output_type, concat_dimension_op.getOutput(), inputs);
}

// For tile sharded inputs to TPU computation, inject split op between the
// input values and TPU computation so that tiled input values are passed in
// as inputs to TPU computations. If more than one dimension is sharded, then
// a tree of connected split ops are added before tf_device.parallel_execute op.

                          PostCalibrationComponent::kName, *function_aliases,
                          *ctx, *module));

  // Remove the `tpu` tag for exporting because the output quantized model is
  // essentially a CPU model.
  tags.erase("tpu");

  py_function_library.SaveExportedModel(
      dst_saved_model_path, post_calibrated_exported_model,
      src_saved_model_path, tags, signature_def_map);