Replace obsolete import contextlib2 with the equivalent stdlib module contextlib

PiperOrigin-RevId: 767547534

Author: oprypin

sonnet/src/BUILD

@@ -195,7 +195,6 @@ snt_py_library(
     srcs = ["custom_getter.py"],
     deps = [
         ":base",
-        # pip: contextlib2
         # pip: tensorflow
         # pip: tree
     ],
@@ -637,7 +636,6 @@ snt_py_library(
     deps = [
         ":custom_getter",
         ":utils",
-        # pip: contextlib2
         # pip: tensorflow
         # pip: tree
     ],

sonnet/src/custom_getter.py

@@ -14,9 +14,9 @@
 # ============================================================================
 """Custom getter for module members."""
 
+import contextlib
 from typing import Any, Callable, ContextManager, Iterable, Optional, Type
 
-import contextlib
 from sonnet.src import base
 import tensorflow as tf
 import tree

sonnet/src/distribute/BUILD

@@ -35,7 +35,6 @@ snt_py_library(
     srcs = ["replicator.py"],
     deps = [
         # pip: absl/logging
-        # pip: contextlib2
         "//sonnet/src:initializers",
         # pip: tensorflow
     ],

sonnet/src/distribute/replicator.py

@@ -14,10 +14,10 @@
 # ============================================================================
 """Replicator Distribution Strategy."""
 
+import contextlib
 from typing import Callable, TypeVar
 
 from absl import logging
-import contextlib
 from sonnet.src import initializers
 import tensorflow as tf
 

sonnet/src/functional/BUILD

@@ -9,7 +9,6 @@ snt_py_library(
     srcs = ["haiku.py"],
     deps = [
         ":utils",
-        # pip: contextlib2
         # pip: tensorflow
     ],
 )

sonnet/src/functional/haiku.py

@@ -15,11 +15,11 @@
 """Implements part of the Haiku ("Sonnet for JAX") API in TensorFlow 2."""
 
 import collections
+import contextlib
 import functools
 import itertools
 import threading
 
-import contextlib
 from sonnet.src.functional import utils
 import tensorflow as tf
 

sonnet/src/mixed_precision.py

@@ -14,8 +14,8 @@
 # ============================================================================
 """Mixed Precision Decorator for Sonnet 2."""
 
-import uuid
 import contextlib
+import uuid
 
 from sonnet.src import custom_getter
 from sonnet.src import utils