Deprecate experimental_from_jax in favor of jax2tf + `from_saved_…

…model`. The recommended way to convert JAX models includes using PAX/SAX/T5X/Orbax Export. For a JAX-native model, `jax2tf` can be used directly. Once a TF SavedModel is procured using the aforementioned tools, then `lite.TFLiteConverter.from_saved_model` can be used without the need to specify any additional flags. PiperOrigin-RevId: 548228454
tensorflow · Jul 14, 2023 · 40f0167 · 40f0167
1 parent 387f5f4
commit 40f0167
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/tensorflow/lite/python/lite.py b/tensorflow/lite/python/lite.py
@@ -2127,6 +2127,11 @@ def from_keras_model(cls, model):
     return TFLiteKerasModelConverterV2(model)
 
   @classmethod
+  @_deprecation.deprecated(
+      None,
+      "Use `jax2tf.convert` and (`lite.TFLiteConverter.from_saved_model`"
+      " or `lite.TFLiteConverter.from_concrete_functions`) instead.",
+  )
   def experimental_from_jax(cls, serving_funcs, inputs):
     # Experimental API, subject to changes.
     # TODO(b/197690428): Currently only support single function.