equinox

equinox copied to clipboard

Best Dataloader practices

6

I was wondering if the equinox community has good advice / best practices for creating dataloaders that work well with jax? I've done some stuff from scratch but I tend...

jejjohnson

VGG support or eqxvision repository?

1

Hi, so with some work I have now VGG16 with imagenet weights working in equinox. Here is the basic model description: ```python import jax import jax.numpy as jp from equinox...

oarriaga

Feature Request: add attention utility functions like flax

3

It would be nice to split up the the [MultiheadAttention](https://github.com/patrick-kidger/equinox/blob/main/equinox/nn/attention.py#L14) Module code into utility functions for calculating attention weights and multihead attention without projection layers. Flax does this -- but...

jenkspt

This looks like it might be a bit of fun to try when displaying PyTrees i.e. models: https://github.com/AharonSambol/PrettyPrintTree

patrick-kidger

Added `eqx.experimental.noinline`

8

TL;DR: XLA sub-graphs! ### Background At present, JAX inlines the entire computation into a single XLA graph. However, many scientific computing applications involve defining some modestly complicated function and then...

patrick-kidger

Out-of-place updates with custom `__init__()` constructor

3

Hi @patrick-kidger! --- Sorry for the long-winded feature request, it got out of hand as I was writing it! The basic question is this: **would you be interested in adding...

Charl-AI

List a few projects using Equinox

3

People are starting to build cool stuff with Equinox. It might be worth adding links to these in the documentation for examples. - Diffrax - sympy2jax - Lucidrains' work -...

patrick-kidger

Equinox equivalent of FuncTorch's make_functional

4

Hi @patrick-kidger, Apologises for opening an issue as this is probably something I've missed, but is it at all possible to separate an `eqx.Module` object into a pure function and...

AlphaBetaGamma96

By adding a few meta tags it's possible to get Twitter to display a custom image etc. when linking to the documentation. Adding this in to the docs sounds like...

patrick-kidger

At the moment, `get_state` and `set_state` demand that their arguments have precisely the same shape, dtype, and choice of batch axes. It would make sense to allow some compatibility between...

patrick-kidger