spec->biases_params[2] = {{0, 1},
                                  quant::GetUniformQuantizedTypeForBias};
      }
    } else if (function_name.contains("gather")) {
      // Note that gather has axis attribute that specifies channel axis.
      spec->coeff_op_quant_dim[0] = -1;
    }
    for (auto quantizable_operand : spec->coeff_op_quant_dim) {
      spec->quantizable_operands.insert(quantizable_operand.first);
    }
  }

	Score int64
}

// BuildBrokenLinearFunction creates a function which is built using linear segments. Segments are defined via shape array.
// Shape[i].Utilization slice represents points on "Utilization" axis where different segments meet.
// Shape[i].Score represents function values at meeting points.
//
// function f(p) is defined as:
//
//	shape[0].Score for p < shape[0].Utilization

  Type packed_type = RankedTensorType::get(
      {bcast.output_batch_size(), rows, cols}, element_type);
  const auto axis = rewriter.getI64IntegerAttr(0);
  auto pack_op =
      rewriter.create<TF::PackOp>(loc, packed_type, /*values=*/matmuls, axis);

  // Reshape the rank-3 tensor into the correct output shape.
  const auto& result_batch_shape = bcast.output_batch_shape().dim_sizes();

  if (!IsZero(scope, grad_grad)) {
    std::vector<int> axis;
    auto logits_softmax = Softmax(scope, logits);

    auto grad_grad_expand = ExpandDims(scope, grad_grad, 1);
    auto logits_softmax_expand = ExpandDims(scope, logits_softmax, 2);
    auto matmul_result =
        BatchMatMul(scope, grad_grad_expand, logits_softmax_expand);
    axis.push_back(1);

              {static_cast<int>(
                  dot_dimensions_info.batch_dimensions().AxesArray().size())},
              builder.getIntegerType(32)),
          operand_shape, batch_axes_tensor, /*axis*/ 0, /*batch_dims*/ 0);
      flattend_shape_values.push_back(batch_dims);
    } else {
      llvm::SmallVector<int32_t> batch_i32_vec;
      for (int64_t element :

// -----

// CHECK-LABEL: @RemoveRedundantUnpackPack
func.func @RemoveRedundantUnpackPack(%arg0: tensor<2x5xf32>) -> tensor<2x5xf32> {
  %0:2 = "tfl.unpack"(%arg0) {axis = 0 : i32, num = 2 : i32} : (tensor<2x5xf32>) -> (tensor<5xf32>, tensor<5xf32>)
  %1 = "tfl.pack"(%0#0, %0#1) {axis = 0 : i32, values_count = 2 : i32} : (tensor<5xf32>, tensor<5xf32>) -> (tensor<2x5xf32>)
  func.return %1: tensor<2x5xf32>
  // CHECK-NOT: pack

                         %weight : tensor<*xi8>, %input : tensor<*xi32>, %axis : tensor<i32>,
                         %weight_scale : tensor<*xf32>, %weight_zp : tensor<*xi32>) -> tensor<*xf32>
      attributes {tf_quant.quantized_ops = ${quantized_ops}}
  {
    %accum_out = "tf.GatherV2"(%weight, %input, %axis) {

  func.return
}

// CHECK-LABEL: testGatherToV2
// Ensures that axis param and batch_dims attr use their default values of 0.
func.func @testGatherToV2(%params: tensor<4x3xf32>, %indices: tensor<1x2xi32>) -> tensor<2x3xf32> {
  // CHECK: %[[AXIS:.*]] = "tf.Const"() <{value = dense<0> : tensor<i32>}> : () -> tensor<i32>

  }

  func.func private @func_2_CPU_FLOAT(%arg0: tensor<1xf32>, %arg1: tensor<1xf32>) -> tensor<2x1xf32> attributes {tac.device = "CPU", tac.inference_type = "FLOAT", tac.interface_name = "func_2"} {
    %0 = "tfl.pack"(%arg0, %arg1) {axis = 0 : i32, tac.device = "CPU", tac.inference_type = "FLOAT", values_count = 2 : i32} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2x1xf32>
    func.return %0 : tensor<2x1xf32>
  }

// CHECK:     %17 = "tfl.gather"(%arg13, %14) <{axis = 0 : i32, batch_dims = 0 : i32}> {tac.device = "DARWINN", tac.inference_type = "FLOAT"} : (tensor<5xi32>, tensor<?xi32>) -> tensor<?xi32>
// CHECK:     %18 = tfl.add %arg14, %17 {fused_activation_function = "NONE", tac.device = "DARWINN", tac.inference_type = "FLOAT"} : tensor<?xi32>