jnp.stack() OOM on TPUv3-8

[AakashKumarNain]

I am using Equinox for one of my workflows. All Equniox modules are just Pytrees. I have a list of these Pytree modules like this:

layers = [some_module for key in keys]  # This works

The total parameters count in layers is 6979584000 where each param is float16 dtype. When I try to stack this into a one large array, I get OOM error on TPUv3-8.

layers = [some_module for key in keys] 
layers = jax.tree_util.tree_map(lambda *x: jnp.stack(x), *layers) # Throws OOM

Can someone provide any pointers to why is this happening?

[jakevdp]

I'm not sure what you mean by "total parameters count": to clarify, how many arrays do you have, and how many elements are in each?

[AakashKumarNain]

I apologize for the vagueness. I will rephrase it in a better way with an example