jax

255c3030 - Fix a bug where treedef.flatten_up_to(...) was overly permissive for None treedefs.

Fix a bug where treedef.flatten_up_to(...) was overly permissive for None treedefs. For example, tree_map(..., None, [2, 3]) did not raise an error, but None is a container and only leaves can be considered tree prefixes in this case. PiperOrigin-RevId: 674019460