Now running tree_check when initialising a Module#478

Open

patrick-kidger wants to merge 18 commits into

mainfrom

Commits on Sep 5, 2023

Renamed eqx.tree_inference -> eqx.nn.inference_mode, as it's really an eqx.nn thing, not an eqx thing.
patrick-kidger
committed
Embedding now supports initialistion with just a weight. It now validates its input as being a scalar index.
patrick-kidger
committed
LayerNorm is now a bit safer, by explicitly validating shapes. (dlwh pointed out that otherwise this can silently normalise over sequence dimensions in transformer models if you miss the vmap)
patrick-kidger
committed

Commits on Sep 7, 2023

Commits on Sep 12, 2023