equinox

equinox copied to clipboard

This is a general code cleanliness thing: it is [not great style](https://docs.kidger.site/equinox/pattern/) to have two classes that are not instantiable, with one inheriting from the other. In this case, `Adaptive{Avg,Max}Pool{1,2,3}d`...

patrick-kidger

Orbax Integration

4

Integrating with Orbax would be great. I actually have a working `class EquinoxCheckPointHandler(ocp.CheckpointHandler)` that seems to work for generic instances of `eqx.Module`. Would there be interest in finding a way...

allen-adastra

How to update the weight of one class in another class

5

self.nn = update will raise a error. How to change the weight of self.nn

jimmylihui

Support for data format NHWC in convolutional layers

5

Hello, First of all, thank you for your great work! I noticed that Equinoxs' convolutional layers use _NCHW_ data format, which, according to [Nvidia Tensor Layout](https://docs.nvidia.com/deeplearning/performance/dl-performance-convolutional/index.html#tensor-layout), is not optimal for...

MathiesW

Can't import equinox 0.10.6 in python 3.9.0

6

When I try to import equinox 0.10.6 in python 3.9.0 I get the following error: ``` >>> import equinox as eqx Traceback (most recent call last): File "", line 1,...

dieterichlawson

[Question] Using eqx.Linear with inferred input size?

1

Hi Patrick! I'm currently converting a Haiku codebase to use equinox, and there's a _lot_ of boilerplate that tries to construct MLPs before knowing the full size of the input...

phinate

Quantization

2

Do you plan on supporting quantization? Perhaps as offered by AQT: https://github.com/google/aqt

jkuntzer

Alternative Sharing

3

I was looking at how sharing was added into equinox, and was a bit surprised at the approach taken. I had a similar challenge where I had a PyTree with...

bowlingmh

Slow JIT compilation when doing teacher forcing

1

Hello Patrick, I recently tried to move from torch to jax/equinox (thanks a lot for your contributions both for diffrax and equinox) and i tried to adapt all my code...

itsakk

Unstable PPO?

7

Hello Patrick, I am doing an implementation of the PPO algorithm for a custom environment and first wanted to test things out with a standard example and I choose CartPole-v1...

stergiosba