patrick-kidger / equinox

Elegant easy-to-use neural networks + scientific computing in JAX. https://docs.kidger.site/equinox/

patrick-kidger/equinox Issues

Recommended way of filtering params for weight decay
Closed 7 days ago20
Calling eqx.Module method invokes user defined __getattr__
Updated 7 days ago
Best way to do complex weight sharing
Updated 7 days ago
Filter jax lax primitives
Closed 8 days ago1
Better error messages for strict
Closed 9 days ago1
better bug hint when writing a simple neural network in equinox
Updated 9 days ago1
Initialization of large models on multi-hosts environment
Updated 10 days ago9
Ensemble model training/running with different architectures
Closed 13 days ago6
Does equinox have an example or opensource model for distributed training on multi-hosts
Closed 14 days ago
best way to write a wrapper function for tree_at
Updated 16 days ago1
MLP "depth" parameter doc
Updated 17 days ago2
test_traceback_runtime_eqx fails with latest jaxtyping (0.2.33)
Updated 19 days ago2
Bug: typing issue due to `__getattribute__`
Updated 21 days ago4
Conv2dTranspose possible bug??
Closed 23 days ago4
Getting a type error when initializing Linear module (unexpected keyword argument 'key')
Updated 23 days ago2
Vmap and lax.scan with different sequence array
Updated a month ago1
Increase minimum required python version to 3.10
Updated a month ago2
How to extract learnable parameters of a subnetwork
Closed a month ago2
Dropout causes `None`s with gradient accumulation
Updated a month ago3
`Conv` doesn't support complex dtype
Closed a month ago2
Comparisons to NNX
Updated a month ago6
`jit`ing large models for inference has bad compilation performance
Updated a month ago3
_TrivialClosureConvert mis-specifies dynamic attributes
Closed a month ago3
`filter_shard` test crashes on JAX 0.4.29
Updated a month ago4
Export attention function similar to torch.nn.functional.scaled_dot_product_attention
Updated a month ago1
Remove references to internal `jax.core.pp_*` functions
Closed a month ago2
Add a `warn_if`
Updated a month ago1
Given a state, how to get the values of its arrays?
Closed a month ago1
Question: `filter_shard` with PartitionSpecs, or other ways to ensure batched output, as in `vmap`.
Updated 2 months ago3
leaked trace error when using custom_jvp, scan, eqx.Module, and jit
Closed 2 months ago13
Question - How to assert that a `filter_spec` is valid?
Closed 2 months ago2
Diffusion example error
Closed 2 months ago2
Test fail with latest jaxtyping (0.2.29)
Closed 2 months ago2
`ValueError` from `filter_jit` when using an `eqx.Module` as a static argument on second call
Closed 2 months ago2
Right way to count cache size
Updated 2 months ago2
Error with data sharding
Closed 2 months ago5
docs request for filter_jacs/hessian
Updated 2 months ago2
A simple DeepMLP is taking 11GB/16GB in GPU RAM
Closed 2 months ago4
Flash Attention
Updated 2 months ago6
Conv Transpose doesn't exactly transpose
Closed 2 months ago4
Unroll in eqx's internal `scan`
Updated 2 months ago11
Kolmogorov Arnold Networks with equinox [kanx]
Updated 2 months ago2
Custom JVP/VJP definition within Module?
Updated 2 months ago3
Static/Dynamic issues with filter_eval_shape
Closed 2 months ago1
Jax 0.4.27: ValueError: safe_map() argument 2 is shorter than argument 1
Updated 3 months ago7
Strange error with multiple tuples as attributes and `equinox.tree_at`?
Closed 3 months ago5
How to serialize a model + state pair?
Updated 3 months ago2
merging multiple eqx.Module classes
Updated 3 months ago2
`eqx.tree_at` fails with `TypeError: Operation undefined, {x} is not a leaf of the pytree` when custom pytree attributes are modified within `__init__`
Closed 3 months ago5
Regarding trainable parameters in equinox
Closed 3 months ago9