Fix Array defaults for struct.PyTreeNode

[timoklein]

I've been working some on this and made a little progress at least: I'm fetching the default values for all jax.Arrays and setting a field with default_factory= lambda: default.

The problem is now that Python only remembers the last value for default from the loop, and thus will initialize all other Array attributes erroneously to this value.

[timoklein]

@epignatelli The code now passes all tests. The solution can probably be made more elegant.

What's the best way to test/benchmark that no funny business occurs during compilation?

I tried running some scripts under test/performance, but they threw errors for me, e.g.,

TypeError: scan body function carry input and carry output must have the same pytree structure, but they differ: The input carry component c.info is a <class 'dict'> with 0 child but the corresponding component of the carry output is a <class 'dict'> with 1 child, so the numbers of children do not match, with the symmetric difference of key sets: {'return'}.

when running python tests/performance/grid.py.