equinox

equinox copied to clipboard

JaxTyping warning and performance issues

1

Hi @patrick-kidger For reasons beyond my control, I have to use an old version of JAX, which means an older version of Equinox as well (0.9.2). I get this error...

ddrous

```error_if``` support for functions and levels

3

Hi Patrick, I recently thought about an interesting feature that you might find useful to implement as part of Equinox. The idea is to allow ```error_if``` to consume functions, so...

knyazer

Pattern for lossless round trip of module serialization without storing side information required by __init__?

14

Woudn't it make sense to add something to wrap all this to make deserialization work? Currently it's lossy and you need to basically either define a init_param_from_params function or store...

cottrell

Second order updates

8

Hello Patrick, I hope you're doing well. I've been working with the Equinox package, and I appreciate the effort that's gone into it. Currently, I'm working on implementing a second-order...

BabaYara

Bayesian NNs with Equinox

2

This is really an exploratory question about implementing a Bayesian layer in Equinox. My research requires Bayesian neural networks, and I've been playing with Equinox to see how it stacks...

cdagher

Using functools.cached_property

1

Hello! I had a question about how Modules work under jit with the functools.cached_property decorator. I recently came across this decorator and thought it was pretty nifty. I have several...

mjo22

This is a fairly experimental idea. It'd be great to run tree_check automatically, to validate that there aren't any accidental repeated layers. In practice I wouldn't be surprised if there's...

patrick-kidger

LSTMCell and

4

The example for recurrent network at https://docs.kidger.site/equinox/api/nn/rnn/ shows that two cells just return final states of the cells but not the outputs.

YongchaoT-1QBit

`is_array` returns `True` for instances of `np.generic`

1

I started looking at PR for https://github.com/patrick-kidger/equinox/issues/486. While doing so, I noticed another potential issue: [`is_array`'s `isinstance` check includes `np.generic`](https://github.com/patrick-kidger/equinox/blob/76a13bd100a8127c7e5176f2dee9d36c837b46ec/equinox/_filters.py#L21). This does not seem to align with the intention for...

colehaus

"Full" `eqx.Module` serialisation

4

Dear @patrick-kidger, this project is awesome, thank you very much for it!! I am currently developing a fitting library for binned likelihood fits ([`dilax`](https://github.com/pfackeldey/dilax/tree/main)) in high energy physics based on...

pfackeldey