code: def jax2tflite(harness: ModelHarness, use_flex_ops: bool = False):
      """Returns a converter with Flex ops linked in iff `use_flex_ops==True`."""
      tf_fn = tf.function(
          jax2tf_convert(harness, enable_xla=False),
          input_signature=harness.tf_input_signature,
          autograph=False,
      )
      apply_tf = tf_fn.get_concrete_function()
      converter = tf.lite.TFLiteConverter.from_concrete_functions(
          [apply_tf], tf_fn
      )
      supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS]
      if use_flex_ops:
          supported_ops.append(tf.lite.OpsSet.SELECT_TF_OPS)
      converter.target_spec.supported_ops = supported_ops
  
      # Convert the model.
      tflite_model = converter.convert()
  
      # Construct an interpreter for doing a numerical comparison.
      interpreter = tf.lite.Interpreter(model_content=tflite_model)
      interpreter.allocate_tensors()
  
      inputs = interpreter.get_input_details()
      output_details = interpreter.get_output_details()
      outputs = tuple(
          interpreter.tensor(out["index"]) for out in output_details
      )
  
      def apply_tflite(*xs):
          assert len(xs) == len(inputs)
          for i, x in enumerate(xs):
              interpreter.set_tensor(inputs[i]["index"], x)
          interpreter.invoke()
          if len(outputs) > 1:
              return tuple(o() for o in outputs)
          else:
              return outputs[0]()
  
      return apply_tflite
  
file_path: jax/jax/experimental/jax2tf/tests/converters.py
imports: import dataclasses
  import functools
  import tempfile
  from typing import Any, Callable
  from jax.experimental import jax2tf
  import tensorflow
  import tensorflowjs
  from jax.experimental.jax2tf.tests.model_harness import ModelHarness
label_desc: Use image node_id 3 for calling a global function with example usage: jax2tflite(harness, use_flex_ops) and returns: apply_tflite, tuple, outputs