equinox

equinox copied to clipboard

merging multiple eqx.Module classes

2

Hi all, I am building a multi-physics solver using JAX and equinox, where I want to merge multiple dataclasses inherited from eqx.Module. If I am using standard Python classes, the...

ToshiyukiBandai

Getting a type error when initializing Linear module (unexpected keyword argument 'key')

1

This may not be the medium for petty user bugs, but I am not able to find any help on general forums. I am not sure what is causing the...

AbhinavRao23

Weird `jax` error when trying vmap twice while using batchnorm

5

I'm trying to 1. Utilize the init/apply design pattern. 2. Vmap a module twice (where the module uses batchnorm) in order to handle inputs of shape `[n_samples, batch_dim, feature_dim]` instead...

PaulScemama

BatchNorm training instability fix

1

This is in reference to issue [659](https://github.com/patrick-kidger/equinox/issues/659). I modified BatchNorm to have two approaches `"batch"` and `"ema"`. `"batch"` just uses the batch statistics during training time. If approach is not...

andrewdipper

lax.scan for equinox Modules ?

9

Hi all, I've create Equinox modules and I want to process them with `jax.lax.scan`, but I don't see any `equinox.filter_scan`. Is there a particular reason for that ? Here's some...

mchagneux

gradient with sequential and jax.nn.<activation functions>

4

Hi Patrick, I am trying to use sequential to make my code clean, and I run into an issue when I try to get gradients I put a simple example...

mfouesneau

Graph Neural Networks

1

Greetings Patrick, I was wondering, what are your thoughts when it comes to Graph Neural Networks? Would a GNNlayer be implemented as a module for Equinox or an advanced example...

stergiosba

Add a [source] to docs [feature request]

3

For users that are interested/regularly investigate the source code, having a method to go from docs to associated source code is very helpful (this is what I am talking about:...

lockwo

Optimally training ensembles.

4

What is the recommended way to train an ensemble model like the one in the documentation? I have re-produced the model here and would appreciate some help with this. ```...

BabaYara

Type annotations for "struct of arrays" ?

1

I love `equinox` and `jaxtyping` and would like to figure out a nice way to annotate stacked or batched `eqx.Modules` but so far I've failed. For example, I'll often do...

nboyd