VAE example outdated

[tisabe]

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]

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]

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]

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.