equinox

equinox copied to clipboard

Key/Value Serialization

3

I recently began using the equinox serialization library to store model checkpoints. My current understanding of that code is that it writes all the non-static pytree fields in a specific...

dieterichlawson

`filter_vmap()` on `eqx.nn.*` Modules with `key` as a keyword-only argument

2

Thanks for the library! If I understand correctly: * Keyword arguments cannot be used with functions wrapped with `eqx.filter_vmap` * The `key` argument of all equinox Module methods is a...

mthiboust

Best practices to convert `torch.nn.Module` to `eqx.Module`

12

Thanks for the great package! I was wondering whether there was some documentation regarding the best practice for converting `torch.nn.Module` to `eqx.Module`? In particular - It is quite clear that...

emilemathieu

using batchnorm/dropout layers from flax.linen along with diffrax package

14

Hi Patrick, I am using the diffrax ode solver in my code and will need to use batchnorm/dropout layers in the function that will be passed to the solver. However...

Negar-Erfanian

Segmentation fault in tree_flatten when subclass method is passed to superclass __init__

5

I have the following code, which causes a segmentation fault when tree_flatten is called on any instance of SubComponent. I think it has to do with the `jtu.Partial(method, self)` in...

tttc3

Obtaining the name, values pairs in Equinox

10

Pytorch has this default function that yields the parameter name-value pairs in a module including all its submodules called [torch.nn.Module.named_parameters()](https://pytorch.org/docs/stable/generated/torch.nn.Module.html?highlight=named_param#torch.nn.Module.named_parameters). This function can be used in a loop as follows:...

sachith-gunasekara

optim.init fails if not filtering by inexact array

1

Why does this fail?: ```python import optax import equinox as eqx m = eqx.nn.MLP(1,1,2,1, key=PRNGKey(0)) optim = optax.adamw(1e-3) optim.init(eqx.filter(m,eqx.is_inexact_array)) # this does not fail optim.init(m) ``` whereas in the [bert...

agpcd

shouldn't we renormalize when Inference=True for the dropout layers?

2

I was checking the [BERT example](https://docs.kidger.site/equinox/examples/bert/#bert-implementation) and realized dropout simply doesn't do anything when called with Inference=True, but according to my understanding, you need to renormalize at inference time to...

agpcd

Multigpu training?

6

So I know for a while there was an issue with equinox not playing well with pmap. With the advent of pjit (and I think now just jit), are there...

rdilip

Include Q&A Wiki?

1

Hi Patrick, a lot of issues are labelled "question". Perhaps it would be better to have a wiki to post these questions to leave the "issues" section for actual issues?...

Artur-Galstyan