Move converter related tflite functions to tensorflow/lite repo

PiperOrigin-RevId: 688270228

Author: ecalubaquib

jax/experimental/jax2tf/README.md

@@ -103,7 +103,6 @@ For more involved examples, please see examples involving:
 
    * SavedModel for archival ([examples below](#usage-saved-model)), including
      saving [batch-polymorphic functions](#shape-polymorphic-conversion),
-   * TensorFlow Lite ([examples](https://github.com/jax-ml/jax/blob/main/jax/experimental/jax2tf/examples/tflite/mnist/README.md)),
    * TensorFlow.js ([examples](https://github.com/jax-ml/jax/blob/main/jax/experimental/jax2tf/examples/tf_js/quickdraw/README.md)),
    * TFX ([examples](https://github.com/tensorflow/tfx/blob/master/tfx/examples/penguin/README.md#instructions-for-using-flax)),
    * TensorFlow Hub and Keras ([examples](https://github.com/jax-ml/jax/blob/main/jax/experimental/jax2tf/examples/README.md)).

jax/experimental/jax2tf/examples/README.md

@@ -9,7 +9,7 @@ This directory contains a number of examples of using the
   * save SavedModel from trained MNIST models, using both Flax and pure JAX.
   * reuse the feature-extractor part of the trained MNIST model
     in a larger TensorFlow Keras model.
-  * use Flax models with TensorFlow Serving, TensorFlow JavaScript, and TensorFlow Lite.
+  * use Flax models with TensorFlow Serving, TensorFlow JavaScript.
 
 You can also find usage examples in other projects:
 
@@ -176,16 +176,15 @@ At the moment, the open-source TensorFlow model server is missing XLA support,
 but the Google version can be used, as shown in the
 [serving examples](https://github.com/jax-ml/jax/blob/main/jax/experimental/jax2tf/examples/serving/README.md).
 
-# Using jax2tf with TensorFlow Lite and TensorFlow JavaScript
+# Using jax2tf with TensorFlow JavaScript
 
 A jax2tf-generated SavedModel can also be converted to a format usable with
-TensorFlow Lite or TensorFlow.js, by using the appropriate converters from SavedModel. 
+TensorFlow.js, by using the appropriate converters from SavedModel. 
 At the moment, these converters may reject some jax2tf-generated SavedModels due to
 some ops not yet implemented in the converters. As a partial workaround, one
 can pass the `enable_xla=False` parameter to `jax2tf.convert` to direct 
 `jax2tf` to avoid problematic ops. This will increase the coverage, and in fact
 most, but not all, Flax examples can be converted this way. 
 
-Check out the [MNIST TensorFlow Lite](https://github.com/jax-ml/jax/blob/main/jax/experimental/jax2tf/examples/tflite/mnist/README.md)
-and the
+Check out
 [Quickdraw TensorFlow.js example](https://github.com/jax-ml/jax/blob/main/jax/experimental/jax2tf/examples/tf_js/quickdraw/README.md).

jax/experimental/jax2tf/g3doc/BUILD

@@ -15,7 +15,10 @@ licenses(["notice"])
 
 package(
     default_applicable_licenses = [],
-    default_visibility = ["//jax/experimental/jax2tf:__subpackages__"],
+    default_visibility = [
+        "//jax/experimental/jax2tf:__subpackages__",
+        "//third_party/tensorflow/lite/experimental/mlir/testing/jax:__subpackages__",
+    ],
 )
 
 filegroup(

jax/experimental/jax2tf/tests/converters.py

@@ -20,7 +20,6 @@
 from typing import Any
 
 from jax.experimental import jax2tf
-import tensorflow as tf
 import tensorflowjs as tfjs
 
 from jax.experimental.jax2tf.tests.model_harness import ModelHarness
@@ -51,56 +50,14 @@ def jax2tfjs(harness: ModelHarness):
         model_dir=model_dir)
 
 
-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
-
-
 ALL_CONVERTERS = [
     # jax2tf with XLA support (enable_xla=True).
     Converter(name='jax2tf_xla', convert_fn=jax2tf_convert),
     # jax2tf without XLA support (enable_xla=False).
     Converter(
         name='jax2tf_noxla',
-        convert_fn=functools.partial(jax2tf_convert, enable_xla=False)),
+        convert_fn=functools.partial(jax2tf_convert, enable_xla=False),
+    ),
     # Convert JAX to Tensorflow.JS.
     Converter(name='jax2tfjs', convert_fn=jax2tfjs, compare_numerics=False),
-    # Convert JAX to TFLIte.
-    Converter(name='jax2tflite', convert_fn=jax2tflite),
-    # Convert JAX to TFLIte with support for Flex ops.
-    Converter(
-        name='jax2tflite+flex',
-        convert_fn=functools.partial(jax2tflite, use_flex_ops=True))
 ]

jax/experimental/jax2tf/tests/flax_models/BUILD

@@ -19,7 +19,10 @@ licenses(["notice"])
 
 package(
     default_applicable_licenses = [],
-    default_visibility = ["//jax/experimental/jax2tf:__subpackages__"],
+    default_visibility = [
+        "//jax/experimental/jax2tf:__subpackages__",
+        "//third_party/tensorflow/lite/experimental/mlir/testing/jax/test_models:__subpackages__",
+    ],
 )
 
 py_library(