_CalibrationMethod = (
    stablehlo_quant_config_pb2.CalibrationOptions.CalibrationMethod
)


class CustomAggregatorTest(test.TestCase):

  def setUp(self):
    super(CustomAggregatorTest, self).setUp()
    ops.disable_eager_execution()

  def testBypassAndMinMax(self):
    with self.session():
      input_tensor = array_ops.constant(
          [1.0, 2.0, 3.0, 4.0, 5.0], dtypes.float32
      )

  @def_function.function(
      input_signature=[tensor_spec.TensorSpec(shape=None, dtype=dtypes.string)])
  def lookup(self, word):
    return self.table.lookup(word)


def get_simple_session():
  ops.disable_eager_execution()
  sess = session_lib.Session()
  variables.Variable(1.)
  sess.run(variables.global_variables_initializer())
  return sess


MODULE_CTORS = {

  # Default TF1.x uses reference variables that are not supported by SavedModel
  # v1 Importer. To use SavedModel V1 Importer, resource variables should be
  # enabled.
  tf.enable_resource_variables()
  tf.compat.v1.disable_eager_execution()


# This function needs to take a "create_module_fn", as opposed to just the
# module itself, because the creation of the module has to be delayed until