stanford-crfm / haliax

Named Tensors for Legible Deep Learning in JAX
Apache License 2.0
141 stars 9 forks source link

Equinox woes: declared class attributes #3

Open dlwh opened 1 year ago

dlwh commented 1 year ago

from JohnH:

Why can’t I set attributes that weren’t declared as class attributes? I probably should have realized this but at this point am just going through and adding them…

This is a constraint in equinox but actually pretty unnecessary, except for declaring static fields. We should maybe raise to Patrick or just fork away from equinox... It wouldn't be hard to make our own Module and do our own thing. (I also think he treats the precise behavior of tree_flatten in modules as internal behavior, but we rely on it...)