Enable T5x to take native_serialization_platforms for jax2tf

PiperOrigin-RevId: 559827754

t5x/export_lib.py

@@ -293,6 +293,7 @@ def create_inference_function(
     enable_xla: bool = True,
     polymorphic_shapes_inputs: Optional[Any] = None,
     native_lowering: bool = False,
+    native_lowering_platforms: Sequence[str] = (),
     model_fn_extra_kwargs: Optional[Mapping[str, Any]] = None,
 ) -> Callable[[Mapping[str, Any], Any], PyTree]:
   """Fetches a model and returns the inference function based on inference_mode."""
@@ -355,6 +356,7 @@ def model_fn(
         model_fn,
         polymorphic_shapes=[None, polymorphic_shapes_inputs],
         native_serialization=native_lowering,
+        native_serialization_platforms=native_lowering_platforms,
         enable_xla=enable_xla,
     )
 
@@ -1297,6 +1299,7 @@ def save(
     mixture_or_task_name: Optional[str] = None,
     validation_examples: Optional[List[Any]] = None,
     native_lowering: bool = False,
+    native_lowering_platforms: Sequence[str] = (),
     enable_xla: bool = True,
     decode_outputs: Optional[bool] = None,
     trailing_shapes: Optional[Mapping[str, Tuple[int, ...]]] = None,
@@ -1344,6 +1347,11 @@ def save(
       model.
     native_lowering: for experimental purposes only -- if True, don't convert
       Jax fns to TF fns.
+    native_lowering_platforms: In conjunction with `native_lowering`, specify
+      the platform(s) for which to lower the code. Must be a tuple of strings,
+      including a subset of: 'cpu', 'cuda', 'rocm', 'tpu'. The default
+      (empty tuple), specifies the JAX default backend on the machine where the
+      lowering is done.
     enable_xla: Defaults to true. If false, jax2tf conversion only emits non-XLA
       ops.
     decode_outputs: Optional bool. If provided, determines whether to decode the
@@ -1417,6 +1425,7 @@ def save(
           input_signature, preprocessor
       ),
       native_lowering=native_lowering,
+      native_lowering_platforms=native_lowering_platforms,
   )
 
   logging.info('Loading parameters from checkpoint...')