equinox
equinox copied to clipboard
patrick-kidger
Elegant easy-to-use neural networks + scientific computing in JAX. https://docs.kidger.site/equinox/
Hi Patrick, I am building a repository of image classification models with Equinox as a personal project. I have added ResNet and LeNet so far and plan to add more....
Equinox's `filter_*` functions are very helpful functions. It would be great if there were`filter_jacfwd` and `filter_jacrev` functions. For my specific use case, I am using `eqx.filter_value_and_grad(fn)(model, ...)`. As far as...
Hello! I have a question. I am using equinox for a project here: https://github.com/mjo22/cryojax. My Modules are very nested, and in general I am interested in using grad and vmap...
Hi, I really enjoy using equinox, but sometimes, when I am trying to process high-dimensional data with `jax.vmap` it can be a quite slow. Therefore, I was wondering, if it...
Hi, not sure if this is really a bug or intended, but it's not possible to pass an `eqx.Module` as carry in a `jax.lax.scan`. Here is the MVP: ```python import...
In the following code: ```python from beartype import beartype import equinox as eqx from jaxtyping import Array, Float @beartype class MyClass(eqx.Module): x: Float[Array, ""] def fn(self, y: bool) -> Float[Array,...
When attempting to use the `Lambda` layer in Equinox, it appears that the `__init__` method's parameters are not being correctly resolved. This results in a 'too many arguments' error during...
Hi Patrick! Thanks for the great work with the library. I'm just getting started out in JAX after years of PyTorch and have a question regarding composable `LSTMCells` in equinox....
I have a code where using jax.jit like this works: `@partial(jax.jit, static_argnums=1)` But `@eqx.filter_jit` in the same place causes jax.clear_caches to crash with either Segmentation fault or sometimes with `double...
Metadata
Owner
Metadata
Elegant easy-to-use neural networks + scientific computing in JAX. https://docs.kidger.site/equinox/