equinox

equinox copied to clipboard

Recover optimised model parameters as a dictionary?

4

As a minimal example, lets say we have something like this (the real situation is custom modules which contain other modules, hence the flatten). If I look at flat_model obviously...

adam-hartshorne

Unensembling

6

I am enjoying Equinox a great deal and it seems a huge ergonomic improvement over raw JAX in many cases. Thank you for building it! I am using `Equinox` to...

JeppeKlitgaard

Implementation of `filter_map` using `jax.lax.map`

4

I am really enjoying using `filter_vmap` and `filter_pmap` to very easily transform functions that act on PyTrees – Truly a game-changer! Playing around on a Google Cloud TPU, I have...

JeppeKlitgaard

Class fields and instance fields behavior

3

Hello, Patrick currently, data classes allow fields to be either class or instance variables for example ```python @dc.dataclass class Test: a:int = 1 def __init__(self) -> None: pass Test().__dict__ #...

ASEM000

- [ ] Have it support classical treeverse, for improved results when using `equinox.internal.scan(..., kind="checkpointed")`. - [ ] Have it support saving the residuals of each step, rather than just...

patrick-kidger

Recompilation of purely static `Module` bound methods

3

Consider the following code: ```python import jax import jax.numpy as jnp import equinox as eqx from functools import partial class Model(eqx.Module): def f(self): return model = Model() # We get...

KeAWang

Easy mechanism to explicitly specify dynamic params

2

It would be nice if there were a mechanism to state which `eqx.Module` members are dynamic and a filter transformation that treats only elements marked dynamic as dynamic. This would...

allen-adastra

Using Frozen Parameters to Freeze Differing Parameters Based Upon Epoch

12

I am trying to figure out if it is possible to efficiently use something like the frozen parameters approach described here https://docs.kidger.site/equinox/examples/frozen_layer/ , to setup up a training loop such...

adam-hartshorne

L2 norm of weights

1

Hi, This week I started using equinox (great project!), and am slowly figuring out the best practices. I'm trying to do a simple L2 norm of the weights, which requires...

tychovdo

Runtime penalty with fields that have expensive inits

5

I was surprised about the runtime of a posterior estimation I was coding. Here is a MWA to reproduce the issue: ```python import jax.numpy as jnp import equinox as eqx...

pmelchior