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
  );