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

// Rewrites RecvTPUEmbeddingActivationsOp and SendTPUEmbeddingGradients ops to
// internal variants by introducing XlaRecvTPUEmbeddingDeduplicationData op.
struct RewriteTPUEmbeddingOps
    : public impl::RewriteTPUEmbeddingOpsPassBase<RewriteTPUEmbeddingOps> {
  void runOnOperation() override;
};

embedding table in the model.}]>:$outputs
  );

  TF_DerivedResultSizeAttr num_tables = TF_DerivedResultSizeAttr<0>;
}

def TF_XlaRecvTPUEmbeddingDeduplicationDataOp : TF_Op<"XlaRecvTPUEmbeddingDeduplicationData", []> {
  let summary = [{
Receives deduplication data (indices and weights) from the embedding core.
  }];

  let description = [{