code: def test_simple_model(dataset, conversion_function, model_cls):
      model = model_cls(weights="imagenet")
  
      # For the untrained models, do a very small amount of training so that the
      # batch norm stats (and, less importantly, the weights) have sensible
      # values.
      model.compile(
          optimizer=tf.keras.optimizers.Adam(1e-4),
          loss="sparse_categorical_crossentropy",
      )
      model.fit(dataset, epochs=1)
  
      model_lce = conversion_function(model)
  
      if model_cls == toy_model_int8:
          # 0.025 atol is chosen because it allows an off-by-one error in the int8
          # model (with +-3 scales) but not off-by-two.
          rtol = 0.01
          atol = 0.025
      else:
          rtol = 0.001
          atol = 0.001
  
      # Test on a single batch of images
      inputs = next(iter(dataset.take(1)))[0]
      outputs = model(inputs).numpy()
      assert_model_output(
          model_lce, inputs.numpy(), outputs, rtol, atol
      )
  
      # Test on some random inputs
      input_shape = (10, *model.input.shape[1:])
      inputs = np.random.uniform(-1, 1, size=input_shape).astype(
          np.float32
      )
      outputs = model(inputs).numpy()
      assert_model_output(model_lce, inputs, outputs, rtol, atol)
  
imports: import functools
  import os
  import sys
  import tempfile
  import larq
  import numpy
  import pytest
  import tensorflow
  import tensorflow_datasets
  from larq_compute_engine.mlir.python.converter import convert_keras_model, convert_saved_model
  from larq_compute_engine.tflite.python.interpreter import Interpreter
  from preprocess import preprocess_image_tensor, IMAGE_SIZE
label_desc: Use image node_id 10 for calling a global function with example usage: test_simple_model(dataset, conversion_function, model_cls) without return types