Building Neural Networks with Equinox

Explain Equinox's philosophy: models as pytrees of arrays with no separate parameter dictionaries.

Define a custom eqx.Module subclass with parameters, submodules, and static fields.

Use eqx.filter, eqx.partition, and eqx.combine to separate trainable arrays from static data.

Apply eqx.filter_jit and eqx.filter_grad to train Equinox models cleanly.

Use an inference-mode flag to conditionally apply operations like dropout, demonstrating correct PRNG key management.

Manage and update evolving state (e.g., BatchNorm running statistics) by partitioning the model and returning updated state from the forward pass.