Jed Homer

Extension to 3D Wavelet Transform?

5

Hi, Is there any future work planned toward allowing 3D DWT transforms? E.g. on an array of shape `(batch, channels, height, width, depth) `. I'd be very interested in working...

Analog to `flax.struct.dataclass`

2

Hi `chex` team, Is there any potential for something like `flax.struct.dataclass` in `chex`? Basically a kind of `dataclass` that can mark static arguments. Two other variations include [`jax_dataclasses`](https://github.com/brentyi/jax_dataclasses) and [`simple-pytree`](https://github.com/cgarciae/simple-pytree/blob/main/simple_pytree/pytree.py)...