Orbax import breaks logging. Again.

google / orbax

Orbax provides common checkpointing and persistence utilities for JAX users

https://orbax.readthedocs.io/

Apache License 2.0

290 stars 29 forks source link

Orbax import breaks logging. Again. #1248

Open qGentry opened 1 hour ago

qGentry commented 1 hour ago

Basically the same issue I've submitted a while ago here https://github.com/google/orbax/issues/699 Even the same repro works :) Tested on orbax 0.7.0, while 0.5.3 (where fix for this issue was implemented) doesn't break logging

Maybe it is possible to add some tests for this specific case?

qGentry commented 1 hour ago

@niketkumar cc

qGentry commented 35 minutes ago

Looks like issue is this log as initialization of ArrayHandler happens on orbax import: https://github.com/google/orbax/blob/main/checkpoint/orbax/checkpoint/type_handlers.py#L1117 Commenting it our fixes the issue.