// backward_ops
    %res_b1 = "tf.Identity"(%res_t) {_embedding_pipelining = "backward", _replication_info = "repl_info"} : (tensor<i32>) -> tensor<i32>
    %res_b2 = "tf.Identity"(%res_t) {_embedding_pipelining = "backward", _replication_info = "repl_info"} : (tensor<i32>) -> tensor<i32>
    %res_b3 = "tf.Identity"(%res_t) {_embedding_pipelining = "backward", _replication_info = "repl_info"} : (tensor<i32>) -> tensor<i32>

#define GEN_PASS_DEF_EMBEDDINGPIPELININGPASS
#include "tensorflow/compiler/mlir/tensorflow/transforms/sparsecore/sparsecore_passes.h.inc"

static constexpr char kEmbeddingPipelining[] = "_embedding_pipelining";
static constexpr char kEmbeddingPipeliningInlineAttr[] =
    "_embedding_pipelining_inline";
static constexpr char kEmbeddingForward[] = "forward";

    ],
    deps = [
        ":embedding_pipelining",
        ":embedding_program_key",
        ":embedding_sequencing",
        ":sparsecore_passes_inc_gen",
        "@llvm-project//llvm:Support",
        "@llvm-project//mlir:FuncDialect",
        "@llvm-project//mlir:IR",
        "@llvm-project//mlir:Pass",
    ],
)

cc_library(
    name = "embedding_pipelining",
    srcs = ["embedding_pipelining.cc"],

    // forward_ops
    %res_f = "tf.Const"() {_embedding_pipelining = "forward", _replication_info = "repl_info", value = dense<2> : tensor<i32>} : () -> tensor<i32>

    // core_tpu ops:
    %res_t = "tf.Identity"(%res_f) {_replication_info = "repl_info"} : (tensor<i32>) -> tensor<i32>

    // backward_ops

  // Disables full embedding pipelining when true. Instead, strict SparseCore
  // TensorCore sequencing will be used.
  bool tf_xla_disable_full_embedding_pipelining;

  // Force the WhileOps in embedding_pipelining and embedding_sequencing to use
  // this many parallel_iterations
  int tf_xla_embedding_parallel_iterations;
};

// Flags for common MLIR configurations.
struct MlirCommonFlags {

       Flag("tf_xla_embedding_parallel_iterations",
            &build_ops_flags->tf_xla_embedding_parallel_iterations,
            "If >0 then use this many parallel iterations in "
            "embedding_pipelining and embedding_sequency. By default, use the "
            "parallel_iterations on the original model WhileOp."),

       Flag("tf_xla_compile_on_demand", &device_flags->tf_xla_compile_on_demand,

#define GEN_PASS_DEF_EMBEDDINGSEQUENCINGPASS
#include "tensorflow/compiler/mlir/tensorflow/transforms/sparsecore/sparsecore_passes.h.inc"

static constexpr char kEmbeddingPipelining[] = "_embedding_pipelining";
static constexpr char kEmbeddingForward[] = "forward";
static constexpr char kEmbeddingBackward[] = "backward";
static constexpr char kDevice[] = "device";

  pm.addNestedPass<FuncOp>(
      mlir::TFTPU::CreateTPUReorderReplicateAndPartitionedInputsPass());
  pm.addNestedPass<FuncOp>(mlir::TF::CreateDecomposeReduceDatasetPass());
  // Only one of EmbeddingSequencing and EmbeddingPipelining will actually
  // run and the logic is in EmbeddingPipeliningPass. If the pipelining pass
  // runs, embedding attributes are stripped and the sequencing pass will have