Repository navigation

All pull requests

Compute float16 and bfloat16 average pooling in float32
#1280AHMETHAKANBEZIR1 opened Oct 26, 2026Contributor
Fix masked attention weights and half-precision logit overflow
#1279AHMETHAKANBEZIR1 opened Oct 26, 2026Contributor
Fix NaN gradients when Dropout drops every entry
#1278AHMETHAKANBEZIR1 opened Oct 25, 2026Contributor
Sum broadcast-input cotangents in vprim transpose
#1275Don-Yin opened Oct 24, 2026
3
Make error_if work for shard_map
#1262spraharsh opened Sep 19, 2026
4
Fix closure-converted input checks within shard map
#1261spraharsh opened Sep 19, 2026
1
Add _FastModuleMeta for skipping instantiation-time checks on BoundMethod
#1260nstarman opened Sep 8, 2026Contributor3 tasks done
Don't leak partially-constructed modules when __init__ raises
#1251nstarman opened Jul 20, 2026Contributor
4
Fix eager error_if hanging when inputs are sharded over multiple devices
#1242discobot opened Jun 13, 2026Contributor
4
Switch eqx.nn.Linear.__call__ to x @ weight.T to eliminate vmap+grad permutes (gated on GPU benchmark)
#1235mvanhorn opened May 19, 2026Contributor
5
perf: add fastpaths in _ModuleMeta.__call__
#1230nstarman opened May 14, 2026Contributor
5
Reformat dependencies in pyproject.toml
#1194dhirschfeld opened Mar 7, 2026
1
Use @jtu.register_static as decorator on _Missing
#1193nstarman opened Mar 4, 2026Contributor
6
Enable eqx.nn.State to track and return Intermediate values
#1165arenaudineau opened Jan 10, 2026
Fixes for mypy + Enumeration
#1122patrick-kidger opened Oct 17, 2025Owner
perf: enable module flatten/unflatten fastpath
#1117nstarman opened Oct 12, 2025Contributor
12
Add exact spectral norm feature
#831lockwo opened Sep 6, 2024Contributor
8
BatchNorm training instability fix
#675andrewdipper opened Mar 7, 2024
21
Now running tree_check when initialising a Module
#478patrick-kidger opened Sep 7, 2023Owner
First pass at making error_if work with sharding.
#342patrick-kidger opened May 8, 2023Owner

Navigation Menu

Repository navigation

All pull requests

Compute float16 and bfloat16 average pooling in float32

Fix masked attention weights and half-precision logit overflow

Fix NaN gradients when Dropout drops every entry

Sum broadcast-input cotangents in vprim transpose

Make error_if work for shard_map

Fix closure-converted input checks within shard map

Add `_FastModuleMeta` for skipping instantiation-time checks on `BoundMethod`

Don't leak partially-constructed modules when `init` raises

Fix eager `error_if` hanging when inputs are sharded over multiple devices

Switch `eqx.nn.Linear.call` to `x @ weight.T` to eliminate vmap+grad permutes (gated on GPU benchmark)

perf: add fastpaths in `_ModuleMeta.call`

Reformat dependencies in pyproject.toml

Use `@jtu.register_static` as decorator on `_Missing`

Enable eqx.nn.State to track and return Intermediate values

Fixes for mypy + Enumeration

perf: enable module flatten/unflatten fastpath

Add exact spectral norm feature

BatchNorm training instability fix

Now running tree_check when initialising a Module

First pass at making error_if work with sharding.

Pull requests