Increase dataclass test coverage of the dm_tree API
Increase dataclass test coverage of the dm_tree API

Add tests for `tree.flatten_up_to`, `tree.flatten_with_path`,
`tree.flatten_with_path_up_to`, `tree.map_structure_up_to`,
`tree.map_structure_with_path`, `tree.map_structure_with_path_up_to`

PiperOrigin-RevId: 379449163

deepmind/chex,deepmind/chex