jax

d45fa224 - Add tests to cover `PyTreeDef.flatten_up_to` error scenarios.

Add tests to cover `PyTreeDef.flatten_up_to` error scenarios. Also improve coverage of `PyTreeDef.flatten_up_to` success scenarios. PiperOrigin-RevId: 570152827