jax

2828bce2 - Add check to `lax.composite` to prevent `DynamicJaxprTracer` type errors.

Add check to `lax.composite` to prevent `DynamicJaxprTracer` type errors. There's some confusion on whether jax arrays can be used inside the attributes, so I made the error more explicit. PiperOrigin-RevId: 724316766