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)...

Question: States, `vmap`-construction-of-layers and `lax.scan`'ing.

7

Hello! I'm having some trouble with states and the construction of a model that `filter_vmap`-constructs layers. I've checked your FAQ's and stateful-ops page, but I can't get this model working...