v0.3.2
JAX JIT Compatibility
All string-based fields (res_dim, chunks, dtype, mode, etc.) across all ORC modules are now marked eqx.field(static=True). This makes ORC models fully compatible with jax.jit when the model is passed as an explicit pytree argument, and resolves ConcretizationTypeError failures in training functions like train_RCForecaster under jax.jit. A new JAX JIT Compatibility example notebook documents the supported patterns and sharp bits for jax.jit vs eqx.filter_jit. The preference for eqx.filter_jit is strongly advertised.
Type Checking
- ty has been added as a development dependency and integrated into the CI pipeline (ty check src/ runs on every push). Many incorrect type annotations across the codebase have been fixed as a result
Documentation
- Notebook outputs are now generated at docs build time via mkdocs-jupyter (execute: true), keeping notebook source files output-free in the repository.
- Fixed active navigation item color in dark mode.
- Fixed a broken figure reference in the JOSS paper (architecture_paper.pdf).
Other
- Removed implicit float64 enabling; users must now call jax.config.update("jax_enable_x64", True) explicitly. A warning is raised if - float64 is not enabled at import time.
- Fixed bug with reversed order of res_dim and in_state for GRU driver.
- Removed plt.show() calls from built-in visualisation utilities.
- Significantly reduced repository clone size by rewriting git history to remove large historical blobs.
- Expanded test coverage for jit, vmap, and autodiff