Correct logging imports

checkpoint/orbax/checkpoint/_src/metadata/sharding.py

@@ -20,7 +20,7 @@
 import dataclasses
 import enum
 import json
-import logging
+from absl import logging
 from typing import Any, List, Mapping, Optional, Sequence, Tuple, Union
 
 import jax

checkpoint/orbax/checkpoint/experimental/emergency/test_utils/dataset_iterator_checkpoint_handler.py

@@ -15,7 +15,7 @@
 """This module provides a PyGrain CheckpointHandler for integration with Orbax."""
 import dataclasses
 import json
-import logging
+from absl import logging
 from typing import Any, Optional
 from etils import epath
 import jax

checkpoint/orbax/checkpoint/experimental/v1/_src/layout/orbax_v0_layout.py

@@ -15,7 +15,7 @@
 """Defines `V0Layout`, a class to handle Orbax V0 checkpoint formats."""
 
 import asyncio
-import logging
+from absl import logging
 from typing import Any, Awaitable
 
 from orbax.checkpoint._src import asyncio_utils

export/orbax/export/modules/obm_module.py

@@ -16,7 +16,7 @@
 
 from collections.abc import Callable, Mapping, Sequence
 import copy
-import logging
+from absl import logging
 from typing import Any, Optional, Union
 import warnings
 

export/orbax/export/obm_configs.py

@@ -18,7 +18,7 @@
 import dataclasses
 import enum
 import itertools
-import logging
+from absl import logging
 import jax
 from jax import export as jax_export