Add check to lax.composite to prevent DynamicJaxprTracer type errors.

jax/_src/lax/lax.py

@@ -1041,6 +1041,15 @@ def _decorator(*args, **kwargs):
                                      args, kwargs)
     flat_args, in_tree = tree_util.tree_flatten(args)
     in_avals = tuple(core.get_aval(x) for x in flat_args)
+    if any(isinstance(v, core.Tracer) for v in kwargs.values()):
+      raise UnexpectedTracerError(
+          "Found a JAX Tracer as an attribute in the decomposition for the "
+          f"composite op '{name}'. This means that the decomposition function "
+          "closes over a value that is involved in a JAX transformation. "
+          "Any values that aren't explicitly known at compile time must be "
+          "explicitly passed as arguments to the composite."
+          "\n\nNote: If you are passing jax arrays as attributes, use numpy "
+          "arrays instead.")
     closed_jaxpr, out_tree = _trace_composite_to_jaxpr(
         partial(decomposition, **kwargs), in_tree, in_avals, name, debug_info
     )

tests/lax_test.py

@@ -4610,6 +4610,25 @@ def decomposition(x, **_):
         mlir_module)
     self.assertIn("func.func private @my.tangent", mlir_module)
 
+  def test_composite_unsupported_attribute_dtypes(self):
+
+    def my_tangent_composite_with_attributes(x):
+      def decomposition(x, **_):
+        return lax.sin(x) / lax.cos(x)
+      return lax.composite(decomposition, "my.tangent")(
+          x, tensor=jnp.zeros((1, 2), dtype=jnp.float32)
+      )
+
+    pi = jnp.pi
+    x = jnp.array([0.0, pi / 4, 3 * pi / 4, pi], dtype=jnp.float32)
+
+    with self.assertRaisesRegex(
+        UnexpectedTracerError,
+        "Note: If you are passing jax arrays as attributes, use numpy arrays "
+        "instead."
+    ):
+      jax.jit(my_tangent_composite_with_attributes).lower(x).as_text()
+
   def test_composite_with_non_default_version(self):
     @partial(lax.composite, name="my.square", version=1)
     def my_square_with_version(x):