Remove 'from __future__ import annotations'

jax/_src/abstract_arrays.py

@@ -12,8 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from __future__ import annotations
-
 from functools import partial
 
 import numpy as np

jax/_src/ad_checkpoint.py

@@ -12,8 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from __future__ import annotations
-
 from collections.abc import Callable, Sequence
 import functools
 from functools import partial

jax/_src/ad_util.py

@@ -11,7 +11,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from __future__ import annotations
 
 from collections.abc import Callable
 import types

jax/_src/api.py

@@ -20,7 +20,6 @@
 tree_util.py), which include nested tuples/lists/dicts, where the leaves are
 arrays.
 """
-from __future__ import annotations
 
 import collections
 from collections.abc import Callable, Generator, Hashable, Iterable, Sequence

jax/_src/api_util.py

@@ -12,8 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from __future__ import annotations
-
 from collections.abc import Callable, Iterable, Sequence
 import inspect
 import operator

jax/_src/array.py

@@ -12,8 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from __future__ import annotations
-
 from collections import defaultdict
 from collections.abc import Callable, Sequence
 import enum

jax/_src/basearray.py

@@ -14,8 +14,6 @@
 
 # Note that type annotations for this file are defined in basearray.pyi
 
-from __future__ import annotations
-
 import abc
 import numpy as np
 from typing import Any, Union

jax/_src/callback.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Module for JAX callbacks."""
-from __future__ import annotations
 
 from collections.abc import Callable, Sequence
 import dataclasses

jax/_src/checkify.py

@@ -11,7 +11,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from __future__ import annotations
 
 from collections.abc import Callable, Sequence
 import dataclasses

jax/_src/clusters/cloud_tpu_cluster.py

@@ -12,8 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from __future__ import annotations
-
 import logging
 import os
 import re

jax/_src/clusters/cluster.py

@@ -12,8 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from __future__ import annotations
-
 from collections.abc import Sequence
 import logging
 from jax._src.cloud_tpu_init import running_in_cloud_tpu_vm

jax/_src/clusters/ompi_cluster.py

@@ -12,8 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from __future__ import annotations
-
 import os
 import re
 from jax._src import clusters

jax/_src/clusters/slurm_cluster.py

@@ -12,8 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from __future__ import annotations
-
 import os
 from jax._src import clusters
 

jax/_src/compilation_cache.py

@@ -12,8 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from __future__ import annotations
-
 import logging
 import threading
 import warnings

jax/_src/compiler.py

@@ -14,8 +14,6 @@
 
 # Interface to the compiler
 
-from __future__ import annotations
-
 from collections.abc import Sequence
 import logging
 import os

jax/_src/compute_on.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from __future__ import annotations
 import threading
 from contextlib import contextmanager
 

jax/_src/config.py

@@ -12,8 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from __future__ import annotations
-
 from collections.abc import Callable, Hashable, Iterator, Sequence
 import contextlib
 import functools

jax/_src/core.py

@@ -11,7 +11,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from __future__ import annotations
 
 from collections import Counter, defaultdict, deque, namedtuple
 from collections.abc import (Callable, Collection, Generator, Hashable,

jax/_src/custom_batching.py

@@ -12,8 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from __future__ import annotations
-
 from collections.abc import Callable
 import functools
 import operator

jax/_src/custom_derivatives.py

@@ -12,8 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from __future__ import annotations
-
 from collections.abc import Callable, Sequence
 import dataclasses
 from functools import update_wrapper, reduce, partial

jax/_src/custom_transpose.py

@@ -12,8 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from __future__ import annotations
-
 from collections.abc import Callable
 import functools
 from typing import Any

jax/_src/debugger/cli_debugger.py

@@ -11,7 +11,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from __future__ import annotations
 
 import cmd
 import pprint

jax/_src/debugger/colab_debugger.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Module for Colab-specific debugger."""
-from __future__ import annotations
 
 import html
 import inspect

jax/_src/debugger/colab_lib.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Module for building interfaces in Colab."""
-from __future__ import annotations
 
 import abc
 import dataclasses

jax/_src/debugger/core.py

@@ -11,7 +11,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from __future__ import annotations
 
 from collections.abc import Hashable
 import dataclasses

jax/_src/debugger/web_debugger.py

@@ -11,7 +11,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from __future__ import annotations
 
 import atexit
 import functools

jax/_src/debugging.py

@@ -13,8 +13,6 @@
 # limitations under the License.
 """Module for JAX debugging primitives and related functionality."""
 
-from __future__ import annotations
-
 import importlib.util
 from collections.abc import Callable, Sequence
 import functools

jax/_src/dispatch.py

@@ -13,7 +13,6 @@
 # limitations under the License.
 
 # Primitive dispatch and jit dispatch.
-from __future__ import annotations
 
 import atexit
 from collections.abc import Callable, Iterator, Sequence

jax/_src/distributed.py

@@ -12,8 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from __future__ import annotations
-
 import atexit
 from collections.abc import Sequence
 import logging

jax/_src/dlpack.py

@@ -12,8 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from __future__ import annotations
-
 from typing import Any
 
 from jax._src.api import device_put

jax/_src/dtypes.py

@@ -19,8 +19,6 @@
 # b) the set of supported types (e.g., bfloat16),
 # so we need our own implementation that deviates from NumPy in places.
 
-from __future__ import annotations
-
 import abc
 import builtins
 import functools

jax/_src/earray.py

@@ -12,8 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from __future__ import annotations
-
 import math
 
 from jax._src import api_util

jax/_src/effects.py

@@ -50,8 +50,6 @@
 https://jax.readthedocs.io/en/latest/jep/10657-sequencing-effects.html.
 """
 
-from __future__ import annotations
-
 from collections.abc import Iterable, Set
 from typing import Any
 

jax/_src/environment_info.py

@@ -12,8 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from __future__ import annotations
-
 import platform
 import subprocess
 import sys

jax/_src/errors.py

@@ -11,7 +11,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from __future__ import annotations
 
 from jax._src import core
 from jax._src.util import set_module

jax/_src/export/_export.py

@@ -15,8 +15,6 @@
 
 """
 
-from __future__ import annotations
-
 from collections.abc import Callable, Sequence
 import copy
 import dataclasses

jax/_src/export/serialization.py

@@ -14,8 +14,6 @@
 
 # Serialization and deserialization of _export.Exported
 
-from __future__ import annotations
-
 from collections.abc import Callable, Sequence
 from functools import partial
 from typing import TypeVar

jax/_src/export/shape_poly.py

@@ -16,8 +16,6 @@
 See documentation at https://jax.readthedocs.io/en/latest/export/shape_poly.html.
 """
 
-from __future__ import annotations
-
 import enum
 from collections.abc import Callable, Sequence
 import dataclasses

jax/_src/export/shape_poly_decision.py

@@ -15,8 +15,6 @@
 
 """
 
-from __future__ import annotations
-
 from collections.abc import Sequence
 import itertools
 import math

jax/_src/extend/ffi.py

@@ -12,8 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from __future__ import annotations
-
 import os
 import ctypes
 from collections.abc import Iterable, Mapping, Sequence

jax/_src/image/scale.py

@@ -12,8 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from __future__ import annotations
-
 from collections.abc import Callable, Sequence
 from functools import partial
 import enum

jax/_src/internal_test_util/export_back_compat_test_util.py

@@ -68,8 +68,6 @@ def func(...): ...
 
 """
 
-from __future__ import annotations
-
 from collections.abc import Callable, Iterable, Sequence
 import dataclasses
 import datetime

jax/_src/internal_test_util/lax_test_util.py

@@ -17,8 +17,6 @@
 # only, and may be changed or removed at any time and without any deprecation
 # cycle.
 
-from __future__ import annotations
-
 import collections
 import itertools
 from typing import Union, cast

jax/_src/internal_test_util/test_harnesses.py

@@ -36,8 +36,6 @@
 to fail. A Limitation is specific to a harness.
 """
 
-from __future__ import annotations
-
 from collections.abc import Callable, Iterable, Sequence
 import operator
 import os

jax/_src/interpreters/ad.py

@@ -12,8 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from __future__ import annotations
-
 from collections.abc import Callable, Sequence
 import contextlib
 import functools

jax/_src/interpreters/batching.py

@@ -11,7 +11,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from __future__ import annotations
 
 import collections
 from collections.abc import Callable, Iterable, Sequence