tensor_spec.TensorSpec(
                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):

      ('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]

          signature_key,
          tags,
          inputs={input_key: in_placeholder},
          outputs={output_key: output_tensor},
      )

      return in_placeholder

  def _create_gather_model(self, input_type, use_variable):
    class GatherModel(autotrackable.AutoTrackable):
      """A simple model with a single gather."""

      def __init__(self, use_variable):

      ('small_gather_use_variable', dtypes.int32, True, False),
  )
  @test_util.run_v2_only
  def test_gather_model(
      self, input_type, use_variable, expect_quantized_gather
  ):
    model = self._create_gather_model(input_type, use_variable)

    saved_model_save.save(model, self._input_saved_model_path)

    tags = {tag_constants.SERVING}

    quantization_options = quant_opts_pb2.QuantizationOptions(