<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
import itertools

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

import numpy as onp

from jax import api
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_array_shapes = [(), (4,), (3, 4), (3, 1), (1, 4), (2, 1, 4), (2, 3, 4)]
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]
int_dtypes = [onp.int32, onp.int64]
unsigned_dtypes = [onp.uint32, onp.uint64]
bool_dtypes = [onp.bool_]
default_dtypes = float_dtypes + int_dtypes
numeric_dtypes = float_dtypes + complex_dtypes + int_dtypes


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


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

JAX_ONE_TO_ONE_OP_RECORDS = [
    op_record("abs", 1, default_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
    op_record("add", 2, default_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
    op_record("ceil", 1, float_dtypes, all_shapes, jtu.rand_default(), []),
    op_record("conj", 1, numeric_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
    op_record("conjugate", 1, numeric_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
    op_record("equal", 2, default_dtypes, all_shapes, jtu.rand_some_equal(), []),
    #op_record("exp", 1, numeric_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
    op_record("floor", 1, float_dtypes, all_shapes, jtu.rand_default(), []),
    op_record("greater", 2, default_dtypes, all_shapes, jtu.rand_some_equal(), []),
    op_record("greater_equal", 2, default_dtypes, all_shapes, jtu.rand_some_equal(), []),
    op_record("less", 2, default_dtypes, all_shapes, jtu.rand_some_equal(), []),
    op_record("less_equal", 2, default_dtypes, all_shapes, jtu.rand_some_equal(), []),
    op_record("log", 1, numeric_dtypes, all_shapes, jtu.rand_positive(), ["rev"]),
    op_record("logical_and", 2, default_dtypes, all_shapes, jtu.rand_bool(), []),
    op_record("logical_not", 1, default_dtypes, all_shapes, jtu.rand_bool(), []),
    op_record("logical_or", 2, default_dtypes, all_shapes, jtu.rand_bool(), []),
    op_record("logical_xor", 2, default_dtypes, all_shapes, jtu.rand_bool(), []),
    op_record("maximum", 2, default_dtypes, all_shapes, jtu.rand_some_inf(), []),
    op_record("minimum", 2, default_dtypes, all_shapes, jtu.rand_some_inf(), []),
    op_record("multiply", 2, default_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
    op_record("negative", 1, default_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
    op_record("not_equal", 2, default_dtypes, all_shapes, jtu.rand_some_equal(), ["rev"]),
    #op_record("power", 2, float_dtypes, all_shapes, jtu.rand_positive(), ["rev"]),
    op_record("subtract", 2, default_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
    op_record("tanh", 1, numeric_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
    op_record("sin", 1, default_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
    op_record("cos", 1, default_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
]

JAX_COMPOUND_OP_RECORDS = [
    op_record("cosh", 1, default_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
    op_record("divide", 2, default_dtypes, all_shapes, jtu.rand_nonzero(), ["rev"]),
    op_record("expm1", 1, numeric_dtypes, all_shapes, jtu.rand_positive(), [],
              test_name="expm1_large"),
    op_record("expm1", 1, numeric_dtypes, all_shapes, jtu.rand_small_positive(), []),
    op_record("floor_divide", 2, default_dtypes, all_shapes, jtu.rand_nonzero(), ["rev"]),
    op_record("isclose", 2, float_dtypes, all_shapes, jtu.rand_small_positive(), []),
    op_record("log1p", 1, numeric_dtypes, all_shapes, jtu.rand_positive(), [],
              test_name="log1p_large"),
    op_record("log1p", 1, numeric_dtypes, all_shapes, jtu.rand_small_positive(), []),
    op_record("logaddexp", 2, float_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
    op_record("ravel", 1, default_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
    op_record("remainder", 2, default_dtypes, all_shapes, jtu.rand_nonzero(), []),
    op_record("sinh", 1, default_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
    op_record("sqrt", 1, default_dtypes, all_shapes, jtu.rand_positive(), ["rev"]),
    op_record("transpose", 1, default_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
    op_record("true_divide", 2, default_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("all", 1, bool_dtypes, all_shapes, jtu.rand_default(), []),
    op_record("any", 1, bool_dtypes, all_shapes, jtu.rand_default(), []),
    op_record("max", 1, default_dtypes, nonempty_shapes, jtu.rand_default(), []),
    #op_record("mean", 1, default_dtypes, all_shapes, jtu.rand_default(), []),
    op_record("min", 1, default_dtypes, nonempty_shapes, jtu.rand_default(), []),
    #op_record("prod", 1, default_dtypes, all_shapes, jtu.rand_small_positive(), []),
    op_record("sum", 1, default_dtypes, all_shapes, jtu.rand_default(), []),
    #op_record("var", 1, default_dtypes, all_shapes, jtu.rand_default(), []),
]

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

CombosWithReplacement = itertools.combinations_with_replacement


def _dtypes_are_compatible_for_bitwise_ops(args):
  if len(args) <= 1:
    return True
  is_signed = lambda dtype: onp.issubdtype(dtype, onp.signedinteger)
  width = lambda dtype: onp.iinfo(dtype).bits
  x, y = args
  if width(x) > width(y):
    x, y = y, x
  # The following condition seems a little ad hoc, but seems to capture what
  # numpy actually implements.
  return (
      is_signed(x) == is_signed(y)
      or (width(x) == 32 and width(y) == 32)<commit_msg>Re-enable tests that were accidentally commented out.<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
import itertools

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

import numpy as onp

from jax import api
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_array_shapes = [(), (4,), (3, 4), (3, 1), (1, 4), (2, 1, 4), (2, 3, 4)]
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]
int_dtypes = [onp.int32, onp.int64]
unsigned_dtypes = [onp.uint32, onp.uint64]
bool_dtypes = [onp.bool_]
default_dtypes = float_dtypes + int_dtypes
numeric_dtypes = float_dtypes + complex_dtypes + int_dtypes


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


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

JAX_ONE_TO_ONE_OP_RECORDS = [
    op_record("abs", 1, default_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
    op_record("add", 2, default_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
    op_record("ceil", 1, float_dtypes, all_shapes, jtu.rand_default(), []),
    op_record("conj", 1, numeric_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
    op_record("conjugate", 1, numeric_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
    op_record("equal", 2, default_dtypes, all_shapes, jtu.rand_some_equal(), []),
    op_record("exp", 1, numeric_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
    op_record("floor", 1, float_dtypes, all_shapes, jtu.rand_default(), []),
    op_record("greater", 2, default_dtypes, all_shapes, jtu.rand_some_equal(), []),
    op_record("greater_equal", 2, default_dtypes, all_shapes, jtu.rand_some_equal(), []),
    op_record("less", 2, default_dtypes, all_shapes, jtu.rand_some_equal(), []),
    op_record("less_equal", 2, default_dtypes, all_shapes, jtu.rand_some_equal(), []),
    op_record("log", 1, numeric_dtypes, all_shapes, jtu.rand_positive(), ["rev"]),
    op_record("logical_and", 2, default_dtypes, all_shapes, jtu.rand_bool(), []),
    op_record("logical_not", 1, default_dtypes, all_shapes, jtu.rand_bool(), []),
    op_record("logical_or", 2, default_dtypes, all_shapes, jtu.rand_bool(), []),
    op_record("logical_xor", 2, default_dtypes, all_shapes, jtu.rand_bool(), []),
    op_record("maximum", 2, default_dtypes, all_shapes, jtu.rand_some_inf(), []),
    op_record("minimum", 2, default_dtypes, all_shapes, jtu.rand_some_inf(), []),
    op_record("multiply", 2, default_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
    op_record("negative", 1, default_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
    op_record("not_equal", 2, default_dtypes, all_shapes, jtu.rand_some_equal(), ["rev"]),
    op_record("power", 2, float_dtypes, all_shapes, jtu.rand_positive(), ["rev"]),
    op_record("subtract", 2, default_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
    op_record("tanh", 1, numeric_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
    op_record("sin", 1, default_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
    op_record("cos", 1, default_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
]

JAX_COMPOUND_OP_RECORDS = [
    op_record("cosh", 1, default_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
    op_record("divide", 2, default_dtypes, all_shapes, jtu.rand_nonzero(), ["rev"]),
    op_record("expm1", 1, numeric_dtypes, all_shapes, jtu.rand_positive(), [],
              test_name="expm1_large"),
    op_record("expm1", 1, numeric_dtypes, all_shapes, jtu.rand_small_positive(), []),
    op_record("floor_divide", 2, default_dtypes, all_shapes, jtu.rand_nonzero(), ["rev"]),
    op_record("isclose", 2, float_dtypes, all_shapes, jtu.rand_small_positive(), []),
    op_record("log1p", 1, numeric_dtypes, all_shapes, jtu.rand_positive(), [],
              test_name="log1p_large"),
    op_record("log1p", 1, numeric_dtypes, all_shapes, jtu.rand_small_positive(), []),
    op_record("logaddexp", 2, float_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
    op_record("ravel", 1, default_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
    op_record("remainder", 2, default_dtypes, all_shapes, jtu.rand_nonzero(), []),
    op_record("sinh", 1, default_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
    op_record("sqrt", 1, default_dtypes, all_shapes, jtu.rand_positive(), ["rev"]),
    op_record("transpose", 1, default_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
    op_record("true_divide", 2, default_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("all", 1, bool_dtypes, all_shapes, jtu.rand_default(), []),
    op_record("any", 1, bool_dtypes, all_shapes, jtu.rand_default(), []),
    op_record("max", 1, default_dtypes, nonempty_shapes, jtu.rand_default(), []),
    op_record("mean", 1, default_dtypes, all_shapes, jtu.rand_default(), []),
    op_record("min", 1, default_dtypes, nonempty_shapes, jtu.rand_default(), []),
    op_record("prod", 1, default_dtypes, all_shapes, jtu.rand_small_positive(), []),
    op_record("sum", 1, default_dtypes, all_shapes, jtu.rand_default(), []),
    op_record("var", 1, default_dtypes, all_shapes, jtu.rand_default(), []),
]

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

CombosWithReplacement = itertools.combinations_with_replacement


def _dtypes_are_compatible_for_bitwise_ops(args):
  if len(args) <= 1:
    return True
  is_signed = lambda dtype: onp.issubdtype(dtype, onp.signedinteger)
  width = lambda dtype: onp.iinfo(dtype).bits
  x, y = args
  if width(x) > width(y):
    x, y = y, x
  # The following condition seems a little ad hoc, but seems to capture what
  # numpy actually implements.
  return (
      is_signed(x) == is_signed(y)
      or (width(x) == 32 and width(y) == 32)
{'fraction_non_alphanumeric': 0.15489614243323443, 'fraction_numerical': 0.014770853939993406, 'mean_word_length': 5.48951647411211, 'pattern_counts': {'":': 0, '<': 5, '<?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': '7356530', 'n_tokens_mistral': 5792, 'n_tokens_neox': 5552, 'n_words': 1185}