public:
  explicit ConvertEmbeddedLookupFunc(func::FuncOp func) : func_(func) {}

  void RewriteFunc() {
    func_->setAttr(kTFImplements,
                   StringAttr::get(func_.getContext(), "embedding_lookup"));
    Value lookup = func_.getArgument(1);
    Value value = func_.getArgument(0);
    auto output_type = func_.getFunctionType().getResult(0);

    OpBuilder builder(func_.getBody());

  %value1 = "tfl.pseudo_const"() {value = dense<[[[1.0, 2.0], [3.0, 4.0]], [[5.0, 6.0], [7.0, 8.0]]]> : tensor<2x2x2xf32>} : () -> tensor<2x2x2xf32>
  %lookup0 = "tfl.embedding_lookup"(%index0, %value0) : (tensor<5xi32>, tensor<3x2xf32>) -> tensor<5x2xf32>
  %lookup1 = "tfl.embedding_lookup"(%index1, %value1) : (tensor<3xi32>, tensor<2x2x2xf32>) -> tensor<3x2x2xf32>
  func.return %lookup0, %lookup1 : tensor<5x2xf32>, tensor<3x2x2xf32>

// LINT.IfChange
enum BuiltinOperator : byte {
  ADD = 0,
  AVERAGE_POOL_2D = 1,
  CONCATENATION = 2,
  CONV_2D = 3,
  DEPTHWISE_CONV_2D = 4,
  DEPTH_TO_SPACE = 5,
  DEQUANTIZE = 6,
  EMBEDDING_LOOKUP = 7,
  FLOOR = 8,
  FULLY_CONNECTED = 9,
  HASHTABLE_LOOKUP = 10,
  L2_NORMALIZATION = 11,
  L2_POOL_2D = 12,
  LOCAL_RESPONSE_NORMALIZATION = 13,
  LOGISTIC = 14,
  LSH_PROJECTION = 15,

  // CHECK: %[[RES:.*]] = "tfl.embedding_lookup"(%arg0, %[[CST]]) : (tensor<2xi32>, tensor<3x5xf32>) -> tensor<2x5xf32>
  // CHECK: return %[[RES]] : tensor<2x5xf32>
}

// CHECK-LABEL: dontReplaceOneHotFullyConnectedWithLookupBadIndexType

// LINT.IfChange
enum BuiltinOperator : int32 {
  ADD = 0,
  AVERAGE_POOL_2D = 1,
  CONCATENATION = 2,
  CONV_2D = 3,
  DEPTHWISE_CONV_2D = 4,
  DEPTH_TO_SPACE = 5,
  DEQUANTIZE = 6,
  EMBEDDING_LOOKUP = 7,
  FLOOR = 8,
  FULLY_CONNECTED = 9,
  HASHTABLE_LOOKUP = 10,
  L2_NORMALIZATION = 11,
  L2_POOL_2D = 12,
  LOCAL_RESPONSE_NORMALIZATION = 13,
  LOGISTIC = 14,
  LSH_PROJECTION = 15,

// LINT.IfChange
enum BuiltinOperator : int32 {
  ADD = 0,
  AVERAGE_POOL_2D = 1,
  CONCATENATION = 2,
  CONV_2D = 3,
  DEPTHWISE_CONV_2D = 4,
  DEPTH_TO_SPACE = 5,
  DEQUANTIZE = 6,
  EMBEDDING_LOOKUP = 7,
  FLOOR = 8,
  FULLY_CONNECTED = 9,
  HASHTABLE_LOOKUP = 10,
  L2_NORMALIZATION = 11,
  L2_POOL_2D = 12,
  LOCAL_RESPONSE_NORMALIZATION = 13,
  LOGISTIC = 14,
  LSH_PROJECTION = 15,

// CHECK-SAME:              [[VAL_0:%.*]]: tensor<*xf32>, [[VAL_1:%.*]]: tensor<*xi32>) -> tensor<*xf32>

// CHECK-LABEL:   attributes  {tf._implements = "embedding_lookup", tf._reference = "mlir"} {
// CHECK:           [[VAL_2:%.*]] = "tfl.embedding_lookup"([[VAL_1]], [[VAL_0]]) : (tensor<*xi32>, tensor<*xf32>) -> tensor<*xf32>
// CHECK:           return [[VAL_2]] : tensor<*xf32>

// CHECK-LABEL:   func @lstmcellsimple(

embedding_indices and aggregation_weights into rows.
It corresponds to ids.row_splits in embedding_lookup(), when ids is a
RaggedTensor.}]>:$sample_splits,
    Arg<Variadic<TF_I32OrI64Tensor>, [{A list of rank 1 Tensors, indices into the embedding tables.
It corresponds to ids.values in embedding_lookup(), when ids is a RaggedTensor.}]>:$embedding_indices,

  }];

  let arguments = (ins TFL_TensorOf<[F32, I8]>:$x);

  let results = (outs TFL_TensorOf<[F32, I8]>:$y);

  let hasOptions = 0;
}

def TFL_EmbeddingLookupOp: TFL_Op<"embedding_lookup",
    [Pure,
     TFL_OperandHasRank<0, 1>,
     TFL_OperandHasRankAtLeast<1, 2>,
     DynamicRangeQuantizedOpInterface,
     QuantizableResult
    ]> {
  let summary = "Embedding lookup operator";

  func.return %2 : tensor<10x!quant.uniform<u8:f32, 4.0>>
}

// -----

func.func @testEmbeddingLookup(%arg0 : tensor<?xi32>, %arg1 : tensor<?x?xf32>) -> tensor<?xf32> {
  %0 = "tfl.embedding_lookup"(%arg0, %arg1) : (tensor<?xi32>,tensor<?x?xf32>) -> tensor<?xf32>
  func.return %0 : tensor<?xf32>
}

// -----