VAE example outdated - Githubissues

google / flax

Flax is a neural network library for JAX that is designed for flexibility.

https://flax.readthedocs.io

Apache License 2.0

6.06k stars 639 forks source link

VAE example outdated #3666

Open tisabe opened 8 months ago

tisabe commented 8 months ago

The VAE example in /flax/examples/vae seems to have some outdated requirements, jax.random usage, and absl.flags usage.

Modifications I made to get the example running:

change jax.random.key to jax.random.PRNGKey
install clu==0.0.10
run main.py without workdir argument (change README to not use --workdir)

System information

jax==0.4.6
flax==0.6.9
other requirements installed from requirements.txt

I'd be happy to open a PR with these changes.

IvyZX commented 8 months ago

jax.random.PRNGKey is a deprecated version of jax.random.key, so the usage of latter is preferred. You are using a very old JAX version that still uses PRNGKey. And it looks like clu was constrained to be below 0.0.10 by https://github.com/google/flax/pull/3394 - maybe this is to fix some CI test? @cgarciae And yes the --workdir argument seems to be mistakenly copied - a PR fixing that would be perfect.

tisabe commented 8 months ago

As far as I can see jax.random.PRNGKey is still referenced in Pseudo Random Numbers in JAX. I also updated my JAX version to 0.4.23, and this had no problem with jax.random.PRNGKey, however, I can't find current documentation for jax.random.key, did I miss this in my search?

IvyZX commented 8 months ago

Ah, JAX documentation is a bit outdated... One of the JAX owners actually wrote the PR that converted all Flax examples to use jax.random.key: https://github.com/google/flax/pull/3337 I might open a PR later to change the JAX doc.