source sync

PiperOrigin-RevId: 222452709

jax/abstract_arrays.py

@@ -13,6 +13,8 @@
 # limitations under the License.
 
 from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
 
 import numpy as onp
 import six

jax/ad_util.py

@@ -13,6 +13,8 @@
 # limitations under the License.
 
 from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
 
 from .core import JaxTuple, lattice_join
 from .interpreters.partial_eval import Primitive

jax/api.py

@@ -13,6 +13,8 @@
 # limitations under the License.
 
 from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
 
 import itertools
 

jax/api_util.py

@@ -13,6 +13,8 @@
 # limitations under the License.
 
 from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
 
 from .core import pack
 from .tree_util import build_tree, process_pytree

jax/core.py

@@ -13,6 +13,8 @@
 # limitations under the License.
 
 from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
 
 from operator import attrgetter
 from contextlib import contextmanager

jax/experimental/lapax.py

@@ -13,7 +13,9 @@
 # limitations under the License.
 
 """A linear algebra library for use with JAX."""
+
 from __future__ import absolute_import
+from __future__ import division
 from __future__ import print_function
 
 import numpy as onp

jax/interpreters/ad.py

@@ -13,6 +13,8 @@
 # limitations under the License.
 
 from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
 
 from . import partial_eval as pe
 from . import xla

jax/interpreters/batching.py

@@ -12,6 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
 from collections import namedtuple
 
 import itertools as it

jax/interpreters/partial_eval.py

@@ -13,6 +13,8 @@
 # limitations under the License.
 
 from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
 
 import itertools as it
 from collections import namedtuple, Counter, defaultdict

jax/interpreters/xla.py

@@ -13,6 +13,8 @@
 # limitations under the License.
 
 from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
 
 from collections import namedtuple
 import itertools as it

jax/linear_util.py

@@ -13,6 +13,8 @@
 # limitations under the License.
 
 from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
 
 from .util import curry, partial
 

jax/pprint_util.py

@@ -12,6 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
 from six.moves import reduce
 
 

jax/test_util.py

@@ -13,6 +13,8 @@
 # limitations under the License.
 
 from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
 
 import functools
 import re

jax/tree_util.py

@@ -13,6 +13,8 @@
 # limitations under the License.
 
 from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
 
 from collections import namedtuple
 import itertools as it

jax/util.py

@@ -13,6 +13,8 @@
 # limitations under the License.
 
 from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
 
 import functools
 import itertools as it

tests/batching_test.py

@@ -13,6 +13,8 @@
 # limitations under the License.
 
 from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
 
 import numpy as onp
 from absl.testing import absltest

tests/lax_numpy_indexing_test.py

@@ -12,6 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
 import collections
 from functools import partial
 import itertools