static bool AreInputDimensionsOneInAxes(Value input,
                                        const mlir::Attribute &axes) {
  RankedTensorType input_type =
      mlir::dyn_cast_or_null<RankedTensorType>(input.getType());
  if (!input_type) return false;
  auto type_shape = input_type.getShape();

  DenseIntElementsAttr axes_attr =
      mlir::dyn_cast_or_null<DenseIntElementsAttr>(axes);
  if (!axes_attr) return false;

class Operation {
 public:
  Operation() : node_(nullptr) {}
  explicit Operation(Node* n);

  int32 num_inputs() const { return node_->num_inputs(); }
  DataType input_type(int32_t o) const { return node_->input_type(o); }
  Output input(int32_t i) const;

  int32 num_outputs() const { return node_->num_outputs(); }
  DataType output_type(int32_t o) const { return node_->output_type(o); }

  return attr_enforced_quantizable || trait_enforced_quantizable;
}

// Returns the quantized type for the
// input_type/min/max/storag_type_width/narrow_range.
// This is entry point to the Quant dialect and used for both quantizing
// activations and weights.
Type GetQuantizedType(Builder builder, const Type input_type,
                      const ArrayRef<double> min, const ArrayRef<double> max,

    auto qbarrier = rewriter.create<QuantizeCastOp>(op.getLoc(), quantizedType,
                                                    op.getInputs());
    rewriter.replaceOpWithNewOp<DequantizeCastOp>(op, converter.input_type,
                                                  qbarrier.getResult());

    return false;
  }
};

class ConstFakeQuantRewrite
    : public FakeQuantRewrite<ConstFakeQuantRewrite, ConstFakeQuant> {

  const tflite::TensorType input_type =
      FromTocoDataTypeToTflitToTensorType(input_data_type);
  const tflite::TensorType output_type =
      FromTocoDataTypeToTflitToTensorType(output_data_type);

  std::string output_model;
  const absl::string_view input_model_buffer(buf, length);
  auto status = mlir::lite::QuantizeModel(
      input_model_buffer, input_type, output_type, inference_tensor_type,

  // Create a tfl.transpose op that performs ZX transpose on `input`.
  auto create_z_x_transpose_op = [&](Value input) -> Value {
    RankedTensorType input_type = mlir::cast<RankedTensorType>(input.getType());
    const int input_rank = input_type.getRank();

    // Create a 1D I32 tensor for representing the dimension permutation.
    auto permuation_tensor_type =

                shape=input_shape, dtype=dtypes.float32, name='input_tensor'
            )
        ),
        options=save_options,
    )
    return model

  def _create_gather_model(self, input_type, use_variable) -> module.Module:
    class GatherModel(module.Module):
      """A simple model with a single gather."""

      def __init__(self, use_variable):
        """Initializes a GatherModel.

// of `input_type`, if the `quant_dim` is valid. On the other hand, the
// symmetry of min and max is not adjusted by this method. The QAT workflow
// should set min/max correctly (and use `narrow_range`=true, `is_signed`=true)
// if symmetric quantization is required.
TypeAttr GetQuantizedTypeAttr(Builder builder, Type input_type, Attribute min,

      ('use_constant_with_int64_input', np.int64, False),
      ('use_variable_with_int64_input', np.int64, True),
  )
  @test_util.run_v2_only
  def test_gather_model(self, input_type, use_variable):
    model = self._create_gather_model(input_type, use_variable)

    save.save(model, self._input_saved_model_path)

    rng = np.random.default_rng(seed=42)
    static_input_shape = [6]

    // If no padding is negative return the input as is.
    if (llvm::all_of(explicit_padding, [](int64_t pad) { return pad >= 0; })) {
      return value;
    }

    auto input_type = mlir::cast<RankedTensorType>(value.getType());
    auto input_shape = input_type.getShape();

    llvm::SmallVector<int64_t, 4> start;
    llvm::SmallVector<int64_t, 4> size;
    start.reserve(explicit_padding.size() / 2);