What's the fastest way to stack pytrees?

[NeilGirdhar]

I was forced to convert a jax.lax.scan into an ordinary loop that produces a list of pytrees. I want to stack them, but it is incredibly slow:

        def stack(*elements: Array) -> Array:
            return jnp.stack(elements)

        with Timer() as timer:
            trajectory = tree_map(stack, *results)  # results has 5000 trees in it
        log.info(f"Stacking took {timer.elapsed_str(3)}")  # Stacking took 71.2 s!!!

What's the fastest way to accomplish this in Jax?

[NeilGirdhar]

Switching to np.stack brought the runtime down to 2s!

[YouJiacheng]

But np.stack cannot jit?

[NeilGirdhar]

@YouJiacheng True, but if I could use the JIT, I would also be using lax.scan.