revise axes_scan to flatten argument pytrees only once

A user has a custom pytree node with the unusual behavior that it introduces
new arrays when flattening. That is, it's as if we had:

```python
# a custom object with two leaf arrays
custom_tree_object = SomeObject(jax_arrray1, jax_array2)
# convert leaves to ShapedArrays
custom_tree_object2 = jax.tree.map(core.typeof, custom_tree_object)
# flatten, should only see ShapedArrays, right?
leaves, treedef = jax.tree.flatten(custom_tree_object2)
print(leaves)
# [ShapedArray(...), ShapedArray(...), np.array(...)]
```

This change makes the `flax.nn.scan` function robust to such behavior. Without it, we were passing non-AbstractValues into JAX where JAX required AbstractValues.

I don't think we want to support this in general, but this fix seemed like the most
expedient way to roll fowrard jax-ml/jax#29273

PiperOrigin-RevId: 768175118

Author: mattjj

flax/core/axes_scan.py

@@ -147,15 +147,16 @@ def body_fn(c, xs, init_mode=False):
 
     broadcast_body = functools.partial(body_fn, init_mode=True)
 
-    carry_avals = jax.tree_util.tree_map(
-        lambda x: core.ShapedArray(jnp.shape(x), jnp.result_type(x)), init
-    )
-    scan_avals = jax.tree_util.tree_map(
-        lambda x: core.ShapedArray(jnp.shape(x)[1:], jnp.result_type(x)), xs
-    )
-    input_avals = (carry_avals, scan_avals)
+    init_flat, carry_tree = jax.tree.flatten(init)
+    xs_flat, scan_tree = jax.tree.flatten(xs)
+    carry_avals = [core.ShapedArray(jnp.shape(x), jnp.result_type(x))
+                   for x in init_flat]
+    scan_avals = [core.ShapedArray(jnp.shape(x)[1:], jnp.result_type(x))
+                  for x in xs_flat]
+    in_avals = [*carry_avals, *scan_avals]
+    in_tree = jax.tree_util.treedef_tuple((carry_tree, scan_tree))
+    assert all(isinstance(a, core.AbstractValue) for a in in_avals), in_avals
 
-    in_avals, in_tree = jax.tree_util.tree_flatten(input_avals)
     debug_info = jax.api_util.debug_info("flax scan", broadcast_body,
                                          (in_tree,), {})
     f_flat, out_tree = jax.api_util.flatten_fun_nokwargs(