VAE example outdated

[Tim Bechtel]

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.

[Ivy Zheng]

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

[Tim Bechtel]

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?

[Ivy Zheng]

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: #3337 I might open a PR later to change the JAX doc.