Location loc = op.getLoc();
    Value scalar_zero = CreateI32SplatConst(loc, &rewriter, {}, 0);

    // Compute the input tensorlist's length and store it in `input_size`.
    IntegerType shape_dtype = rewriter.getIntegerType(32);
    auto input_size = rewriter.create<TF::TensorListLengthOp>(
        loc, tensorflow::GetTypeFromTFTensorShape({}, shape_dtype),
        op.getOperand(0));

        mlir::dyn_cast<RankedTensorType>(input.getType());
    if (!input_type) {
      return failure();
    }
    const ArrayRef<int64_t> input_shape = input_type.getShape();

    const int input_size = input_shape.size();

    RankedTensorType begin_type = mlir::cast<RankedTensorType>(begin.getType());
    const ArrayRef<int64_t> begin_shape = begin_type.getShape();
    const int begin_dim = begin_shape.size();

    TF_Float32Tensor:$sorted_gains,
    TF_Float32Tensor:$embedding_table,
    TF_Int32Tensor:$num_minibatches_per_physical_sparse_core,

    ConfinedAttr<I64Attr, [IntMinValue<0>]>:$input_size,
    OptionalAttr<F32Attr>:$quantization_config_low,
    OptionalAttr<F32Attr>:$quantization_config_high,
    OptionalAttr<I64Attr>:$quantization_config_num_buckets,
    StrAttr:$table_name
  );

    std::set<EdgeSpec> a_edges;  // actual edges
    // Get edges from inputs to body nodes and between body nodes
    for (const NodeDef& node_def : fdef.node_def()) {
      for (int i = 0; i < node_def.input_size(); ++i) {
        const string& in = node_def.input(i);
        const auto& v =
            a_edges.insert({in, strings::StrCat(node_def.name(), ":", i)});
        ASSERT_TRUE(v.second) << "Duplicate edge " << in << " -> "