Merge pull request #1456 from jheek:better-transform-error

PiperOrigin-RevId: 390124613

Author: Flax Authors

flax/core/tracers.py

@@ -16,6 +16,8 @@
 
 import jax
 
+from .. import errors
+
 
 def current_trace():
   """Returns the innermost Jax tracer."""
@@ -32,4 +34,4 @@ def trace_level(main):
 def check_trace_level(base_level):
   level = trace_level(current_trace())
   if level != base_level:
-    raise ValueError('Jax transforms and modules cannot be mixed.')
+    raise errors.JaxTransformError()

flax/errors.py

@@ -263,6 +263,19 @@ def __init__(self, col, variable_name, scope_path):
                      f'"{scope_path}" because collection "{col}" is immutable.')
 
 
+class JaxTransformError(FlaxError):
+  """
+  JAX transforms and Flax modules cannot be mixed.
+
+  JAX's functional transformations expect pure function.
+  When you want to use JAX transformations **inside** Flax models,
+  you should make use of the Flax transformation wrappers
+  (e.g.: ``flax.linen.vmap``, ``flax.linen.scan``, etc.).
+  """
+  def __init__(self):
+    super().__init__('Jax transforms and Flax models cannot be mixed.')
+
+
 #################################################
 # module.py errors                              #
 #################################################