%6 = "quantfork.dcast"(%5) : (tensor<3x4x2x!quant.uniform<i8:f32, 0.13170163023705575:-1>>) -> tensor<3x4x2xf32>
    %7 = "stablehlo.gather"(%6, %arg2) {
    dimension_numbers = #stablehlo.gather<
      offset_dims = [2, 3],
      collapsed_slice_dims = [0],
      start_index_map = [1, 0],
      index_vector_dim = 2>,
    slice_sizes = array<i64: 1, 2, 2>,

  }

  Value dot_result =
      builder
          .create<TF::XlaDotV2Op>(
              loc, /*output=*/output.getType(),
              /*lhs=*/input,
              /*rhs=*/weight,
              /*dimension_numbers=*/
              builder.getStringAttr(dnums.SerializeAsString()),
              /*precision_config=*/builder.getStringAttr(precision_config_str))
          .getResult();

  func.func private @composite_gather_fn(%arg0: tensor<3x4x2xf32>, %arg1: tensor<2x3x2xi32>) -> tensor<2x3x2x2xf32> attributes {_from_xla_call_module} {
    %0 = "stablehlo.gather"(%arg0, %arg1) {
      dimension_numbers = #stablehlo.gather<
        offset_dims = [2, 3],
        collapsed_slice_dims = [0],
        start_index_map = [1, 0],
        index_vector_dim = 2>,
      slice_sizes = array<i64: 1, 2, 2>,

    // dimension, assuming the filter format is `[0, 1, i, o]`.
    // TODO: b/291029962 - Lift the assumption above and retrieve the
    // quantization dimension from the `dimension_numbers` attribute.
    UniformQuantizedPerAxisType filter_quantized_element_type =
        CreateI8F32UniformQuantizedPerAxisType(
            filter_op->getLoc(), *rewriter.getContext(), filter_scale_values,

func.func @gather_fn() -> tensor<2x3x2x2xi32> {
  %0 = stablehlo.constant dense<1> : tensor<3x4x2xi32>
  %1 = stablehlo.constant dense<1> : tensor<2x3x2xi64>
  %2 = "stablehlo.gather"(%0, %1) {
  dimension_numbers = #stablehlo.gather<
    offset_dims = [2, 3],
    collapsed_slice_dims = [0],
    start_index_map = [1, 0],
    index_vector_dim = 2>,
  slice_sizes = array<i64: 1, 2, 2>,
  indices_are_sorted = false

  kernel_spatial_dimensions.resize(num_spatial_dims);
  std::iota(kernel_spatial_dimensions.begin(), kernel_spatial_dimensions.end(),
            0);

  return builder->getNamedAttr(
      "dimension_numbers",
      ConvDimensionNumbersAttr::get(
          builder->getContext(), batch_dim, feature_dim, spatial_dims,
          kernel_input_feature_dim, kernel_output_feature_dim,

  %strides = "tf.Const"() {value = dense<[3, 1, 1]> : tensor<3xi32>} : () -> tensor<3xi32>

    DefaultValuedOptionalAttr<I64Attr, "1">:$batch_group_count,
    DefaultValuedOptionalAttr<I64Attr, "1">:$feature_group_count,
    DefaultValuedOptionalAttr<StrAttr, "\"\"">:$dimension_numbers,
    DefaultValuedOptionalAttr<I64Attr, "-1">:$lhs_quantization_axis,
    I64Attr:$lhs_quantization_min_val,
    I64Attr:$lhs_quantization_max_val,