generics.maybe_jit
generics.maybe_jit(func=None, **jit_kwargs)Decorator to JIT compile a function using JAX, if available.
Parameters
| Name | Type | Description | Default |
|---|---|---|---|
| func | Callable | None | The function to be JIT compiled. | None |
| jit_kwargs | dict | Keyword arguments to be passed to jax.jit. | {} |
Returns
| Name | Type | Description |
|---|---|---|
| Callable | The decorated function or method |