how to assert two frozendicts' equality?

[franchesoni]

assert paramsp.tree_flatten() == params.tree_flatten()

does not work, how should I do this?

[matthias-wright]

Assuming all your leaves are jax arrays, you could do this:

import jax.numpy as jnp
import numpy as np

assert np.all(jax.tree_leaves(jax.tree_multimap(lambda x, y: jnp.all(jnp.equal(x, y)), paramsp, params)))

Hope that helps!

[marcvanzee]

Or slightly shorter, I think this also works:

import jax.numpy as jnp
from jax.tree_util import tree_all

assert tree_all(jax.tree_multimap(lambda x, y: (x==y).all(), paramsp, params))