<commit_before># Copyright 2018 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# 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 absolute_import
from __future__ import division
from __future__ import print_function

import collections
import functools
from functools import partial
import itertools
from unittest import skip

from absl.testing import absltest
from absl.testing import parameterized
import six

import numpy as onp

from jax import api
from jax import lax
from jax import numpy as lnp
from jax import test_util as jtu

from jax.config import config
config.parse_flags_with_absl()
FLAGS = config.FLAGS

nonempty_nonscalar_array_shapes = [(4,), (3, 4), (3, 1), (1, 4), (2, 1, 4), (2, 3, 4)]
nonempty_array_shapes = [()] + nonempty_nonscalar_array_shapes
empty_array_shapes = [(0,), (0, 4), (3, 0),]

scalar_shapes = [jtu.NUMPY_SCALAR_SHAPE]
array_shapes = nonempty_array_shapes + empty_array_shapes
nonempty_shapes = scalar_shapes + nonempty_array_shapes
all_shapes =  scalar_shapes + array_shapes

float_dtypes = [onp.float32, onp.float64]
complex_dtypes = [onp.complex64, onp.complex128]
int_dtypes = [onp.int32, onp.int64]
unsigned_dtypes = [onp.uint32, onp.uint64]
bool_dtypes = [onp.bool_]
default_dtypes = float_dtypes + int_dtypes
inexact_dtypes = float_dtypes + complex_dtypes
number_dtypes = float_dtypes + complex_dtypes + int_dtypes
all_dtypes = number_dtypes + bool_dtypes

OpRecord = collections.namedtuple(
  "OpRecord",
  ["name", "nargs", "dtypes", "shapes", "rng", "diff_modes", "test_name",
   "check_dtypes"])


def op_record(name, nargs, dtypes, shapes, rng, diff_modes, test_name=None,
              check_dtypes=True):
  test_name = test_name or name
  return OpRecord(name, nargs, dtypes, shapes, rng, diff_modes, test_name,
                  check_dtypes)

JAX_ONE_TO_ONE_OP_RECORDS = [
    op_record("abs", 1, number_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
    op_record("add", 2, number_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
    op_record("ceil", 1, float_dtypes, all_shapes, jtu.rand_default(), []),
    op_record("conj", 1, number_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
    op_record("equal", 2, all_dtypes, all_shapes, jtu.rand_some_equal(), []),
    op_record("exp", 1, number_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
    op_record("fabs", 1, float_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
    op_record("floor", 1, float_dtypes, all_shapes, jtu.rand_default(), []),
    op_record("greater", 2, number_dtypes, all_shapes, jtu.rand_some_equal(), []),
    op_record("greater_equal", 2, number_dtypes, all_shapes, jtu.rand_some_equal(), []),
    op_record("isfinite", 1, number_dtypes, all_shapes, jtu.rand_some_inf(), []),
    op_record("less", 2, number_dtypes, all_shapes, jtu.rand_some_equal(), []),
    op_record("less_equal", 2, number_dtypes, all_shapes, jtu.rand_some_equal(), []),
    op_record("log", 1, number_dtypes, all_shapes, jtu.rand_positive(), ["rev"]),
    op_record("logical_and", 2, all_dtypes, all_shapes, jtu.rand_bool(), []),
    op_record("logical_not", 1, all_dtypes, all_shapes, jtu.rand_bool(), []),
    op_record("logical_or", 2, all_dtypes, all_shapes, jtu.rand_bool(), []),
    op_record("logical_xor", 2, all_dtypes, all_shapes, jtu.rand_bool(), []),
    op_record("maximum", 2, number_dtypes, all_shapes, jtu.rand_some_inf(), []),
    op_record("minimum", 2, number_dtypes, all_shapes, jtu.rand_some_inf(), []),
    op_record("multiply", 2, number_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
    op_record("negative", 1, number_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
    op_record("not_equal", 2, number_dtypes, all_shapes, jtu.rand_some_equal(), ["rev"]),
    op_record("power", 2, inexact_dtypes, all_shapes, jtu.rand_positive(), ["rev"]),
    op_record("reciprocal", 1, inexact_dtypes, all_shapes, jtu.rand_default(), []),
    op_record("subtract", 2, number_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
    op_record("sin", 1, number_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
    op_record("cos", 1, number_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
    op_record("tan", 1, number_dtypes, all_shapes, jtu.rand_uniform(-1.5, 1.5),
              ["rev"]),
    op_record("sinh", 1, number_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
    op_record("cosh", 1, number_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
    op_record("tanh", 1, number_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
    op_record("arcsin", 1, float_dtypes, all_shapes, jtu.rand_small(), ["rev"]),
    op_record("arccos", 1, float_dtypes, all_shapes, jtu.rand_small(), ["rev"]),
    op_record("arctan", 1, float_dtypes, all_shapes, jtu.rand_small(), ["rev"]),
    op_record("arctan2", 2, float_dtypes, all_shapes, jtu.rand_small(), ["rev"]),
    op_record("arcsinh", 1, number_dtypes, all_shapes, jtu.rand_positive(), ["rev"]),
    op_record("arccosh", 1, number_dtypes, all_shapes, jtu.rand_positive(), ["rev"]),
    op_record("arctanh", 1, number_dtypes, all_shapes, jtu.rand_small(), ["rev"]),
]

JAX_COMPOUND_OP_RECORDS = [
    # angle has inconsistent 32/64-bit return types across numpy versions.
    op_record("angle", 1, number_dtypes, all_shapes, jtu.rand_default(), [],
              check_dtypes=False),
    op_record("atleast_1d", 1, default_dtypes, all_shapes, jtu.rand_default(), []),
    op_record("atleast_2d", 1, default_dtypes, all_shapes, jtu.rand_default(), []),
    op_record("atleast_3d", 1, default_dtypes, all_shapes, jtu.rand_default(), []),
    op_record("conjugate", 1, number_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
    op_record("deg2rad", 1, float_dtypes, all_shapes, jtu.rand_default(), []),
    op_record("divide", 2, number_dtypes, all_shapes, jtu.rand_nonzero(), ["rev"]),
    op_record("exp2", 1, number_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
    op_record("expm1", 1, number_dtypes, all_shapes, jtu.rand_positive(), [],
              test_name="expm1_large"),
    op_record("expm1", 1, number_dtypes, all_shapes, jtu.rand_small_positive(), []),
    op_record("floor_divide", 2, number_dtypes, all_shapes, jtu.rand_nonzero(), ["rev"]),
    op_record("gcd", 2, int_dtypes, all_shapes, jtu.rand_default(), []),
    op_record("heaviside", 2, default_dtypes, all_shapes, jtu.rand_default(), []),
    op_record("hypot", 2, default_dtypes, all_shapes, jtu.rand_default(), []),
    op_record("kron", 2, number_dtypes, nonempty_shapes, jtu.rand_default(), []),
    op_record("outer", 2, number_dtypes, all_shapes, jtu.rand_default(), []),
    op_record("imag", 1, number_dtypes, all_shapes, jtu.rand_some_inf(), []),
    op_record("isclose", 2, all_dtypes, all_shapes, jtu.rand_small_positive(), []),
    op_record("iscomplex", 1, number_dtypes, all_shapes, jtu.rand_some_inf(), []),
    op_record("isreal", 1, number_dtypes, all_shapes, jtu.rand_some_inf(), []),
    op_record("lcm", 2, int_dtypes, all_shapes, jtu.rand_default(), []),
    op_record("log2", 1, number_dtypes, all_shapes, jtu.rand_positive(), ["rev"]),
    op_record("log10", 1, number_dtypes, all_shapes, jtu.rand_positive(), ["rev"]),
    op_record("log1p", 1, number_dtypes, all_shapes, jtu.rand_positive(), [],
              test_name="log1p_large"),
    op_record("log1p", 1, number_dtypes, all_shapes, jtu.rand_small_positive(), []),
    op_record("logaddexp", 2, float_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
    op_record("logaddexp2", 2, float_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
    op_record("polyval", 2, number_dtypes, nonempty_nonscalar_array_shapes, jtu.rand_default(), []),
    op_record("positive", 1, number_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
    op_record("rad2deg", 1, float_dtypes, all_shapes, jtu.rand_default(), []),
    op_record("ravel", 1, all_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
    op_record("real", 1, number_dtypes, all_shapes, jtu.rand_some_inf(), []),
    op_record("remainder", 2, default_dtypes, all_shapes, jtu.rand_nonzero(), []),
    op_record("mod", 2, default_dtypes, all_shapes, jtu.rand_nonzero(), []),
    op_record("sinc", 1, number_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
    op_record("square", 1, number_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
    op_record("sqrt", 1, number_dtypes, all_shapes, jtu.rand_positive(), ["rev"]),
    op_record("transpose", 1, all_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
    op_record("true_divide", 2, all_dtypes, all_shapes, jtu.rand_nonzero(), ["rev"]),
    op_record("where", 3, (onp.float32, onp.int64), all_shapes, jtu.rand_some_zero(), []),
]

JAX_BITWISE_OP_RECORDS = [
    op_record("bitwise_and", 2, int_dtypes + unsigned_dtypes, all_shapes,
              jtu.rand_bool(), []),
    op_record("bitwise_not", 1, int_dtypes + unsigned_dtypes, all_shapes,
              jtu.rand_bool(), []),
    op_record("bitwise_or", 2, int_dtypes + unsigned_dtypes, all_shapes,
              jtu.rand_bool(), []),
    op_record("bitwise_xor", 2, int_dtypes + unsigned_dtypes, all_shapes,
              jtu.rand_bool(), []),
]

JAX_REDUCER_RECORDS = [
    op_record("mean", 1, number_dtypes, nonempty_shapes, jtu.rand_default(), []),
    op_record("prod", 1, number_dtypes, all_shapes, jtu.rand_small_positive(), []),
    op_record("sum", 1, number_dtypes, all_shapes, jtu.rand_default(), []),
    op_record("var", 1, number_dtypes, nonempty_shapes, jtu.rand_default(), []),
    op_record("std", 1, inexact_dtypes, nonempty_shapes, jtu.rand_default(), []),
]

JAX_REDUCER_NO_DTYPE_RECORDS = [
    op_record("all", 1, all_dtypes, all_shapes, jtu.rand_some_zero(), []),
    op_record("any", 1, all_dtypes, all_shapes, jtu.rand_some_zero(), []),
    op_record("max", 1, all_dtypes, nonempty_shapes, jtu.rand_default(), []),
    op_record("min", 1, all_dtypes, nonempty_shapes, jtu.rand_default(), []),
]

numpy_version = tuple(map(int, onp.version.version.split('.')))
if numpy_version >= (1, 15):
  JAX_REDUCER_NO_DTYPE_RECORDS += [
      op_record("ptp", 1, number_dtypes, nonempty_shapes, jtu.rand_default(), []),
  ]


JAX_ARGMINMAX_RECORDS = [
    op_record("argmin", 1, all_dtypes, nonempty_shapes, jtu.rand_some_equal(), []),
    op_record("argmax", 1, all_dtypes, nonempty_shapes, jtu.rand_some_equal(), []),
]

JAX_OPERATOR_OVERLOADS = [
    op_record("__add__", 2, number_dtypes, all_shapes, jtu.rand_default(), []),
    op_record("__radd__", 2, number_dtypes, all_shapes, jtu.rand_default(), []),
    op_record("__sub__", 2, number_dtypes, all_shapes, jtu.rand_default(), []),
    op_record("__rsub__", 2, number_dtypes, all_shapes, jtu.rand_default(), []),
    op_record("__mul__", 2, number_dtypes, all_shapes, jtu.rand_default(), []),
    op_record("__rmul__", 2, number_dtypes, all_shapes, jtu.rand_default(), []),
    op_record("__eq__", 2, number_dtypes, all_shapes, jtu.rand_default(), []),
    op_record("__ne__", 2, number_dtypes, all_shapes, jtu.rand_default(), []),
    op_record("__lt__", 2, number_dtypes, all_shapes, jtu.rand_default(), []),
    op_record("__gt__", 2, number_dtypes, all_shapes, jtu.rand_default(), []),
    op_record("__ge__", 2, number_dtypes, all_shapes, jtu.rand_default(), []),
    op_record("__neg__", 1, number_dtypes, all_shapes, jtu.rand_default(), []),<commit_msg>guard lcm and gcd tests against numpy version<commit_after># Copyright 2018 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# 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 absolute_import
from __future__ import division
from __future__ import print_function

import collections
import functools
from functools import partial
import itertools
from unittest import skip

from absl.testing import absltest
from absl.testing import parameterized
import six

import numpy as onp

from jax import api
from jax import lax
from jax import numpy as lnp
from jax import test_util as jtu

from jax.config import config
config.parse_flags_with_absl()
FLAGS = config.FLAGS

nonempty_nonscalar_array_shapes = [(4,), (3, 4), (3, 1), (1, 4), (2, 1, 4), (2, 3, 4)]
nonempty_array_shapes = [()] + nonempty_nonscalar_array_shapes
empty_array_shapes = [(0,), (0, 4), (3, 0),]

scalar_shapes = [jtu.NUMPY_SCALAR_SHAPE]
array_shapes = nonempty_array_shapes + empty_array_shapes
nonempty_shapes = scalar_shapes + nonempty_array_shapes
all_shapes =  scalar_shapes + array_shapes

float_dtypes = [onp.float32, onp.float64]
complex_dtypes = [onp.complex64, onp.complex128]
int_dtypes = [onp.int32, onp.int64]
unsigned_dtypes = [onp.uint32, onp.uint64]
bool_dtypes = [onp.bool_]
default_dtypes = float_dtypes + int_dtypes
inexact_dtypes = float_dtypes + complex_dtypes
number_dtypes = float_dtypes + complex_dtypes + int_dtypes
all_dtypes = number_dtypes + bool_dtypes

OpRecord = collections.namedtuple(
  "OpRecord",
  ["name", "nargs", "dtypes", "shapes", "rng", "diff_modes", "test_name",
   "check_dtypes"])


def op_record(name, nargs, dtypes, shapes, rng, diff_modes, test_name=None,
              check_dtypes=True):
  test_name = test_name or name
  return OpRecord(name, nargs, dtypes, shapes, rng, diff_modes, test_name,
                  check_dtypes)

JAX_ONE_TO_ONE_OP_RECORDS = [
    op_record("abs", 1, number_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
    op_record("add", 2, number_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
    op_record("ceil", 1, float_dtypes, all_shapes, jtu.rand_default(), []),
    op_record("conj", 1, number_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
    op_record("equal", 2, all_dtypes, all_shapes, jtu.rand_some_equal(), []),
    op_record("exp", 1, number_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
    op_record("fabs", 1, float_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
    op_record("floor", 1, float_dtypes, all_shapes, jtu.rand_default(), []),
    op_record("greater", 2, number_dtypes, all_shapes, jtu.rand_some_equal(), []),
    op_record("greater_equal", 2, number_dtypes, all_shapes, jtu.rand_some_equal(), []),
    op_record("isfinite", 1, number_dtypes, all_shapes, jtu.rand_some_inf(), []),
    op_record("less", 2, number_dtypes, all_shapes, jtu.rand_some_equal(), []),
    op_record("less_equal", 2, number_dtypes, all_shapes, jtu.rand_some_equal(), []),
    op_record("log", 1, number_dtypes, all_shapes, jtu.rand_positive(), ["rev"]),
    op_record("logical_and", 2, all_dtypes, all_shapes, jtu.rand_bool(), []),
    op_record("logical_not", 1, all_dtypes, all_shapes, jtu.rand_bool(), []),
    op_record("logical_or", 2, all_dtypes, all_shapes, jtu.rand_bool(), []),
    op_record("logical_xor", 2, all_dtypes, all_shapes, jtu.rand_bool(), []),
    op_record("maximum", 2, number_dtypes, all_shapes, jtu.rand_some_inf(), []),
    op_record("minimum", 2, number_dtypes, all_shapes, jtu.rand_some_inf(), []),
    op_record("multiply", 2, number_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
    op_record("negative", 1, number_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
    op_record("not_equal", 2, number_dtypes, all_shapes, jtu.rand_some_equal(), ["rev"]),
    op_record("power", 2, inexact_dtypes, all_shapes, jtu.rand_positive(), ["rev"]),
    op_record("reciprocal", 1, inexact_dtypes, all_shapes, jtu.rand_default(), []),
    op_record("subtract", 2, number_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
    op_record("sin", 1, number_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
    op_record("cos", 1, number_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
    op_record("tan", 1, number_dtypes, all_shapes, jtu.rand_uniform(-1.5, 1.5),
              ["rev"]),
    op_record("sinh", 1, number_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
    op_record("cosh", 1, number_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
    op_record("tanh", 1, number_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
    op_record("arcsin", 1, float_dtypes, all_shapes, jtu.rand_small(), ["rev"]),
    op_record("arccos", 1, float_dtypes, all_shapes, jtu.rand_small(), ["rev"]),
    op_record("arctan", 1, float_dtypes, all_shapes, jtu.rand_small(), ["rev"]),
    op_record("arctan2", 2, float_dtypes, all_shapes, jtu.rand_small(), ["rev"]),
    op_record("arcsinh", 1, number_dtypes, all_shapes, jtu.rand_positive(), ["rev"]),
    op_record("arccosh", 1, number_dtypes, all_shapes, jtu.rand_positive(), ["rev"]),
    op_record("arctanh", 1, number_dtypes, all_shapes, jtu.rand_small(), ["rev"]),
]

JAX_COMPOUND_OP_RECORDS = [
    # angle has inconsistent 32/64-bit return types across numpy versions.
    op_record("angle", 1, number_dtypes, all_shapes, jtu.rand_default(), [],
              check_dtypes=False),
    op_record("atleast_1d", 1, default_dtypes, all_shapes, jtu.rand_default(), []),
    op_record("atleast_2d", 1, default_dtypes, all_shapes, jtu.rand_default(), []),
    op_record("atleast_3d", 1, default_dtypes, all_shapes, jtu.rand_default(), []),
    op_record("conjugate", 1, number_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
    op_record("deg2rad", 1, float_dtypes, all_shapes, jtu.rand_default(), []),
    op_record("divide", 2, number_dtypes, all_shapes, jtu.rand_nonzero(), ["rev"]),
    op_record("exp2", 1, number_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
    op_record("expm1", 1, number_dtypes, all_shapes, jtu.rand_positive(), [],
              test_name="expm1_large"),
    op_record("expm1", 1, number_dtypes, all_shapes, jtu.rand_small_positive(), []),
    op_record("floor_divide", 2, number_dtypes, all_shapes, jtu.rand_nonzero(), ["rev"]),
    op_record("heaviside", 2, default_dtypes, all_shapes, jtu.rand_default(), []),
    op_record("hypot", 2, default_dtypes, all_shapes, jtu.rand_default(), []),
    op_record("kron", 2, number_dtypes, nonempty_shapes, jtu.rand_default(), []),
    op_record("outer", 2, number_dtypes, all_shapes, jtu.rand_default(), []),
    op_record("imag", 1, number_dtypes, all_shapes, jtu.rand_some_inf(), []),
    op_record("isclose", 2, all_dtypes, all_shapes, jtu.rand_small_positive(), []),
    op_record("iscomplex", 1, number_dtypes, all_shapes, jtu.rand_some_inf(), []),
    op_record("isreal", 1, number_dtypes, all_shapes, jtu.rand_some_inf(), []),
    op_record("log2", 1, number_dtypes, all_shapes, jtu.rand_positive(), ["rev"]),
    op_record("log10", 1, number_dtypes, all_shapes, jtu.rand_positive(), ["rev"]),
    op_record("log1p", 1, number_dtypes, all_shapes, jtu.rand_positive(), [],
              test_name="log1p_large"),
    op_record("log1p", 1, number_dtypes, all_shapes, jtu.rand_small_positive(), []),
    op_record("logaddexp", 2, float_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
    op_record("logaddexp2", 2, float_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
    op_record("polyval", 2, number_dtypes, nonempty_nonscalar_array_shapes, jtu.rand_default(), []),
    op_record("positive", 1, number_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
    op_record("rad2deg", 1, float_dtypes, all_shapes, jtu.rand_default(), []),
    op_record("ravel", 1, all_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
    op_record("real", 1, number_dtypes, all_shapes, jtu.rand_some_inf(), []),
    op_record("remainder", 2, default_dtypes, all_shapes, jtu.rand_nonzero(), []),
    op_record("mod", 2, default_dtypes, all_shapes, jtu.rand_nonzero(), []),
    op_record("sinc", 1, number_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
    op_record("square", 1, number_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
    op_record("sqrt", 1, number_dtypes, all_shapes, jtu.rand_positive(), ["rev"]),
    op_record("transpose", 1, all_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
    op_record("true_divide", 2, all_dtypes, all_shapes, jtu.rand_nonzero(), ["rev"]),
    op_record("where", 3, (onp.float32, onp.int64), all_shapes, jtu.rand_some_zero(), []),
]

JAX_BITWISE_OP_RECORDS = [
    op_record("bitwise_and", 2, int_dtypes + unsigned_dtypes, all_shapes,
              jtu.rand_bool(), []),
    op_record("bitwise_not", 1, int_dtypes + unsigned_dtypes, all_shapes,
              jtu.rand_bool(), []),
    op_record("bitwise_or", 2, int_dtypes + unsigned_dtypes, all_shapes,
              jtu.rand_bool(), []),
    op_record("bitwise_xor", 2, int_dtypes + unsigned_dtypes, all_shapes,
              jtu.rand_bool(), []),
]

JAX_REDUCER_RECORDS = [
    op_record("mean", 1, number_dtypes, nonempty_shapes, jtu.rand_default(), []),
    op_record("prod", 1, number_dtypes, all_shapes, jtu.rand_small_positive(), []),
    op_record("sum", 1, number_dtypes, all_shapes, jtu.rand_default(), []),
    op_record("var", 1, number_dtypes, nonempty_shapes, jtu.rand_default(), []),
    op_record("std", 1, inexact_dtypes, nonempty_shapes, jtu.rand_default(), []),
]

JAX_REDUCER_NO_DTYPE_RECORDS = [
    op_record("all", 1, all_dtypes, all_shapes, jtu.rand_some_zero(), []),
    op_record("any", 1, all_dtypes, all_shapes, jtu.rand_some_zero(), []),
    op_record("max", 1, all_dtypes, nonempty_shapes, jtu.rand_default(), []),
    op_record("min", 1, all_dtypes, nonempty_shapes, jtu.rand_default(), []),
]

numpy_version = tuple(map(int, onp.version.version.split('.')))
if numpy_version >= (1, 15):
  JAX_REDUCER_NO_DTYPE_RECORDS += [
      op_record("gcd", 2, int_dtypes, all_shapes, jtu.rand_default(), []),
      op_record("lcm", 2, int_dtypes, all_shapes, jtu.rand_default(), []),
      op_record("ptp", 1, number_dtypes, nonempty_shapes, jtu.rand_default(), []),
  ]


JAX_ARGMINMAX_RECORDS = [
    op_record("argmin", 1, all_dtypes, nonempty_shapes, jtu.rand_some_equal(), []),
    op_record("argmax", 1, all_dtypes, nonempty_shapes, jtu.rand_some_equal(), []),
]

JAX_OPERATOR_OVERLOADS = [
    op_record("__add__", 2, number_dtypes, all_shapes, jtu.rand_default(), []),
    op_record("__radd__", 2, number_dtypes, all_shapes, jtu.rand_default(), []),
    op_record("__sub__", 2, number_dtypes, all_shapes, jtu.rand_default(), []),
    op_record("__rsub__", 2, number_dtypes, all_shapes, jtu.rand_default(), []),
    op_record("__mul__", 2, number_dtypes, all_shapes, jtu.rand_default(), []),
    op_record("__rmul__", 2, number_dtypes, all_shapes, jtu.rand_default(), []),
    op_record("__eq__", 2, number_dtypes, all_shapes, jtu.rand_default(), []),
    op_record("__ne__", 2, number_dtypes, all_shapes, jtu.rand_default(), []),
    op_record("__lt__", 2, number_dtypes, all_shapes, jtu.rand_default(), []),
    op_record("__gt__", 2, number_dtypes, all_shapes, jtu.rand_default(), []),
    op_record("__ge__", 2, number_dtypes, all_shapes, jtu.rand_default(), []),
    op_record("__neg__", 1, number_dtypes, all_shapes, jtu.rand_default(), []),
{'fraction_non_alphanumeric': 0.16639672662177138, 'fraction_numerical': 0.01551444889608729, 'mean_word_length': 5.665625, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 2, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7213824', 'n_tokens_mistral': 9202, 'n_tokens_neox': 8837, 'n_words': 1707}