#define GEN_PASS_DEF_TPUUPDATEEMBEDDINGENQUEUEOPINPUTSPASS
#include "tensorflow/compiler/mlir/tensorflow/transforms/tf_passes.h.inc"

struct TPUUpdateEmbeddingEnqueueOpInputsPass
    : public impl::TPUUpdateEmbeddingEnqueueOpInputsPassBase<
          TPUUpdateEmbeddingEnqueueOpInputsPass> {
  void runOnOperation() override;
};

// Extracts `_tpu_embedding_layer` attribute from TPU embedding ops and

    at the head of TPU computation if it's only used by outside compiled ops.
  }];

  let constructor = "TFTPU::CreateTPUHostComputationExpansionPass()";
}

def TPUUpdateEmbeddingEnqueueOpInputsPass : Pass<"tf-tpu-update-embedding-enqueue-op-inputs", "mlir::func::FuncOp"> {
  let summary = "Updates inputs to TPU embedding enqueue ops depending on whether graph is in training mode or in evaluation mode.";