AttributeError: module 'jax' has no attribute 'flatten_util'. In jax.version 0.6.0

[mathliuqiang]

Hi, I'm using JAX version 0.6.0 and I've noticed that flatten_util seems to be missing, as I'm getting the error: AttributeError: module 'jax' has no attribute 'flatten_util'.

I need the functionality provided by ravel_pytree(), which I believe was part of flatten_util. Could you please let me know how I can achieve the ravel_pytree() functionality in JAX 0.6.0? Is there an alternative or a new location for this function?

Thanks!

[jakevdp]

You need to explicitly import the submodule:

import jax
print(jax.__version__)  # 0.6.0
import jax.flatten_util
print(jax.flatten_util.ravel_pytree) # <function ravel_pytree at 0x10c307ec0>

[mathliuqiang]

Thank you very very much!

[mathliuqiang]

3x

[jakevdp]

If you have a new question, it’s best to open a new discussion rather than post it as an answer to another question.

[mathliuqiang]

Thanks!