Raise a more useful error when doing a dtensor export with pspecs but not all params are jax arrays.

PiperOrigin-RevId: 875283768

Author: BrianWieder

export/orbax/export/modules/tensorflow_module.py

@@ -276,11 +276,22 @@ def jax_params_to_tf_variables(
       if not all(
           isinstance(x, jax.Array) for x in jax.tree_util.tree_leaves(params)
       ):
-        logging.warning(
-            'Some params are not jax.Array, DTensor export will not take'
-            ' effect.Falling back to traditional TF export.'
+        flattened_params = jax.tree_util.tree_leaves_with_path(params)
+        non_jax_array_info = [
+            f'{jax.tree_util.keystr(path)} (type: {type(x).__name__})'
+            for path, x in flattened_params
+            if not isinstance(x, jax.Array)
+        ]
+        if len(non_jax_array_info) > 10:
+          omitted = len(non_jax_array_info) - 10
+          non_jax_array_info = non_jax_array_info[:10] + [
+              f'...and {omitted} more'
+          ]
+        raise ValueError(
+            'All parameters must be JAX arrays when DTensor export is enabled.'
+            ' Found non-JAX array parameters at:'
+            f' {", ".join(non_jax_array_info)}'
         )
-        mesh = None
 
     if mesh is None and pspecs is not None:
       raise ValueError(

export/orbax/export/modules/tensorflow_module_test.py

@@ -15,6 +15,7 @@
 import collections
 import logging
 import os
+from unittest import mock
 
 from absl.testing import parameterized
 import chex
@@ -180,6 +181,30 @@ def test_jax_array(self):
     self.assertEqual(variables[0].name, 'arr:0')
     self.assertAllEqual(variables[0], global_input_data)
 
+  def test_dtensor_export_non_jax_array_error(self):
+    global_mesh = jax.sharding.Mesh(
+        np.array(jax.local_devices(backend='cpu')), 'x'
+    )
+    mesh_axes = jax.sharding.PartitionSpec('x')
+
+    arr = jnp.array([1, 2, 3])
+    with mock.patch.object(
+        tensorflow_module.dtensor_utils,
+        'get_current_mesh',
+        return_value=global_mesh,
+    ):
+      with self.assertRaisesRegex(
+          ValueError,
+          r'All parameters must be JAX arrays when DTensor export is enabled.*'
+          r"Found non-JAX array parameters at: \['some_non_array'\] \(type:"
+          r' ellipsis\)',
+      ):
+        TensorFlowModule(
+            params={'arr': arr, 'some_non_array': ...},
+            apply_fn=DEFAULT_APPLY_FN,
+            pspecs={'arr': mesh_axes, 'some_non_array': None},
+        )
+
   @parameterized.parameters(True, False)
   def test_call_tf_module_methods(self, jit_compile):