lax_numpy_test.py

# 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.


import collections
import copy
import functools
from functools import partial
import inspect
import io
import itertools
import operator
from typing import cast, Iterator, Optional, List, Tuple
import unittest
from unittest import SkipTest
import warnings

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

import numpy as np
try:
  import numpy_dispatch
except ImportError:
  numpy_dispatch = None

import jax
import jax.ops
from jax import lax
from jax import numpy as jnp
from jax import tree_util
from jax.test_util import check_grads

from jax._src import device_array
from jax._src import dtypes
from jax._src import test_util as jtu
from jax._src.lax import lax as lax_internal
from jax._src.numpy.lax_numpy import _promote_dtypes, _promote_dtypes_inexact
from jax._src.numpy.util import _parse_numpydoc, ParsedDoc, _wraps
from jax._src.util import prod, safe_zip

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

numpy_version = tuple(map(int, np.__version__.split('.')[:3]))

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

scalar_shapes = [jtu.NUMPY_SCALAR_SHAPE, jtu.PYTHON_SCALAR_SHAPE]
array_shapes = nonempty_array_shapes + empty_array_shapes
nonzerodim_shapes = nonempty_nonscalar_array_shapes + empty_array_shapes
nonempty_shapes = scalar_shapes + nonempty_array_shapes
all_shapes = scalar_shapes + array_shapes

float_dtypes = jtu.dtypes.all_floating
complex_dtypes = jtu.dtypes.complex
int_dtypes = jtu.dtypes.all_integer
unsigned_dtypes = jtu.dtypes.all_unsigned
bool_dtypes = jtu.dtypes.boolean
default_dtypes = float_dtypes + int_dtypes
inexact_dtypes = float_dtypes + complex_dtypes
number_dtypes = float_dtypes + complex_dtypes + int_dtypes + unsigned_dtypes
all_dtypes = number_dtypes + bool_dtypes


python_scalar_dtypes = [jnp.bool_, jnp.int_, jnp.float_, jnp.complex_]

# uint64 is problematic because with any uint type it promotes to float:
int_dtypes_no_uint64 = [d for d in int_dtypes + unsigned_dtypes if d != np.uint64]

def _indexer_with_default_outputs(indexer, use_defaults=True):
  """Like jtu.with_jax_dtype_defaults, but for __getitem__ APIs"""
  class Indexer:
    @partial(jtu.with_jax_dtype_defaults, use_defaults=use_defaults)
    def __getitem__(self, *args):
      return indexer.__getitem__(*args)
  return Indexer()

def _valid_dtypes_for_shape(shape, dtypes):
  # Not all (shape, dtype) pairs are valid. In particular, Python scalars only
  # have one type in each category (float, bool, etc.)
  if shape is jtu.PYTHON_SCALAR_SHAPE:
    return [t for t in dtypes if t in python_scalar_dtypes]
  return dtypes

def _shape_and_dtypes(shapes, dtypes):
  for shape in shapes:
    for dtype in _valid_dtypes_for_shape(shape, dtypes):
      yield (shape, dtype)

def _compatible_shapes(shape):
  if shape in scalar_shapes or np.ndim(shape) == 0:
    return [shape]
  return (shape[n:] for n in range(len(shape) + 1))

def _get_y_shapes(y_dtype, shape, rowvar):
  # Helper function for testCov.
  if y_dtype is None:
    return [None]
  if len(shape) == 1:
    return [shape]
  elif rowvar or shape[0] == 1:
    return [(1, shape[-1]), (2, shape[-1]), (5, shape[-1])]
  return [(shape[0], 1), (shape[0], 2), (shape[0], 5)]

OpRecord = collections.namedtuple(
  "OpRecord",
  ["name", "nargs", "dtypes", "shapes", "rng_factory", "diff_modes",
   "test_name", "check_dtypes", "tolerance", "inexact", "kwargs"])

def op_record(name, nargs, dtypes, shapes, rng_factory, diff_modes,
              test_name=None, check_dtypes=True,
              tolerance=None, inexact=False, kwargs=None):
  test_name = test_name or name
  return OpRecord(name, nargs, dtypes, shapes, rng_factory, diff_modes,
                  test_name, check_dtypes, tolerance, inexact, kwargs)

JAX_ONE_TO_ONE_OP_RECORDS = [
    op_record("abs", 1, all_dtypes,
              all_shapes, jtu.rand_default, ["rev"]),
    op_record("add", 2, all_dtypes, all_shapes, jtu.rand_default, ["rev"]),
    op_record("ceil", 1, float_dtypes, all_shapes, jtu.rand_default, []),
    op_record("ceil", 1, int_dtypes + unsigned_dtypes, all_shapes,
              jtu.rand_default, [], check_dtypes=False),
    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"],
              inexact=True),
    op_record("fabs", 1, float_dtypes, all_shapes, jtu.rand_default, ["rev"]),
    op_record("float_power", 2, inexact_dtypes, all_shapes,
              partial(jtu.rand_default, scale=1), ["rev"],
              tolerance={jnp.bfloat16: 1e-2, np.float32: 1e-3,
                         np.float64: 1e-12, np.complex64: 2e-4,
                         np.complex128: 1e-12}, check_dtypes=False),
    op_record("floor", 1, float_dtypes, all_shapes, jtu.rand_default, []),
    op_record("floor", 1, int_dtypes + unsigned_dtypes, all_shapes,
              jtu.rand_default, [], check_dtypes=False),
    op_record("greater", 2, all_dtypes, all_shapes, jtu.rand_some_equal, []),
    op_record("greater_equal", 2, all_dtypes, all_shapes, jtu.rand_some_equal, []),
    op_record("i0", 1, float_dtypes, all_shapes, jtu.rand_default, [],
              check_dtypes=False),
    op_record("ldexp", 2, int_dtypes, all_shapes, jtu.rand_default, [], check_dtypes=False),
    op_record("less", 2, all_dtypes, all_shapes, jtu.rand_some_equal, []),
    op_record("less_equal", 2, all_dtypes, all_shapes, jtu.rand_some_equal, []),
    op_record("log", 1, number_dtypes, all_shapes, jtu.rand_positive, ["rev"],
              inexact=True),
    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, all_dtypes, all_shapes, jtu.rand_some_inf, []),
    op_record("minimum", 2, all_dtypes, all_shapes, jtu.rand_some_inf, []),
    op_record("multiply", 2, all_dtypes, all_shapes, jtu.rand_default, ["rev"]),
    op_record("negative", 1, number_dtypes, all_shapes, jtu.rand_default, ["rev"]),
    op_record("nextafter", 2, [f for f in float_dtypes if f != jnp.bfloat16],
              all_shapes, jtu.rand_default, ["rev"], inexact=True, tolerance=0),
    op_record("not_equal", 2, all_dtypes, all_shapes, jtu.rand_some_equal, ["rev"]),
    op_record("array_equal", 2, number_dtypes, all_shapes, jtu.rand_some_equal, ["rev"]),
    op_record("array_equiv", 2, number_dtypes, all_shapes, jtu.rand_some_equal, ["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("signbit", 1, default_dtypes + bool_dtypes, all_shapes,
              jtu.rand_some_inf_and_nan, ["rev"]),
    op_record("trunc", 1, float_dtypes, all_shapes, jtu.rand_some_inf_and_nan, []),
    op_record("trunc", 1, int_dtypes + unsigned_dtypes, all_shapes,
              jtu.rand_some_inf_and_nan, [], check_dtypes=False),
    op_record("sin", 1, number_dtypes, all_shapes, jtu.rand_default, ["rev"],
              inexact=True),
    op_record("cos", 1, number_dtypes, all_shapes, jtu.rand_default, ["rev"],
              inexact=True),
    op_record("tan", 1, number_dtypes, all_shapes,
              partial(jtu.rand_uniform, low=-1.5, high=1.5), ["rev"],
              inexact=True),
    op_record("sinh", 1, number_dtypes, all_shapes, jtu.rand_default, ["rev"],
              inexact=True),
    op_record("cosh", 1, number_dtypes, all_shapes, jtu.rand_default, ["rev"],
              inexact=True),
    # TODO(b/142975473): on CPU, tanh for complex128 is only accurate to
    # ~float32 precision.
    # TODO(b/143135720): on GPU, tanh has only ~float32 precision.
    op_record("tanh", 1, number_dtypes, all_shapes, jtu.rand_default, ["rev"],
              tolerance={np.float64: 1e-7, np.complex128: 1e-7},
              inexact=True),
    op_record("arcsin", 1, number_dtypes, all_shapes, jtu.rand_small, ["rev"],
              inexact=True),
    op_record("arccos", 1, number_dtypes, all_shapes, jtu.rand_small, ["rev"],
              inexact=True),
    op_record("arctan", 1, number_dtypes, all_shapes, jtu.rand_small, ["rev"],
              inexact=True),
    op_record("arctan2", 2, float_dtypes, all_shapes, jtu.rand_small, ["rev"],
              inexact=True),
    op_record("arcsinh", 1, number_dtypes, all_shapes, jtu.rand_default, ["rev"],
              inexact=True, tolerance={np.complex64: 2E-4, np.complex128: 2E-14}),
    op_record("arccosh", 1, number_dtypes, all_shapes, jtu.rand_default, ["rev"],
              inexact=True, tolerance={np.complex64: 2E-2, np.complex128: 2E-12}),
    op_record("arctanh", 1, number_dtypes, all_shapes, jtu.rand_small, ["rev"],
              inexact=True, tolerance={np.float64: 1e-9}),
]

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, inexact=True),
    op_record("angle", 1, number_dtypes, all_shapes, jtu.rand_default, [],
              check_dtypes=False, inexact=True, test_name="angle_deg", kwargs={'deg': True}),
    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("cbrt", 1, default_dtypes, all_shapes, jtu.rand_some_inf, ["rev"],
              inexact=True),
    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"],
              inexact=True),
    op_record("divmod", 2, int_dtypes + float_dtypes, all_shapes,
              jtu.rand_nonzero, []),
    op_record("exp2", 1, number_dtypes, all_shapes, jtu.rand_default, ["rev"],
              tolerance={jnp.bfloat16: 4e-2, np.float16: 1e-2}, inexact=True),
    # TODO(b/142975473): on CPU, expm1 for float64 is only accurate to ~float32
    # precision.
    op_record("expm1", 1, number_dtypes, all_shapes, jtu.rand_positive, [],
              test_name="expm1_large", tolerance={np.float64: 1e-8}, inexact=True),
    op_record("expm1", 1, number_dtypes, all_shapes, jtu.rand_small_positive,
              [], tolerance={np.float64: 1e-8}, inexact=True),
    op_record("fix", 1, float_dtypes, all_shapes, jtu.rand_default, []),
    op_record("fix", 1, int_dtypes + unsigned_dtypes, all_shapes,
              jtu.rand_default, [], check_dtypes=False),
    op_record("floor_divide", 2, default_dtypes + unsigned_dtypes,
              all_shapes, jtu.rand_nonzero, ["rev"]),
    op_record("fmin", 2, number_dtypes, all_shapes, jtu.rand_some_nan, []),
    op_record("fmax", 2, number_dtypes, all_shapes, jtu.rand_some_nan, []),
    op_record("fmod", 2, default_dtypes, all_shapes, jtu.rand_some_nan, []),
    op_record("heaviside", 2, default_dtypes, all_shapes, jtu.rand_default, [],
              inexact=True),
    op_record("hypot", 2, default_dtypes, all_shapes, jtu.rand_default, [],
              inexact=True),
    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("iscomplex", 1, number_dtypes, all_shapes, jtu.rand_some_inf, []),
    op_record("isfinite", 1, inexact_dtypes, all_shapes, jtu.rand_some_inf_and_nan, []),
    op_record("isinf", 1, inexact_dtypes, all_shapes, jtu.rand_some_inf_and_nan, []),
    op_record("isnan", 1, inexact_dtypes, all_shapes, jtu.rand_some_inf_and_nan, []),
    op_record("isneginf", 1, float_dtypes, all_shapes, jtu.rand_some_inf_and_nan, []),
    op_record("isposinf", 1, float_dtypes, all_shapes, jtu.rand_some_inf_and_nan, []),
    op_record("isreal", 1, number_dtypes, all_shapes, jtu.rand_some_inf, []),
    op_record("isrealobj", 1, number_dtypes, all_shapes, jtu.rand_some_inf, []),
    op_record("log2", 1, number_dtypes, all_shapes, jtu.rand_positive, ["rev"],
              inexact=True),
    op_record("log10", 1, number_dtypes, all_shapes, jtu.rand_positive, ["rev"],
              inexact=True),
    op_record("log1p", 1, number_dtypes, all_shapes, jtu.rand_positive, [],
              test_name="log1p_large", tolerance={np.float64: 1e-12},
              inexact=True),
    op_record("log1p", 1, number_dtypes, all_shapes, jtu.rand_small_positive, [],
              tolerance={np.float64: 1e-12}, inexact=True),
    op_record("logaddexp", 2, float_dtypes, all_shapes,
              jtu.rand_some_inf_and_nan, ["rev"],
              tolerance={np.float64: 1e-12}, inexact=True),
    op_record("logaddexp2", 2, float_dtypes, all_shapes,
              jtu.rand_some_inf_and_nan, ["rev"],
              tolerance={np.float16: 1e-2, np.float64: 2e-14}, inexact=True),
    op_record("polyval", 2, number_dtypes, nonempty_nonscalar_array_shapes,
              jtu.rand_default, [], check_dtypes=False,
              tolerance={dtypes.bfloat16: 4e-2, np.float16: 1e-2,
                         np.float64: 1e-12}),
    op_record("positive", 1, number_dtypes, all_shapes, jtu.rand_default, ["rev"]),
    op_record("power", 2, number_dtypes, all_shapes, jtu.rand_positive, ["rev"],
              tolerance={np.complex128: 1e-14}, check_dtypes=False),
    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, [],
              tolerance={np.float16: 1e-2}),
    op_record("mod", 2, default_dtypes, all_shapes, jtu.rand_nonzero, []),
    op_record("modf", 1, float_dtypes, all_shapes, jtu.rand_default, []),
    op_record("modf", 1, int_dtypes + unsigned_dtypes, all_shapes,
              jtu.rand_default, [], check_dtypes=False),
    op_record("rint", 1, inexact_dtypes, all_shapes, jtu.rand_some_inf_and_nan,
              []),
    op_record("rint", 1, int_dtypes + unsigned_dtypes, all_shapes,
              jtu.rand_default, [], check_dtypes=False),
    op_record("sign", 1, number_dtypes, all_shapes, jtu.rand_some_inf_and_nan, []),
    # numpy 1.16 has trouble mixing uint and bfloat16, so we test these separately.
    op_record("copysign", 2, default_dtypes + unsigned_dtypes,
              all_shapes, jtu.rand_some_inf_and_nan, [], check_dtypes=False),
    op_record("sinc", 1, [t for t in number_dtypes if t != jnp.bfloat16],
              all_shapes, jtu.rand_default, ["rev"],
              tolerance={np.complex64: 1e-5}, inexact=True,
              check_dtypes=False),
    op_record("square", 1, number_dtypes, all_shapes, jtu.rand_default, ["rev"]),
    op_record("sqrt", 1, number_dtypes, all_shapes, jtu.rand_positive, ["rev"],
              inexact=True),
    op_record("transpose", 1, all_dtypes, all_shapes, jtu.rand_default, ["rev"],
              check_dtypes=False),
    op_record("true_divide", 2, all_dtypes, all_shapes, jtu.rand_nonzero,
              ["rev"], inexact=True),
    op_record("ediff1d", 3, [np.int32], all_shapes, jtu.rand_default, [], check_dtypes=False),
    # TODO(phawkins): np.unwrap does not correctly promote its default period
    # argument under NumPy 1.21 for bfloat16 inputs. It works fine if we
    # explicitly pass a bfloat16 value that does not need promition. We should
    # probably add a custom test harness for unwrap that tests the period
    # argument anyway.
    op_record("unwrap", 1, [t for t in float_dtypes if t != dtypes.bfloat16],
              nonempty_nonscalar_array_shapes,
              jtu.rand_default, ["rev"],
              # numpy.unwrap always returns float64
              check_dtypes=False,
              # numpy cumsum is inaccurate, see issue #3517
              tolerance={dtypes.bfloat16: 1e-1, np.float16: 1e-1}),
    op_record("isclose", 2, [t for t in all_dtypes if t != jnp.bfloat16],
              all_shapes, jtu.rand_small_positive, []),
    op_record("gcd", 2, int_dtypes_no_uint64, all_shapes, jtu.rand_default, []),
    op_record("lcm", 2, int_dtypes_no_uint64, all_shapes, jtu.rand_default, []),
]

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("invert", 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, [],
              inexact=True),
    op_record("prod", 1, all_dtypes, all_shapes, jtu.rand_small_positive, []),
    op_record("sum", 1, all_dtypes, all_shapes, jtu.rand_default, []),
    op_record("nanmean", 1, inexact_dtypes, nonempty_shapes, jtu.rand_some_nan,
              [], inexact=True),
    op_record("nanprod", 1, all_dtypes, all_shapes, jtu.rand_some_nan, []),
    op_record("nansum", 1, number_dtypes, all_shapes, jtu.rand_some_nan, []),
]

JAX_REDUCER_INITIAL_RECORDS = [
    op_record("prod", 1, all_dtypes, all_shapes, jtu.rand_small_positive, []),
    op_record("sum", 1, all_dtypes, all_shapes, jtu.rand_default, []),
    op_record("max", 1, all_dtypes, all_shapes, jtu.rand_default, []),
    op_record("min", 1, all_dtypes, all_shapes, jtu.rand_default, []),
]
if numpy_version >= (1, 22):  # initial & where keywords added in numpy 1.22
  JAX_REDUCER_INITIAL_RECORDS += [
      op_record("nanprod", 1, inexact_dtypes, all_shapes, jtu.rand_small_positive, []),
      op_record("nansum", 1, inexact_dtypes, all_shapes, jtu.rand_default, []),
      op_record("nanmax", 1, inexact_dtypes, all_shapes, jtu.rand_default, []),
      op_record("nanmin", 1, inexact_dtypes, all_shapes, jtu.rand_default, []),
  ]

JAX_REDUCER_WHERE_NO_INITIAL_RECORDS = [
    op_record("all", 1, bool_dtypes, all_shapes, jtu.rand_some_zero, []),
    op_record("any", 1, bool_dtypes, all_shapes, jtu.rand_some_zero, []),
    op_record("mean", 1, all_dtypes, nonempty_shapes, jtu.rand_default, [],
              inexact=True),
    op_record("var", 1, all_dtypes, nonempty_shapes, jtu.rand_default, [],
              inexact=True),
    op_record("std", 1, all_dtypes, nonempty_shapes, jtu.rand_default, [],
              inexact=True),
]
if numpy_version >= (1, 22):  # where keyword added in numpy 1.22
  JAX_REDUCER_WHERE_NO_INITIAL_RECORDS += [
      op_record("nanmean", 1, inexact_dtypes, nonempty_shapes, jtu.rand_default, [],
                inexact=True),
      op_record("nanvar", 1, inexact_dtypes, nonempty_shapes, jtu.rand_default, [],
                inexact=True),
      op_record("nanstd", 1, inexact_dtypes, nonempty_shapes, jtu.rand_default, [],
                inexact=True),
  ]

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, []),
    op_record("var", 1, all_dtypes, nonempty_shapes, jtu.rand_default, [],
              inexact=True),
    op_record("std", 1, all_dtypes, nonempty_shapes, jtu.rand_default, [],
              inexact=True),
    op_record("nanmax", 1, all_dtypes, nonempty_shapes, jtu.rand_some_nan, []),
    op_record("nanmin", 1, all_dtypes, nonempty_shapes, jtu.rand_some_nan, []),
    op_record("nanvar", 1, all_dtypes, nonempty_shapes, jtu.rand_some_nan,
              [], inexact=True),
    op_record("nanstd", 1, all_dtypes, nonempty_shapes, jtu.rand_some_nan,
              [], inexact=True),
    op_record("ptp", 1, number_dtypes, nonempty_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, []),
    op_record("nanargmin", 1, default_dtypes, nonempty_shapes, jtu.rand_some_nan, []),
    op_record("nanargmax", 1, default_dtypes, nonempty_shapes, jtu.rand_some_nan, []),
]

JAX_OPERATOR_OVERLOADS = [
    op_record("__add__", 2, number_dtypes, all_shapes, jtu.rand_default, []),
    op_record("__sub__", 2, number_dtypes, all_shapes, jtu.rand_default, []),
    op_record("__mul__", 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, default_dtypes, all_shapes, jtu.rand_default, []),
    op_record("__le__", 2, default_dtypes, all_shapes, jtu.rand_default, []),
    op_record("__gt__", 2, default_dtypes, all_shapes, jtu.rand_default, []),
    op_record("__ge__", 2, default_dtypes, all_shapes, jtu.rand_default, []),
    op_record("__pos__", 1, number_dtypes, all_shapes, jtu.rand_default, []),
    op_record("__neg__", 1, number_dtypes, all_shapes, jtu.rand_default, []),
    op_record("__pow__", 2, inexact_dtypes, all_shapes, jtu.rand_positive, [],
              tolerance={np.float32: 2e-4, np.complex64: 2e-4, np.complex128: 1e-14}),
    op_record("__mod__", 2, default_dtypes, all_shapes, jtu.rand_nonzero, [],
              tolerance={np.float16: 1e-1}),
    op_record("__floordiv__", 2, default_dtypes, all_shapes,
              jtu.rand_nonzero, []),
    op_record("__truediv__", 2, number_dtypes, all_shapes, jtu.rand_nonzero, [],
              inexact=True),
    op_record("__abs__", 1, number_dtypes, all_shapes, jtu.rand_default, []),
    # TODO(mattjj): __invert__ fails on bool dtypes because ~True == -2
    op_record("__invert__", 1, int_dtypes, all_shapes, jtu.rand_default, []),
    # TODO(mattjj): investigate these failures
    # op_record("__or__", 2, number_dtypes, all_shapes, jtu.rand_bool, []),
    # op_record("__and__", 2, number_dtypes, all_shapes, jtu.rand_default, []),
    # op_record("__xor__", 2, number_dtypes, all_shapes, jtu.rand_bool, []),
    # op_record("__divmod__", 2, number_dtypes, all_shapes, jtu.rand_nonzero, []),
    op_record("__lshift__", 2, int_dtypes_no_uint64, all_shapes, partial(jtu.rand_int, high=8), []),
    op_record("__rshift__", 2, int_dtypes_no_uint64, all_shapes, partial(jtu.rand_int, high=8), []),
]

JAX_RIGHT_OPERATOR_OVERLOADS = [
    op_record("__radd__", 2, number_dtypes, all_shapes, jtu.rand_default, []),
    op_record("__rsub__", 2, number_dtypes, all_shapes, jtu.rand_default, []),
    op_record("__rmul__", 2, number_dtypes, all_shapes, jtu.rand_default, []),
    op_record("__rpow__", 2, inexact_dtypes, all_shapes, jtu.rand_positive, [],
              tolerance={np.float32: 2e-4, np.complex64: 1e-3}),
    op_record("__rmod__", 2, default_dtypes, all_shapes, jtu.rand_nonzero, [],
              tolerance={np.float16: 1e-1}),
    op_record("__rfloordiv__", 2, default_dtypes, all_shapes,
              jtu.rand_nonzero, []),
    op_record("__rtruediv__", 2, number_dtypes, all_shapes, jtu.rand_nonzero, [],
              inexact=True),
    # op_record("__ror__", 2, number_dtypes, all_shapes, jtu.rand_bool, []),
    # op_record("__rand__", 2, number_dtypes, all_shapes, jtu.rand_default, []),
    # op_record("__rxor__", 2, number_dtypes, all_shapes, jtu.rand_bool, []),
    # op_record("__rdivmod__", 2, number_dtypes, all_shapes, jtu.rand_nonzero, []),
    op_record("__rlshift__", 2, int_dtypes_no_uint64, all_shapes, partial(jtu.rand_int, high=8), []),
    op_record("__rrshift__", 2, int_dtypes_no_uint64, all_shapes, partial(jtu.rand_int, high=8), [])
]

class _OverrideEverything:
  pass

for rec in JAX_OPERATOR_OVERLOADS + JAX_RIGHT_OPERATOR_OVERLOADS:
  if rec.nargs == 2:
    setattr(_OverrideEverything, rec.name, lambda self, other: self)

class _OverrideNothing:
  pass

for rec in JAX_OPERATOR_OVERLOADS + JAX_RIGHT_OPERATOR_OVERLOADS:
  if rec.nargs == 2:
    setattr(_OverrideNothing, rec.name, lambda self, other: NotImplemented)


def _dtypes_are_compatible_for_bitwise_ops(args):
  if len(args) <= 1:
    return True
  is_signed = lambda dtype: jnp.issubdtype(dtype, np.signedinteger)
  width = lambda dtype: jnp.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)
      or (width(x) == 32 and width(y) == 64 and is_signed(y)))

def _shapes_are_broadcast_compatible(shapes):
  try:
    lax.broadcast_shapes(*(() if s in scalar_shapes else s for s in shapes))
  except ValueError:
    return False
  else:
    return True

def _shapes_are_equal_length(shapes):
  return all(len(shape) == len(shapes[0]) for shape in shapes[1:])


def _promote_like_jnp(fun, inexact=False):
  """Decorator that promotes the arguments of `fun` to `jnp.result_type(*args)`.

  jnp and np have different type promotion semantics; this decorator allows
  tests make an np reference implementation act more like an jnp
  implementation.
  """
  _promote = _promote_dtypes_inexact if inexact else _promote_dtypes
  def wrapper(*args, **kw):
    flat_args, tree = tree_util.tree_flatten(args)
    args = tree_util.tree_unflatten(tree, _promote(*flat_args))
    return fun(*args, **kw)
  return wrapper


class LaxBackedNumpyTests(jtu.JaxTestCase):
  """Tests for LAX-backed Numpy implementation."""

  def _GetArgsMaker(self, rng, shapes, dtypes, np_arrays=True):
    def f():
      out = [rng(shape, dtype or jnp.float_)
             for shape, dtype in zip(shapes, dtypes)]
      if np_arrays:
        return out
      return [jnp.asarray(a) if isinstance(a, (np.ndarray, np.generic)) else a
              for a in out]
    return f

  def testNotImplemented(self):
    for name in jnp._NOT_IMPLEMENTED:
      func = getattr(jnp, name)
      with self.assertRaises(NotImplementedError):
        func()

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": f"_{dtype}_allow_picke={allow_pickle}",
       "dtype": dtype, "allow_pickle": allow_pickle}
      for dtype in float_dtypes + [object]
      for allow_pickle in [True, False]))
  def testLoad(self, dtype, allow_pickle):
    if dtype == object and not allow_pickle:
      self.skipTest("dtype=object requires allow_pickle=True")
    rng = jtu.rand_default(self.rng())
    arr = rng((10), dtype)
    with io.BytesIO() as f:
      jnp.save(f, arr)
      f.seek(0)
      arr_out = jnp.load(f, allow_pickle=allow_pickle)
    self.assertArraysEqual(arr, arr_out)

  @parameterized.named_parameters(itertools.chain.from_iterable(
      jtu.cases_from_list(
        {"testcase_name": jtu.format_test_name_suffix(rec.test_name, shapes,
                                                      dtypes),
         "rng_factory": rec.rng_factory, "shapes": shapes, "dtypes": dtypes,
         "np_op": getattr(np, rec.name), "jnp_op": getattr(jnp, rec.name),
         "check_dtypes": rec.check_dtypes, "tolerance": rec.tolerance,
         "inexact": rec.inexact, "kwargs": rec.kwargs or {}}
        for shapes in filter(
          _shapes_are_broadcast_compatible,
          itertools.combinations_with_replacement(rec.shapes, rec.nargs))
        for dtypes in itertools.product(
          *(_valid_dtypes_for_shape(s, rec.dtypes) for s in shapes)))
      for rec in itertools.chain(JAX_ONE_TO_ONE_OP_RECORDS,
                                 JAX_COMPOUND_OP_RECORDS)))
  @jax.numpy_rank_promotion('allow')  # This test explicitly exercises implicit rank promotion.
  def testOp(self, np_op, jnp_op, rng_factory, shapes, dtypes, check_dtypes,
             tolerance, inexact, kwargs):
    np_op = partial(np_op, **kwargs)
    jnp_op = partial(jnp_op, **kwargs)
    np_op = jtu.ignore_warning(category=RuntimeWarning,
                               message="invalid value.*")(np_op)
    np_op = jtu.ignore_warning(category=RuntimeWarning,
                               message="divide by zero.*")(np_op)

    rng = rng_factory(self.rng())
    args_maker = self._GetArgsMaker(rng, shapes, dtypes, np_arrays=False)
    tol = max(jtu.tolerance(dtype, tolerance) for dtype in dtypes)
    tol = functools.reduce(jtu.join_tolerance,
                           [tolerance, tol, jtu.default_tolerance()])
    self._CheckAgainstNumpy(_promote_like_jnp(np_op, inexact), jnp_op,
                            args_maker, check_dtypes=check_dtypes, tol=tol)
    self._CompileAndCheck(jnp_op, args_maker, check_dtypes=check_dtypes,
                          atol=tol, rtol=tol)

  @parameterized.named_parameters(itertools.chain.from_iterable(
      jtu.cases_from_list(
        {"testcase_name": jtu.format_test_name_suffix(rec.test_name, shapes,
                                                      dtypes),
         "rng_factory": rec.rng_factory, "shapes": shapes, "dtypes": dtypes, "name": rec.name,
         "tol": rec.tolerance}
        for shapes in filter(
          _shapes_are_broadcast_compatible,
          itertools.combinations_with_replacement(rec.shapes, rec.nargs))
        for dtypes in itertools.product(
          *(_valid_dtypes_for_shape(s, rec.dtypes) for s in shapes)))
      for rec in JAX_OPERATOR_OVERLOADS))
  @jax.numpy_rank_promotion('allow')  # This test explicitly exercises implicit rank promotion.
  def testOperatorOverload(self, name, rng_factory, shapes, dtypes, tol):
    rng = rng_factory(self.rng())
    # np and jnp arrays have different type promotion rules; force the use of
    # jnp arrays.
    args_maker = self._GetArgsMaker(rng, shapes, dtypes, np_arrays=False)
    fun = lambda *xs: getattr(operator, name.strip('_'))(*xs)
    self._CompileAndCheck(fun, args_maker, atol=tol, rtol=tol)

  @parameterized.named_parameters(itertools.chain.from_iterable(
      jtu.cases_from_list(
        {"testcase_name": jtu.format_test_name_suffix(rec.test_name, shapes,
                                                      dtypes),
         "rng_factory": rec.rng_factory, "shapes": shapes, "dtypes": dtypes, "name": rec.name,
         "op_tolerance": rec.tolerance}
        for shapes in filter(
          _shapes_are_broadcast_compatible,
          itertools.combinations_with_replacement(rec.shapes, rec.nargs))
        for dtypes in itertools.product(
          *(_valid_dtypes_for_shape(s, rec.dtypes) for s in shapes)))
      for rec in JAX_RIGHT_OPERATOR_OVERLOADS))
  @jax.numpy_rank_promotion('allow')  # This test explicitly exercises implicit rank promotion.
  def testRightOperatorOverload(self, name, rng_factory, shapes, dtypes,
                                op_tolerance):
    if shapes[1] is jtu.PYTHON_SCALAR_SHAPE:
      raise SkipTest("scalars not implemented")  # TODO(mattjj): clean up
    rng = rng_factory(self.rng())
    args_maker = self._GetArgsMaker(rng, shapes, dtypes, np_arrays=False)
    fun = lambda fst, snd: getattr(snd, name)(fst)
    tol = max(jtu.tolerance(dtype, op_tolerance) for dtype in dtypes)
    self._CompileAndCheck( fun, args_maker, atol=tol, rtol=tol)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": rec.test_name + f"_{dtype}",
       "rng_factory": rec.rng_factory,
       "op_name": rec.name, "dtype": dtype}
      for rec in JAX_OPERATOR_OVERLOADS if rec.nargs == 2
      for dtype in rec.dtypes))
  def testBinaryOperatorDefers(self, op_name, rng_factory, dtype):
    rng = rng_factory(self.rng())
    arg = jax.device_put(rng((), dtype))
    op = getattr(operator, op_name)

    other = _OverrideEverything()
    assert op(other, arg) is other
    assert op(arg, other) is other

    other = _OverrideNothing()
    if op_name == "__eq__":
      assert op(other, arg) is False
      assert op(arg, other) is False
    elif op_name == "__ne__":
      assert op(other, arg) is True
      assert op(arg, other) is True
    else:
      with self.assertRaises(TypeError):
        op(other, arg)
      with self.assertRaises(TypeError):
        op(arg, other)

  def testArrayEqualExamples(self):
    # examples from the array_equal() docstring.
    self.assertTrue(jnp.array_equal([1, 2], [1, 2]))
    self.assertTrue(jnp.array_equal(np.array([1, 2]), np.array([1, 2])))
    self.assertFalse(jnp.array_equal([1, 2], [1, 2, 3]))
    self.assertFalse(jnp.array_equal([1, 2], [1, 4]))

    a = np.array([1, np.nan])
    self.assertFalse(jnp.array_equal(a, a))
    self.assertTrue(jnp.array_equal(a, a, equal_nan=True))

    a = np.array([1 + 1j])
    b = a.copy()
    a.real = np.nan
    b.imag = np.nan
    self.assertTrue(jnp.array_equal(a, b, equal_nan=True))

  def testArrayEquivExamples(self):
    # examples from the array_equiv() docstring.
    self.assertTrue(jnp.array_equiv([1, 2], [1, 2]))
    self.assertFalse(jnp.array_equiv([1, 2], [1, 3]))
    with jax.numpy_rank_promotion('allow'):
      self.assertTrue(jnp.array_equiv([1, 2], [[1, 2], [1, 2]]))
      self.assertFalse(jnp.array_equiv([1, 2], [[1, 2, 1, 2], [1, 2, 1, 2]]))
      self.assertFalse(jnp.array_equiv([1, 2], [[1, 2], [1, 3]]))

  def testArrayModule(self):
    if numpy_dispatch is None:
      raise SkipTest('requires https://github.com/seberg/numpy-dispatch')

    jnp_array = jnp.array(1.0)
    np_array = np.array(1.0)

    module = numpy_dispatch.get_array_module(jnp_array)
    self.assertIs(module, jnp)

    module = numpy_dispatch.get_array_module(jnp_array, np_array)
    self.assertIs(module, jnp)

    def f(x):
      module = numpy_dispatch.get_array_module(x)
      self.assertIs(module, jnp)
      return x
    jax.jit(f)(jnp_array)
    jax.grad(f)(jnp_array)

  @parameterized.named_parameters(itertools.chain.from_iterable(
      jtu.cases_from_list(
        {"testcase_name": jtu.format_test_name_suffix(
            rec.test_name, shapes, dtypes),
         "rng_factory": rec.rng_factory, "shapes": shapes, "dtypes": dtypes,
         "np_op": getattr(np, rec.name), "jnp_op": getattr(jnp, rec.name)}
        for shapes in filter(
          _shapes_are_broadcast_compatible,
          itertools.combinations_with_replacement(rec.shapes, rec.nargs))
        for dtypes in filter(
          _dtypes_are_compatible_for_bitwise_ops,
          itertools.combinations_with_replacement(rec.dtypes, rec.nargs)))
      for rec in JAX_BITWISE_OP_RECORDS))
  @jax.numpy_rank_promotion('allow')  # This test explicitly exercises implicit rank promotion.
  def testBitwiseOp(self, np_op, jnp_op, rng_factory, shapes, dtypes):
    rng = rng_factory(self.rng())
    if not config.x64_enabled and any(
        jnp.iinfo(dtype).bits == 64 for dtype in dtypes):
      self.skipTest("x64 types are disabled by jax_enable_x64")
    args_maker = self._GetArgsMaker(rng, shapes, dtypes)
    self._CheckAgainstNumpy(np_op, jnp_op, args_maker,
                            check_dtypes=jtu.PYTHON_SCALAR_SHAPE not in shapes)
    self._CompileAndCheck(jnp_op, args_maker)

  @parameterized.named_parameters(jtu.cases_from_list(
    {"testcase_name": jtu.format_test_name_suffix(op.__name__, shapes, dtypes),
     "op": op, "dtypes": dtypes, "shapes": shapes}
    for op in [jnp.left_shift, jnp.right_shift]
    for shapes in filter(
      _shapes_are_broadcast_compatible,
      # TODO numpy always promotes to shift dtype for zero-dim shapes:
      itertools.combinations_with_replacement(nonzerodim_shapes, 2))
    for dtypes in itertools.product(
      *(_valid_dtypes_for_shape(s, int_dtypes_no_uint64) for s in shapes))))
  @jax.numpy_rank_promotion('allow')  # This test explicitly exercises implicit rank promotion.
  def testShiftOpAgainstNumpy(self, op, dtypes, shapes):
    dtype, shift_dtype = dtypes
    signed_mix = np.issubdtype(dtype, np.signedinteger) != \
                 np.issubdtype(shift_dtype, np.signedinteger)
    has_32 = any(np.iinfo(d).bits == 32 for d in dtypes)
    promoting_to_64 = has_32 and signed_mix
    if promoting_to_64 and not config.x64_enabled:
      self.skipTest("np.right_shift/left_shift promoting to int64"
                    "differs from jnp in 32 bit mode.")

    info, shift_info = map(np.iinfo, dtypes)
    x_rng = jtu.rand_int(self.rng(), low=info.min, high=info.max + 1)
    # NumPy requires shifts to be non-negative and below the bit width:
    shift_rng = jtu.rand_int(self.rng(), high=max(info.bits, shift_info.bits))
    args_maker = lambda: (x_rng(shapes[0], dtype), shift_rng(shapes[1], shift_dtype))
    self._CompileAndCheck(op, args_maker)
    np_op = getattr(np, op.__name__)
    self._CheckAgainstNumpy(np_op, op, args_maker)

  @parameterized.named_parameters(itertools.chain.from_iterable(
      jtu.cases_from_list(
        {"testcase_name": "{}_inshape={}_axis={}_dtype={}_keepdims={}".format(
            rec.test_name.capitalize(),
            jtu.format_shape_dtype_string(shape, dtype), axis,
            "None" if out_dtype is None else np.dtype(out_dtype).name, keepdims),
         "rng_factory": rec.rng_factory, "shape": shape, "dtype": dtype, "out_dtype": out_dtype,
         "np_op": getattr(np, rec.name), "jnp_op": getattr(jnp, rec.name),
         "axis": axis, "keepdims": keepdims, "inexact": rec.inexact}
        for shape in rec.shapes for dtype in rec.dtypes
        for out_dtype in [None] + rec.dtypes if out_dtype not in unsigned_dtypes
        for axis in list(range(-len(shape), len(shape))) + [None]
        for keepdims in [False, True])
      for rec in JAX_REDUCER_RECORDS))
  def testReducer(self, np_op, jnp_op, rng_factory, shape, dtype, out_dtype,
                  axis, keepdims, inexact):
    rng = rng_factory(self.rng())
    @jtu.ignore_warning(category=np.ComplexWarning)
    @jtu.ignore_warning(category=RuntimeWarning,
                        message="mean of empty slice.*")
    @jtu.ignore_warning(category=RuntimeWarning,
                        message="overflow encountered.*")
    def np_fun(x):
      x_cast = x if dtype != jnp.bfloat16 else x.astype(np.float32)
      t = out_dtype if out_dtype != jnp.bfloat16 else np.float32
      return np_op(x_cast, axis, dtype=t, keepdims=keepdims)

    jnp_fun = lambda x: jnp_op(x, axis, dtype=out_dtype, keepdims=keepdims)
    jnp_fun = jtu.ignore_warning(category=jnp.ComplexWarning)(jnp_fun)
    args_maker = lambda: [rng(shape, dtype)]
    tol_spec = {np.float16: 1e-2, np.int32: 1E-3, np.float32: 1e-3,
                np.complex64: 1e-3, np.float64: 1e-5, np.complex128: 1e-5}
    tol = jtu.tolerance(dtype, tol_spec)
    tol = max(tol, jtu.tolerance(out_dtype, tol_spec)) if out_dtype else tol
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker,
                            check_dtypes=jnp.bfloat16 not in (dtype, out_dtype),
                            tol=tol)
    self._CompileAndCheck(jnp_fun, args_maker, atol=tol,
                          rtol=tol)

  @parameterized.named_parameters(itertools.chain.from_iterable(
      jtu.cases_from_list(
        {"testcase_name": "{}_inshape={}_axis={}_keepdims={}".format(
            rec.test_name.capitalize(),
            jtu.format_shape_dtype_string(shape, dtype), axis, keepdims),
        "rng_factory": rec.rng_factory, "shape": shape, "dtype": dtype,
        "np_op": getattr(np, rec.name), "jnp_op": getattr(jnp, rec.name),
        "axis": axis, "keepdims": keepdims, "inexact": rec.inexact}
        for shape in rec.shapes for dtype in rec.dtypes
        for axis in list(range(-len(shape), len(shape))) + [None]
        for keepdims in [False, True])
      for rec in JAX_REDUCER_NO_DTYPE_RECORDS))
  def testReducerNoDtype(self, np_op, jnp_op, rng_factory, shape, dtype, axis,
                         keepdims, inexact):
    rng = rng_factory(self.rng())
    is_bf16_nan_test = dtype == jnp.bfloat16 and rng_factory.__name__ == 'rand_some_nan'
    @jtu.ignore_warning(category=RuntimeWarning,
                        message="Degrees of freedom <= 0 for slice.*")
    @jtu.ignore_warning(category=RuntimeWarning,
                        message="All-NaN (slice|axis) encountered.*")
    def np_fun(x):
      x_cast = x if not is_bf16_nan_test else x.astype(np.float32)
      res = np_op(x_cast, axis, keepdims=keepdims)
      res = res if not is_bf16_nan_test else res.astype(jnp.bfloat16)
      return res

    jnp_fun = lambda x: jnp_op(x, axis, keepdims=keepdims)
    args_maker = lambda: [rng(shape, dtype)]
    tol = {np.float16: 0.002}
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker, tol=tol)
    self._CompileAndCheck(jnp_fun, args_maker, rtol=tol, atol=tol)

  @parameterized.named_parameters(itertools.chain.from_iterable(
      jtu.cases_from_list(
        {"testcase_name": "{}_inshape={}_axis={}_keepdims={}_initial={}".format(
            rec.test_name.capitalize(),
            jtu.format_shape_dtype_string(shape, dtype), axis, keepdims, initial),
        "rng_factory": rec.rng_factory, "shape": shape, "dtype": dtype,
        "np_op": getattr(np, rec.name), "jnp_op": getattr(jnp, rec.name),
        "initial": initial, "axis": axis, "keepdims": keepdims, "inexact": rec.inexact}
        for shape in rec.shapes for dtype in rec.dtypes
        for axis in list(range(-len(shape), len(shape))) + [None]
        for initial in [0, 1] for keepdims in [False, True])
      for rec in JAX_REDUCER_INITIAL_RECORDS))
  def testReducerInitial(self, np_op, jnp_op, rng_factory, shape, dtype, axis,
                         keepdims, initial, inexact):
    rng = rng_factory(self.rng())
    is_bf16_nan_test = dtype == jnp.bfloat16 and rng_factory.__name__ == 'rand_some_nan'
    @jtu.ignore_warning(category=RuntimeWarning,
                        message="Degrees of freedom <= 0 for slice.*")
    @jtu.ignore_warning(category=np.ComplexWarning)
    def np_fun(x):
      x_cast = x if not is_bf16_nan_test else x.astype(np.float32)
      res = np_op(x_cast, axis, keepdims=keepdims, initial=initial)
      res = res if not is_bf16_nan_test else res.astype(jnp.bfloat16)
      return res

    jnp_fun = lambda x: jnp_op(x, axis, keepdims=keepdims, initial=initial)
    jnp_fun = jtu.ignore_warning(category=jnp.ComplexWarning)(jnp_fun)
    args_maker = lambda: [rng(shape, dtype)]
    tol = {jnp.bfloat16: 3E-2}
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker, rtol=tol)
    self._CompileAndCheck(jnp_fun, args_maker)

  @parameterized.named_parameters(itertools.chain.from_iterable(
      jtu.cases_from_list(
        {"testcase_name": "{}_inshape={}_axis={}_keepdims={}_initial={}_whereshape={}".format(
            rec.test_name.capitalize(),
            jtu.format_shape_dtype_string(shape, dtype), axis, keepdims, initial,
            jtu.format_shape_dtype_string(whereshape, bool)),
        "rng_factory": rec.rng_factory, "shape": shape, "dtype": dtype,
        "np_op": getattr(np, rec.name), "jnp_op": getattr(jnp, rec.name), "whereshape": whereshape,
        "initial": initial, "axis": axis, "keepdims": keepdims, "inexact": rec.inexact}
        for shape in rec.shapes for dtype in rec.dtypes
        for whereshape in _compatible_shapes(shape)
        for axis in list(range(-len(shape), len(shape))) + [None]
        for initial in [0, 1] for keepdims in [False, True])
      for rec in JAX_REDUCER_INITIAL_RECORDS))
  def testReducerWhere(self, np_op, jnp_op, rng_factory, shape, dtype, axis,
                       keepdims, initial, inexact, whereshape):
    if (shape in [()] + scalar_shapes and
        dtype in [jnp.int16, jnp.uint16] and
        jnp_op in [jnp.min, jnp.max]):
      self.skipTest("Known XLA failure; see https://github.com/google/jax/issues/4971.")
    rng = rng_factory(self.rng())
    is_bf16_nan_test = dtype == jnp.bfloat16 and rng_factory.__name__ == 'rand_some_nan'
    # Do not pass where via args_maker as that is incompatible with _promote_like_jnp.
    where = jtu.rand_bool(self.rng())(whereshape, np.bool_)
    @jtu.ignore_warning(category=RuntimeWarning,
                        message="Degrees of freedom <= 0 for slice.*")
    @jtu.ignore_warning(category=np.ComplexWarning)
    def np_fun(x):
      x_cast = x if not is_bf16_nan_test else x.astype(np.float32)
      res = np_op(x_cast, axis, keepdims=keepdims, initial=initial, where=where)
      res = res if not is_bf16_nan_test else res.astype(jnp.bfloat16)
      return res

    jnp_fun = lambda x: jnp_op(x, axis, keepdims=keepdims, initial=initial, where=where)
    jnp_fun = jtu.ignore_warning(category=jnp.ComplexWarning)(jnp_fun)
    args_maker = lambda: [rng(shape, dtype)]
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)
    self._CompileAndCheck(jnp_fun, args_maker)

  @unittest.skipIf(numpy_version < (1, 20), "where parameter not supported in older numpy")
  @parameterized.named_parameters(itertools.chain.from_iterable(
    jtu.cases_from_list(
      {"testcase_name": "{}_inshape={}_axis={}_keepdims={}_whereshape={}".format(
        rec.test_name.capitalize(),
        jtu.format_shape_dtype_string(shape, dtype), axis, keepdims,
        jtu.format_shape_dtype_string(whereshape, bool)),
        "rng_factory": rec.rng_factory, "shape": shape, "dtype": dtype,
        "np_op": getattr(np, rec.name), "jnp_op": getattr(jnp, rec.name), "whereshape": whereshape,
        "axis": axis, "keepdims": keepdims, "inexact": rec.inexact}
      for shape in rec.shapes for dtype in rec.dtypes
      for whereshape in _compatible_shapes(shape)
      for axis in list(range(-len(shape), len(shape))) + [None]
      for keepdims in [False, True])
    for rec in JAX_REDUCER_WHERE_NO_INITIAL_RECORDS))
  def testReducerWhereNoInitial(self, np_op, jnp_op, rng_factory, shape, dtype, axis,
                                keepdims, inexact, whereshape):
    rng = rng_factory(self.rng())
    is_bf16_nan_test = dtype == jnp.bfloat16
    # Do not pass where via args_maker as that is incompatible with _promote_like_jnp.
    where = jtu.rand_bool(self.rng())(whereshape, np.bool_)
    @jtu.ignore_warning(category=RuntimeWarning,
                        message="Degrees of freedom <= 0 for slice.*")
    @jtu.ignore_warning(category=RuntimeWarning,
                        message="Mean of empty slice.*")
    @jtu.ignore_warning(category=RuntimeWarning,
                        message="invalid value encountered.*")
    @jtu.ignore_warning(category=np.ComplexWarning)
    def np_fun(x):
      x_cast = x if not is_bf16_nan_test else x.astype(np.float32)
      res = np_op(x_cast, axis, keepdims=keepdims, where=where)
      res = res if not is_bf16_nan_test else res.astype(jnp.bfloat16)
      return res

    jnp_fun = lambda x: jnp_op(x, axis, keepdims=keepdims, where=where)
    jnp_fun = jtu.ignore_warning(category=jnp.ComplexWarning)(jnp_fun)
    args_maker = lambda: [rng(shape, dtype)]
    if numpy_version >= (1, 20, 2) or np_op.__name__ in ("all", "any"):
      self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)
    self._CompileAndCheck(jnp_fun, args_maker)


  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_shape={}_axis={}_discont={}_period={}".format(
          jtu.format_shape_dtype_string(shape, dtype), axis, discont, period),
       "shape": shape, "dtype": dtype, "axis": axis, "discont": discont, "period": period}
      for shape in all_shapes for dtype in default_dtypes
      for discont in [None, "pi", 2]
      for period in ["2pi", "pi"]
      for axis in list(range(-len(shape), len(shape)))))
  def testUnwrap(self, shape, dtype, axis, discont, period):
    if numpy_version < (1, 21) and period != "2pi":
      self.skipTest("numpy < 1.21 does not support the period argument to unwrap()")
    special_vals = {"pi": np.pi, "2pi": 2 * np.pi}
    period = special_vals.get(period, period)
    discont = special_vals.get(discont, discont)

    rng = jtu.rand_default(self.rng())
    if numpy_version < (1, 21):
      np_fun = partial(np.unwrap, axis=axis, discont=discont)
    else:
      np_fun = partial(np.unwrap, axis=axis, discont=discont, period=period)
    jnp_fun = partial(jnp.unwrap, axis=axis, discont=discont, period=period)
    args_maker = lambda: [rng(shape, dtype)]
    if dtype != jnp.bfloat16:  # numpy crashes on bfloat16
      self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker, check_dtypes=False)
    self._CompileAndCheck(jnp_fun, args_maker)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_shape={}_axis={}".format(
          jtu.format_shape_dtype_string(shape, dtype), axis),
       "shape": shape, "dtype": dtype, "axis": axis}
      for shape in all_shapes for dtype in all_dtypes
      for axis in list(range(-len(shape), len(shape))) + [None]))
  def testCountNonzero(self, shape, dtype, axis):
    rng = jtu.rand_some_zero(self.rng())
    np_fun = lambda x: np.count_nonzero(x, axis)
    jnp_fun = lambda x: jnp.count_nonzero(x, axis)
    args_maker = lambda: [rng(shape, dtype)]
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker, check_dtypes=False)
    self._CompileAndCheck(jnp_fun, args_maker)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_shape={}".format(
          jtu.format_shape_dtype_string(shape, dtype)),
       "shape": shape, "dtype": dtype}
      for shape in all_shapes for dtype in all_dtypes))
  def testNonzero(self, shape, dtype):
    rng = jtu.rand_some_zero(self.rng())
    np_fun = lambda x: np.nonzero(x)
    np_fun = jtu.ignore_warning(
      category=DeprecationWarning,
      message="Calling nonzero on 0d arrays.*")(np_fun)
    jnp_fun = lambda x: jnp.nonzero(x)
    args_maker = lambda: [rng(shape, dtype)]
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker, check_dtypes=False)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_shape={}_size={}_fill_value={}".format(
          jtu.format_shape_dtype_string(shape, dtype), size, fill_value),
       "shape": shape, "dtype": dtype, "size": size, "fill_value": fill_value}
      for shape in nonempty_array_shapes
      for dtype in all_dtypes
      for fill_value in [None, -1, shape or (1,)]
      for size in [1, 5, 10]))
  def testNonzeroSize(self, shape, dtype, size, fill_value):
    rng = jtu.rand_some_zero(self.rng())
    args_maker = lambda: [rng(shape, dtype)]
    @jtu.ignore_warning(category=DeprecationWarning, message="Calling nonzero on 0d arrays.*")
    def np_fun(x):
      result = np.nonzero(x)
      if size <= len(result[0]):
        return tuple(arg[:size] for arg in result)
      else:
        fillvals = fill_value if np.ndim(fill_value) else len(result) * [fill_value or 0]
        return tuple(np.concatenate([arg, np.full(size - len(arg), fval, arg.dtype)])
                     for fval, arg in safe_zip(fillvals, result))
    jnp_fun = lambda x: jnp.nonzero(x, size=size, fill_value=fill_value)
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker, check_dtypes=False)
    self._CompileAndCheck(jnp_fun, args_maker)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_shape={}".format(
          jtu.format_shape_dtype_string(shape, dtype)),
       "shape": shape, "dtype": dtype}
      for shape in all_shapes for dtype in all_dtypes))
  def testFlatNonzero(self, shape, dtype):
    rng = jtu.rand_some_zero(self.rng())
    np_fun = jtu.ignore_warning(
      category=DeprecationWarning,
      message="Calling nonzero on 0d arrays.*")(np.flatnonzero)
    jnp_fun = jnp.flatnonzero
    args_maker = lambda: [rng(shape, dtype)]
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker, check_dtypes=False)

    # JIT compilation requires specifying the size statically:
    jnp_fun = lambda x: jnp.flatnonzero(x, size=np.size(x) // 2)
    self._CompileAndCheck(jnp_fun, args_maker)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_shape={}_size={}_fill_value={}".format(
          jtu.format_shape_dtype_string(shape, dtype), size, fill_value),
       "shape": shape, "dtype": dtype, "size": size, "fill_value": fill_value}
      for shape in nonempty_array_shapes
      for dtype in all_dtypes
      for fill_value in [None, -1, 10, (-1,), (10,)]
      for size in [1, 5, 10]))
  def testFlatNonzeroSize(self, shape, dtype, size, fill_value):
    rng = jtu.rand_some_zero(self.rng())
    args_maker = lambda: [rng(shape, dtype)]
    @jtu.ignore_warning(category=DeprecationWarning, message="Calling nonzero on 0d arrays.*")
    def np_fun(x):
      result = np.flatnonzero(x)
      if size <= len(result):
        return result[:size]
      else:
        fill_val = fill_value or 0
        return np.concatenate([result, np.full(size - len(result), fill_val, result.dtype)])
    jnp_fun = lambda x: jnp.flatnonzero(x, size=size, fill_value=fill_value)
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker, check_dtypes=False)
    self._CompileAndCheck(jnp_fun, args_maker)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_shape={}".format(
          jtu.format_shape_dtype_string(shape, dtype)),
       "shape": shape, "dtype": dtype}
      for shape in all_shapes for dtype in all_dtypes))
  def testArgWhere(self, shape, dtype):
    rng = jtu.rand_some_zero(self.rng())
    np_fun = jtu.ignore_warning(
      category=DeprecationWarning,
      message="Calling nonzero on 0d arrays.*")(np.argwhere)
    jnp_fun = jnp.argwhere
    args_maker = lambda: [rng(shape, dtype)]
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker, check_dtypes=False)

    # JIT compilation requires specifying a size statically. Full test of this
    # behavior is in testNonzeroSize().
    jnp_fun = lambda x: jnp.argwhere(x, size=np.size(x) // 2)
    self._CompileAndCheck(jnp_fun, args_maker)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_shape={}_size={}_fill_value={}".format(
          jtu.format_shape_dtype_string(shape, dtype), size, fill_value),
       "shape": shape, "dtype": dtype, "size": size, "fill_value": fill_value}
      for shape in nonempty_array_shapes
      for dtype in all_dtypes
      for fill_value in [None, -1, shape or (1,)]
      for size in [1, 5, 10]))
  def testArgWhereSize(self, shape, dtype, size, fill_value):
    rng = jtu.rand_some_zero(self.rng())
    args_maker = lambda: [rng(shape, dtype)]
    @jtu.ignore_warning(category=DeprecationWarning, message="Calling nonzero on 0d arrays.*")
    def np_fun(x):
      result = np.argwhere(x)
      if size <= len(result):
        return result[:size]
      else:
        fillvals = fill_value if np.ndim(fill_value) else result.shape[-1] * [fill_value or 0]
        return np.empty((size, 0), dtype=int) if np.ndim(x) == 0 else np.stack([np.concatenate([arg, np.full(size - len(arg), fval, arg.dtype)])
                        for fval, arg in safe_zip(fillvals, result.T)]).T
    jnp_fun = lambda x: jnp.argwhere(x, size=size, fill_value=fill_value)
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker, check_dtypes=False)
    self._CompileAndCheck(jnp_fun, args_maker)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "{}_inshape={}_axis={}_keepdims={}".format(
          rec.test_name.capitalize(),
          jtu.format_shape_dtype_string(shape, dtype), axis, keepdims),
       "rng_factory": rec.rng_factory, "shape": shape, "dtype": dtype,
       "np_op": getattr(np, rec.name), "jnp_op": getattr(jnp, rec.name),
       "axis": axis, "keepdims": keepdims}
      for rec in JAX_ARGMINMAX_RECORDS
      for shape, dtype in _shape_and_dtypes(rec.shapes, rec.dtypes)
      for axis in range(-len(shape), len(shape))
      for keepdims in [True, False]))
  def testArgMinMax(self, np_op, jnp_op, rng_factory, shape, dtype, axis, keepdims):
    rng = rng_factory(self.rng())
    if dtype == np.complex128 and jtu.device_under_test() == "gpu":
      raise unittest.SkipTest("complex128 reductions not supported on GPU")
    if "nan" in np_op.__name__ and dtype == jnp.bfloat16:
      raise unittest.SkipTest("NumPy doesn't correctly handle bfloat16 arrays")
    if numpy_version < (1, 22) and keepdims:
      raise unittest.SkipTest("NumPy < 1.22 does not support keepdims argument to argmin/argmax")
    kwds = {"keepdims": True} if keepdims else {}

    np_fun = jtu.with_jax_dtype_defaults(partial(np_op, axis=axis, **kwds))
    jnp_fun = partial(jnp_op, axis=axis, **kwds)

    args_maker = lambda: [rng(shape, dtype)]
    try:
      self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)
    except ValueError as e:
      if str(e) == "All-NaN slice encountered":
        self.skipTest("JAX doesn't support checking for all-NaN slices")
      else:
        raise
    self._CompileAndCheck(jnp_fun, args_maker)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": rec.test_name.capitalize(), "name": rec.name,
       "np_op": getattr(np, rec.name), "jnp_op": getattr(jnp, rec.name)}
      for rec in JAX_ARGMINMAX_RECORDS))
  def testArgMinMaxEmpty(self, name, np_op, jnp_op):
    name = name[3:] if name.startswith("nan") else name
    msg = f"attempt to get {name} of an empty sequence"
    with self.assertRaises(ValueError, msg=msg):
      jnp_op(np.array([]))
    with self.assertRaises(ValueError, msg=msg):
      jnp_op(np.zeros((2, 0)), axis=1)
    np_fun = jtu.with_jax_dtype_defaults(partial(np_op, axis=0))
    jnp_fun = partial(jnp_op, axis=0)
    args_maker = lambda: [np.zeros((2, 0))]
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)
    self._CompileAndCheck(jnp_fun, args_maker)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_{}_{}_{}".format(
          jtu.format_shape_dtype_string(lhs_shape, lhs_dtype),
          jtu.format_shape_dtype_string(rhs_shape, rhs_dtype),
          axes),
       "lhs_shape": lhs_shape, "lhs_dtype": lhs_dtype,
       "rhs_shape": rhs_shape, "rhs_dtype": rhs_dtype,
       "axes": axes}
      for lhs_shape, rhs_shape, axes in [
          [(2,), (2,), (-1, -1, -1, None)], # scalar output
          [(2, 4), (2, 4), (-1, -1, -1, 0)], # 2D vectors
          [(3, 4), (3, 4), (-1, -1, -1, 0)], # 3D vectors
          [(3, 4), (3, 6, 5, 4), (-1, -1, -1, 0)], # broadcasting
          [(4, 3), (3, 6, 5, 4), (1, 0, -1, None)], # different axes
          [(6, 1, 3), (5, 3), (-1, -1, -1, None)], # more broadcasting
          [(6, 1, 2), (5, 3), (-1, -1, -1, None)], # mixed 2D and 3D vectors
          [(10, 5, 2, 8), (1, 5, 1, 3), (-2, -1, -3, None)], # axes/broadcasting
          [(4, 5, 2), (4, 5, 2), (-1, -1, 0, None)], # axisc should do nothing
          [(4, 5, 2), (4, 5, 2), (-1, -1, -1, None)] # same as before
      ]
      for lhs_dtype, rhs_dtype in itertools.combinations_with_replacement(number_dtypes, 2)))
  @jax.numpy_rank_promotion('allow')  # This test explicitly exercises implicit rank promotion.
  def testCross(self, lhs_shape, lhs_dtype, rhs_shape, rhs_dtype, axes):
    rng = jtu.rand_default(self.rng())
    args_maker = lambda: [rng(lhs_shape, lhs_dtype), rng(rhs_shape, rhs_dtype)]
    axisa, axisb, axisc, axis = axes
    jnp_fun = lambda a, b: jnp.cross(a, b, axisa, axisb, axisc, axis)
    def np_fun(a, b):
      a = a.astype(np.float32) if lhs_dtype == jnp.bfloat16 else a
      b = b.astype(np.float32) if rhs_dtype == jnp.bfloat16 else b
      out = np.cross(a, b, axisa, axisb, axisc, axis)
      return out.astype(jnp.promote_types(lhs_dtype, rhs_dtype))
    tol_spec = {dtypes.bfloat16: 3e-1, np.float16: 0.15}
    tol = max(jtu.tolerance(lhs_dtype, tol_spec),
              jtu.tolerance(rhs_dtype, tol_spec))
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker,
                            tol=tol)
    self._CompileAndCheck(jnp_fun, args_maker, atol=tol,
                          rtol=tol)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_{}_{}_{}".format(
          name,
          jtu.format_shape_dtype_string(lhs_shape, lhs_dtype),
          jtu.format_shape_dtype_string(rhs_shape, rhs_dtype)),
       "lhs_shape": lhs_shape, "lhs_dtype": lhs_dtype,
       "rhs_shape": rhs_shape, "rhs_dtype": rhs_dtype}
      for name, lhs_shape, rhs_shape in [
          ("matrix-scalar", (3, 3), ()),
          ("scalar-matrix", (), (3, 3)),
          ("matrix-vector", (4, 5), (5,)),
          ("vector-matrix", (6,), (6, 4)),
          ("matrix-matrix", (3, 4), (4, 5)),
          ("tensor-vector", (4, 3, 2), (2,)),
          ("vector-tensor", (2,), (3, 2, 4)),
          ("tensor-matrix", (4, 3, 2), (2, 5)),
          ("matrix-tensor", (5, 2), (3, 2, 4)),
          ("tensor-tensor", (2, 3, 4), (5, 4, 1))]
      for lhs_dtype, rhs_dtype in itertools.combinations_with_replacement(number_dtypes, 2)))
  def testDot(self, lhs_shape, lhs_dtype, rhs_shape, rhs_dtype):
    rng = jtu.rand_default(self.rng())
    args_maker = lambda: [rng(lhs_shape, lhs_dtype), rng(rhs_shape, rhs_dtype)]
    tol = {np.float16: 1e-2, np.float32: 1e-5, np.float64: 1e-14,
           np.complex128: 1e-14}
    if jtu.device_under_test() == "tpu":
      tol[np.float16] = tol[np.float32] = tol[np.complex64] = 2e-1
    def np_dot(x, y):
      x = x.astype(np.float32) if lhs_dtype == jnp.bfloat16 else x
      y = y.astype(np.float32) if rhs_dtype == jnp.bfloat16 else y
      return np.dot(x, y).astype(jnp.promote_types(lhs_dtype, rhs_dtype))
    self._CheckAgainstNumpy(np_dot, jnp.dot, args_maker,
                            tol=tol)
    self._CompileAndCheck(jnp.dot, args_maker, atol=tol,
                          rtol=tol)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_{}_{}_{}".format(
          name,
          jtu.format_shape_dtype_string(lhs_shape, lhs_dtype),
          jtu.format_shape_dtype_string(rhs_shape, rhs_dtype)),
       "lhs_shape": lhs_shape, "lhs_dtype": lhs_dtype,
       "rhs_shape": rhs_shape, "rhs_dtype": rhs_dtype}
      for name, lhs_shape, rhs_shape in [
          ("vector-vector", (3,), (3,)),
          ("matrix-vector", (3, 3), (3,)),
          ("vector-matrix", (3,), (3, 3)),
          ("matrix-matrix", (3, 3), (3, 3)),
          ("vector-tensor", (3,), (5, 3, 2)),
          ("tensor-vector", (5, 3, 2), (2,)),
          ("matrix-tensor", (5, 2), (3, 2, 4)),
          ("tensor-matrix", (5, 2, 3), (3, 2)),
          ("tensor-tensor", (5, 3, 4), (5, 4, 1)),
          ("tensor-tensor-broadcast", (3, 1, 3, 4), (5, 4, 1))]
      for lhs_dtype, rhs_dtype in itertools.combinations_with_replacement(number_dtypes, 2)))
  def testMatmul(self, lhs_shape, lhs_dtype, rhs_shape, rhs_dtype):
    rng = jtu.rand_default(self.rng())
    def np_fun(x, y):
      dtype = jnp.promote_types(lhs_dtype, rhs_dtype)
      return np.matmul(x, y).astype(dtype)
    args_maker = lambda: [rng(lhs_shape, lhs_dtype), rng(rhs_shape, rhs_dtype)]
    tol = {np.float16: 1e-2, np.float32: 2e-2, np.float64: 1e-12,
           np.complex128: 1e-12}
    if jtu.device_under_test() == "tpu":
      tol[np.float16] = tol[np.float32] = tol[np.complex64] = 4e-2
    self._CheckAgainstNumpy(np_fun, jnp.matmul, args_maker, tol=tol)
    self._CompileAndCheck(jnp.matmul, args_maker, atol=tol, rtol=tol)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_{}_{}_{}".format(
          jtu.format_shape_dtype_string(lhs_shape, lhs_dtype),
          jtu.format_shape_dtype_string(rhs_shape, rhs_dtype),
          axes),
       "lhs_shape": lhs_shape, "lhs_dtype": lhs_dtype,
       "rhs_shape": rhs_shape, "rhs_dtype": rhs_dtype,
       "axes": axes}
      for lhs_shape, rhs_shape, axes in [
          [(3,), (), 0],
          [(2, 3, 4), (5, 6, 7), 0],  # from issue #740
          [(2, 3, 4), (3, 4, 5, 6), 2],
          [(2, 3, 4), (5, 4, 3, 6), [1, 2]],
          [(2, 3, 4), (5, 4, 3, 6), [[1, 2], [2, 1]]],
          [(1, 2, 3, 4), (4, 5, 3, 6), [[2, 3], [2, 0]]],
      ]
      for lhs_dtype, rhs_dtype in itertools.combinations_with_replacement(number_dtypes, 2)))
  def testTensordot(self, lhs_shape, lhs_dtype, rhs_shape, rhs_dtype, axes):
    rng = jtu.rand_default(self.rng())
    args_maker = lambda: [rng(lhs_shape, lhs_dtype), rng(rhs_shape, rhs_dtype)]
    jnp_fun = lambda a, b: jnp.tensordot(a, b, axes)
    def np_fun(a, b):
      a = a if lhs_dtype != jnp.bfloat16 else a.astype(np.float32)
      b = b if rhs_dtype != jnp.bfloat16 else b.astype(np.float32)
      dtype = jnp.promote_types(lhs_dtype, rhs_dtype)
      return np.tensordot(a, b, axes).astype(dtype)
    tol = {np.float16: 1e-1, np.float32: 1e-3, np.float64: 1e-12,
           np.complex64: 1e-3, np.complex128: 1e-12}
    if jtu.device_under_test() == "tpu":
      tol[np.float16] = tol[np.float32] = tol[np.complex64] = 2e-1
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker,
                            tol=tol)
    self._CompileAndCheck(jnp_fun, args_maker)

  def testTensordotErrors(self):
    a = self.rng().random((3, 2, 2))
    b = self.rng().random((2,))
    self.assertRaisesRegex(
      TypeError, "Number of tensordot axes.*exceeds input ranks.*",
      lambda: jnp.tensordot(a, b, axes=2))

    self.assertRaisesRegex(
      TypeError, "tensordot requires axes lists to have equal length.*",
      lambda: jnp.tensordot(a, b, axes=([0], [0, 1])))

    self.assertRaisesRegex(
      TypeError, "tensordot requires both axes lists to be either ints, tuples or lists.*",
      lambda: jnp.tensordot(a, b, axes=('bad', 'axes')))

    self.assertRaisesRegex(
      TypeError, "tensordot axes argument must be an int, a pair of ints, or a pair of lists.*",
      lambda: jnp.tensordot(a, b, axes='badaxes'))


  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_{}_{}_invert={}".format(
          jtu.format_shape_dtype_string(element_shape, dtype),
          jtu.format_shape_dtype_string(test_shape, dtype), invert),
       "element_shape": element_shape, "test_shape": test_shape,
       "dtype": dtype, "invert": invert}
      for element_shape in all_shapes
      for test_shape in all_shapes
      for dtype in default_dtypes
      for invert in [True, False]))
  def testIsin(self, element_shape, test_shape, dtype, invert):
    rng = jtu.rand_default(self.rng())
    args_maker = lambda: [rng(element_shape, dtype), rng(test_shape, dtype)]
    jnp_fun = lambda e, t: jnp.isin(e, t, invert=invert)
    np_fun = lambda e, t: np.isin(e, t, invert=invert)
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)
    self._CompileAndCheck(jnp_fun, args_maker)


  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_{}_{}_invert={}".format(
          jtu.format_shape_dtype_string(element_shape, dtype),
          jtu.format_shape_dtype_string(test_shape, dtype), invert),
       "element_shape": element_shape, "test_shape": test_shape,
       "dtype": dtype, "invert": invert}
      for element_shape in all_shapes
      for test_shape in all_shapes
      for dtype in default_dtypes
      for invert in [True, False]))
  def testIn1d(self, element_shape, test_shape, dtype, invert):
    rng = jtu.rand_default(self.rng())
    args_maker = lambda: [rng(element_shape, dtype), rng(test_shape, dtype)]
    jnp_fun = lambda e, t: jnp.in1d(e, t, invert=invert)
    np_fun = lambda e, t: np.in1d(e, t, invert=invert)
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)
    self._CompileAndCheck(jnp_fun, args_maker)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_{}_{}".format(
       jtu.format_shape_dtype_string(shape1, dtype1),
       jtu.format_shape_dtype_string(shape2, dtype2)),
       "shape1": shape1, "shape2": shape2, "dtype1": dtype1, "dtype2": dtype2}
      for dtype1 in [s for s in default_dtypes if s != jnp.bfloat16]
      for dtype2 in [s for s in default_dtypes if s != jnp.bfloat16]
      for shape1 in all_shapes
      for shape2 in all_shapes))
  def testSetdiff1d(self, shape1, shape2, dtype1, dtype2):
    rng = jtu.rand_default(self.rng())
    args_maker = lambda: [rng(shape1, dtype1), rng(shape2, dtype2)]
    self._CheckAgainstNumpy(np.setdiff1d, jnp.setdiff1d, args_maker)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_{}_{}_size={}_fill_value={}".format(
       jtu.format_shape_dtype_string(shape1, dtype1),
       jtu.format_shape_dtype_string(shape2, dtype2),
       size, fill_value),
       "shape1": shape1, "shape2": shape2, "dtype1": dtype1, "dtype2": dtype2,
       "size": size, "fill_value": fill_value}
      for dtype1 in [s for s in default_dtypes if s != jnp.bfloat16]
      for dtype2 in [s for s in default_dtypes if s != jnp.bfloat16]
      for shape1 in all_shapes
      for shape2 in all_shapes
      for size in [1, 5, 10]
      for fill_value in [None, -1]))
  def testSetdiff1dSize(self, shape1, shape2, dtype1, dtype2, size, fill_value):
    rng = jtu.rand_default(self.rng())
    args_maker = lambda: [rng(shape1, dtype1), rng(shape2, dtype2)]
    def np_fun(arg1, arg2):
      result = np.setdiff1d(arg1, arg2)
      if size <= len(result):
        return result[:size]
      else:
        return np.pad(result, (0, size-len(result)), constant_values=fill_value or 0)
    def jnp_fun(arg1, arg2):
      return jnp.setdiff1d(arg1, arg2, size=size, fill_value=fill_value)
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)
    self._CompileAndCheck(jnp_fun, args_maker)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_{}_{}".format(
       jtu.format_shape_dtype_string(shape1, dtype1),
       jtu.format_shape_dtype_string(shape2, dtype2)),
       "shape1": shape1, "shape2": shape2, "dtype1": dtype1, "dtype2": dtype2}
      for dtype1 in [s for s in default_dtypes if s != jnp.bfloat16]
      for dtype2 in [s for s in default_dtypes if s != jnp.bfloat16]
      for shape1 in nonempty_nonscalar_array_shapes
      for shape2 in nonempty_nonscalar_array_shapes))
  def testUnion1d(self, shape1, shape2, dtype1, dtype2):
    rng = jtu.rand_default(self.rng())
    args_maker = lambda: [rng(shape1, dtype1), rng(shape2, dtype2)]
    def np_fun(arg1, arg2):
      dtype = jnp.promote_types(arg1.dtype, arg2.dtype)
      return np.union1d(arg1, arg2).astype(dtype)
    self._CheckAgainstNumpy(np_fun, jnp.union1d, args_maker)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_{}_{}_size={}_fill_value={}".format(
       jtu.format_shape_dtype_string(shape1, dtype1),
       jtu.format_shape_dtype_string(shape2, dtype2), size, fill_value),
       "shape1": shape1, "shape2": shape2, "dtype1": dtype1, "dtype2": dtype2,
       "size": size, "fill_value": fill_value}
      for dtype1 in [s for s in default_dtypes if s != jnp.bfloat16]
      for dtype2 in [s for s in default_dtypes if s != jnp.bfloat16]
      for shape1 in nonempty_nonscalar_array_shapes
      for shape2 in nonempty_nonscalar_array_shapes
      for size in [1, 5, 10]
      for fill_value in [None, -1]))
  def testUnion1dSize(self, shape1, shape2, dtype1, dtype2, size, fill_value):
    rng = jtu.rand_default(self.rng())
    args_maker = lambda: [rng(shape1, dtype1), rng(shape2, dtype2)]
    def np_fun(arg1, arg2):
      dtype = jnp.promote_types(arg1.dtype, arg2.dtype)
      result = np.union1d(arg1, arg2).astype(dtype)
      fv = result.min() if fill_value is None else fill_value
      if size <= len(result):
        return result[:size]
      else:
        return np.concatenate([result, np.full(size - len(result), fv, result.dtype)])
    def jnp_fun(arg1, arg2):
      return jnp.union1d(arg1, arg2, size=size, fill_value=fill_value)
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)
    self._CompileAndCheck(jnp_fun, args_maker)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_{}_{}_assume_unique={}".format(
       jtu.format_shape_dtype_string(shape1, dtype1),
       jtu.format_shape_dtype_string(shape2, dtype2),
       assume_unique),
       "shape1": shape1, "dtype1": dtype1, "shape2": shape2, "dtype2": dtype2,
       "assume_unique": assume_unique}
      for dtype1 in [s for s in default_dtypes if s != jnp.bfloat16]
      for dtype2 in [s for s in default_dtypes if s != jnp.bfloat16]
      for shape1 in all_shapes
      for shape2 in all_shapes
      for assume_unique in [False, True]))
  def testSetxor1d(self, shape1, dtype1, shape2, dtype2, assume_unique):
    rng = jtu.rand_default(self.rng())
    args_maker = lambda: [rng(shape1, dtype1), rng(shape2, dtype2)]
    jnp_fun = lambda ar1, ar2: jnp.setxor1d(ar1, ar2, assume_unique=assume_unique)
    def np_fun(ar1, ar2):
      if assume_unique:
        # pre-flatten the arrays to match with jax implementation
        ar1 = np.ravel(ar1)
        ar2 = np.ravel(ar2)
      return np.setxor1d(ar1, ar2, assume_unique)
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker, check_dtypes=False)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_{}_{}_assume_unique={}_return_indices={}".format(
       jtu.format_shape_dtype_string(shape1, dtype1),
       jtu.format_shape_dtype_string(shape2, dtype2),
       assume_unique,
       return_indices),
       "shape1": shape1, "dtype1": dtype1, "shape2": shape2, "dtype2": dtype2,
       "assume_unique": assume_unique, "return_indices": return_indices}
      for dtype1 in [s for s in default_dtypes if s != jnp.bfloat16]
      for dtype2 in [s for s in default_dtypes if s != jnp.bfloat16]
      for shape1 in all_shapes
      for shape2 in all_shapes
      for assume_unique in [False, True]
      for return_indices in [False, True]))
  def testIntersect1d(self, shape1, dtype1, shape2, dtype2, assume_unique, return_indices):
    rng = jtu.rand_default(self.rng())
    args_maker = lambda: [rng(shape1, dtype1), rng(shape2, dtype2)]
    jnp_fun = lambda ar1, ar2: jnp.intersect1d(ar1, ar2, assume_unique=assume_unique, return_indices=return_indices)
    np_fun = lambda ar1, ar2: np.intersect1d(ar1, ar2, assume_unique=assume_unique, return_indices=return_indices)
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker, check_dtypes=False)


  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_{}_{}".format(
          jtu.format_shape_dtype_string(lhs_shape, lhs_dtype),
          jtu.format_shape_dtype_string(rhs_shape, rhs_dtype)),
       "lhs_shape": lhs_shape, "lhs_dtype": lhs_dtype,
       "rhs_shape": rhs_shape, "rhs_dtype": rhs_dtype}
      # TODO(phawkins): support integer dtypes too.
      for lhs_shape, lhs_dtype in _shape_and_dtypes(all_shapes, inexact_dtypes)
      for rhs_shape, rhs_dtype in _shape_and_dtypes(all_shapes, inexact_dtypes)
      if len(jtu._dims_of_shape(lhs_shape)) == 0
      or len(jtu._dims_of_shape(rhs_shape)) == 0
      or lhs_shape[-1] == rhs_shape[-1]))
  def testInner(self, lhs_shape, lhs_dtype, rhs_shape, rhs_dtype):
    rng = jtu.rand_default(self.rng())
    args_maker = lambda: [rng(lhs_shape, lhs_dtype), rng(rhs_shape, rhs_dtype)]
    def np_fun(lhs, rhs):
      lhs = lhs if lhs_dtype != jnp.bfloat16 else lhs.astype(np.float32)
      rhs = rhs if rhs_dtype != jnp.bfloat16 else rhs.astype(np.float32)
      dtype = jnp.promote_types(lhs_dtype, rhs_dtype)
      return np.inner(lhs, rhs).astype(dtype)
    jnp_fun = lambda lhs, rhs: jnp.inner(lhs, rhs)
    tol_spec = {np.float16: 1e-2, np.float32: 1e-5, np.float64: 1e-13,
                np.complex64: 1e-5}
    if jtu.device_under_test() == "tpu":
      tol_spec[np.float32] = tol_spec[np.complex64] = 2e-1
    tol = max(jtu.tolerance(lhs_dtype, tol_spec),
              jtu.tolerance(rhs_dtype, tol_spec))
    # TODO(phawkins): there are float32/float64 disagreements for some inputs.
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker, check_dtypes=False, tol=tol)
    self._CompileAndCheck(jnp_fun, args_maker, check_dtypes=False, atol=tol, rtol=tol)


  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_{}_deg={}_rcond={}_full={}_w={}_cov={}".format(
          jtu.format_shape_dtype_string(shape, dtype),
          deg,
          rcond,
          full,
          w,
          cov),
      "shape": shape, "dtype": dtype, "deg": deg,
      "rcond": rcond, "full": full, "w":w, "cov":cov}
      for dtype in [dt for dt in float_dtypes if dt not in [jnp.float16, jnp.bfloat16]]
      for shape in [shape for shape in one_dim_array_shapes if shape != (1,)]
      for deg in [1, 2, 3]
      for rcond in [None, -1, 10e-3, 10e-5, 10e-10]
      for full in [False, True]
      for w in [False, True]
      for cov in [False, True, "unscaled"]))
  def testPolyfit(self, shape, dtype, deg, rcond, full, w, cov):
    rng = jtu.rand_default(self.rng())
    tol_spec = {np.float32: 1e-3, np.float64: 1e-13, np.complex64: 1e-5}
    if jtu.device_under_test() == "tpu":
      tol_spec[np.float32] = tol_spec[np.complex64] = 2e-1
    tol = jtu.tolerance(dtype, tol_spec)
    _w = lambda a: abs(a) if w else None
    args_maker = lambda: [rng(shape, dtype), rng(shape, dtype), rng(shape, dtype)]
    jnp_fun = lambda x, y, a: jnp.polyfit(x, y, deg=deg, rcond=rcond, full=full, w=_w(a), cov=cov)
    np_fun = jtu.ignore_warning(
      message="Polyfit may be poorly conditioned*")(lambda x, y, a: np.polyfit(x, y, deg=deg, rcond=rcond, full=full, w=_w(a), cov=cov))
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker, check_dtypes=False, tol=tol)
    self._CompileAndCheck(jnp_fun, args_maker, check_dtypes=False, atol=tol, rtol=tol)


  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_{}_amin={}_amax={}".format(
          jtu.format_shape_dtype_string(shape, dtype), a_min, a_max),
       "shape": shape, "dtype": dtype, "a_min": a_min, "a_max": a_max}
      for shape in all_shapes for dtype in number_dtypes
      for a_min, a_max in [(-1, None), (None, 1), (-0.9, 1),
                           (-np.ones(1), None),
                           (None, np.ones(1)),
                           (np.full(1, -0.9), np.ones(1))]))
  @jax.numpy_rank_promotion('allow')  # This test explicitly exercises implicit rank promotion.
  def testClipStaticBounds(self, shape, dtype, a_min, a_max):
    if np.issubdtype(dtype, np.unsignedinteger):
      a_min = None if a_min is None else abs(a_min)
      a_max = None if a_max is None else abs(a_max)
    rng = jtu.rand_default(self.rng())
    np_fun = lambda x: np.clip(x, a_min=a_min, a_max=a_max)
    jnp_fun = lambda x: jnp.clip(x, a_min=a_min, a_max=a_max)
    args_maker = lambda: [rng(shape, dtype)]
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker, check_dtypes=False)
    self._CompileAndCheck(jnp_fun, args_maker)

  def testClipError(self):
    with self.assertRaisesRegex(ValueError, "At most one of a_min and a_max.*"):
      jnp.clip(jnp.zeros((3,)))

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_{}_decimals={}".format(
          jtu.format_shape_dtype_string(shape, dtype), decimals),
       "shape": shape, "dtype": dtype, "decimals": decimals}
      for shape, dtype in _shape_and_dtypes(all_shapes, number_dtypes)
      for decimals in [0, 1, -2]))
  def testRoundStaticDecimals(self, shape, dtype, decimals):
    rng = jtu.rand_default(self.rng())
    if jnp.issubdtype(dtype, np.integer) and decimals < 0:
      self.skipTest("Integer rounding with decimals < 0 not implemented")
    np_fun = lambda x: np.round(x, decimals=decimals)
    jnp_fun = lambda x: jnp.round(x, decimals=decimals)
    args_maker = lambda: [rng(shape, dtype)]
    tol = {jnp.bfloat16: 5e-2, np.float16: 1e-2}
    check_dtypes = shape is not jtu.PYTHON_SCALAR_SHAPE
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker,
                            check_dtypes=check_dtypes, tol=tol)
    self._CompileAndCheck(jnp_fun, args_maker, check_dtypes=check_dtypes,
                          atol=tol, rtol=tol)

  def testOperatorRound(self):
    self.assertAllClose(round(np.float32(7.532), 1),
                        round(jnp.float32(7.5), 1))
    self.assertAllClose(round(np.float32(1.234), 2),
                        round(jnp.float32(1.234), 2))
    self.assertAllClose(round(np.float32(1.234)),
                        round(jnp.float32(1.234)), check_dtypes=False)
    self.assertAllClose(round(np.float32(7.532), 1),
                        round(jnp.array(7.5, jnp.float32), 1))
    self.assertAllClose(round(np.float32(1.234), 2),
                        round(jnp.array(1.234, jnp.float32), 2))
    self.assertAllClose(round(np.float32(1.234)),
                        round(jnp.array(1.234, jnp.float32)),
                        check_dtypes=False)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_shape={}_mode={}_padwidth={}_constantvalues={}".format(
          jtu.format_shape_dtype_string(shape, dtype), mode, pad_width,
          constant_values),
       "shape": shape, "dtype": dtype, "mode": mode,
       "pad_width": pad_width, "constant_values": constant_values}
      for mode, shapes in [
          ('constant', all_shapes),
          ('wrap', nonempty_shapes),
          ('edge', nonempty_shapes),
      ]
      for shape, dtype in _shape_and_dtypes(shapes, all_dtypes)
      for constant_values in [
          # None is used for modes other than 'constant'
          None,
          # constant
          0, 1,
          # (constant,)
          (0,), (2.718,),
          # ((before_const, after_const),)
          ((0, 2),), ((-1, 3.14),),
          # ((before_1, after_1), ..., (before_N, after_N))
          tuple((i / 2, -3.14 * i) for i in range(len(shape))),
      ]
      for pad_width in [
        # ((before_1, after_1), ..., (before_N, after_N))
        tuple((i % 3, (i + 1) % 3) for i in range(len(shape))),
        # ((before, after),)
        ((1, 2),), ((2, 0),),
        # (before, after)  (not in the docstring but works in numpy)
        (2, 0), (0, 0),
        # (pad,)
        (1,), (2,),
        # pad
        0, 1,
      ]
      if (pad_width != () and constant_values != () and
          ((mode == 'constant' and constant_values is not None) or
           (mode != 'constant' and constant_values is None)))))
  def testPad(self, shape, dtype, mode, pad_width, constant_values):
    if np.issubdtype(dtype, np.unsignedinteger):
      constant_values = tree_util.tree_map(abs, constant_values)
    rng = jtu.rand_default(self.rng())
    args_maker = lambda: [rng(shape, dtype)]
    if constant_values is None:
      np_fun = partial(np.pad, pad_width=pad_width, mode=mode)
      jnp_fun = partial(jnp.pad, pad_width=pad_width, mode=mode)
    else:
      np_fun = partial(np.pad, pad_width=pad_width, mode=mode,
                       constant_values=constant_values)
      jnp_fun = partial(jnp.pad, pad_width=pad_width, mode=mode,
                        constant_values=constant_values)

    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker,
                            check_dtypes=shape is not jtu.PYTHON_SCALAR_SHAPE)
    self._CompileAndCheck(jnp_fun, args_maker)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_shape={}_mode={}_pad_width={}_stat_length={}".format(
          jtu.format_shape_dtype_string(shape, dtype), mode, pad_width, stat_length),
       "shape": shape, "dtype": dtype, "mode": mode, "pad_width": pad_width,
       "stat_length": stat_length}
      for mode in ['maximum', 'minimum', 'mean', 'median']
      for shape, dtype in _shape_and_dtypes(nonempty_shapes, all_dtypes)
      for pad_width in [
          # ((before_1, after_1), ..., (before_N, after_N))
          tuple((i % 3, (i + 1) % 3) for i in range(len(shape))),
          # ((before, after),)
          ((1, 2),), ((2, 0),),
          # (before, after)  (not in the docstring but works in numpy)
          (2, 0), (0, 0),
          # (pad,)
          (1,), (2,),
          # pad
          0, 1,
      ]
      for stat_length in [
          None,
          # ((before_1, after_1), ..., (before_N, after_N))
          tuple(((i % 3 + 1), ((i + 1) % 3) + 1) for i in range(len(shape))),
          # ((before, after),)
          ((1, 2),), ((2, 2),),
          # (before, after)  (not in the docstring but works in numpy)
          (1, 1), (3, 4),
          # (pad,)
          (1,), (2,),
          # pad
          1, 2
      ]
      if (pad_width != () and stat_length != () and
          not (dtype in bool_dtypes and mode == 'mean'))))
  def testPadStatValues(self, shape, dtype, mode, pad_width, stat_length):
    if mode == 'median' and np.issubdtype(dtype, np.complexfloating):
      self.skipTest("median statistic is not supported for dtype=complex.")
    rng = jtu.rand_default(self.rng())
    args_maker = lambda: [rng(shape, dtype)]

    np_fun = partial(np.pad, pad_width=pad_width, mode=mode, stat_length=stat_length)
    jnp_fun = partial(jnp.pad, pad_width=pad_width, mode=mode, stat_length=stat_length)

    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker,
                            check_dtypes=shape is not jtu.PYTHON_SCALAR_SHAPE)
    self._CompileAndCheck(jnp_fun, args_maker)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_shape={}_mode={}_pad_width={}_reflect_type={}".format(
          jtu.format_shape_dtype_string(shape, dtype), mode, pad_width, reflect_type),
       "shape": shape, "dtype": dtype, "mode": mode, "pad_width": pad_width,
       "reflect_type": reflect_type}
      for mode in ['symmetric', 'reflect']
      for shape, dtype in _shape_and_dtypes(nonempty_shapes, all_dtypes)
      for pad_width in [
          # ((before_1, after_1), ..., (before_N, after_N))
          tuple((i % 3, (i + 1) % 3) for i in range(len(shape))),
          # ((before, after),)
          ((1, 2),), ((2, 3),),
          # (before, after)  (not in the docstring but works in numpy)
          (2, 1), (1, 2),
          # (pad,)
          (1,), (2,), (3,),
          # pad
          0, 5, 7, 10
      ]
      for reflect_type in ['even', 'odd']
      if (pad_width != () and
          # following types lack precision when calculating odd values
          (reflect_type != 'odd' or dtype not in [np.bool_, np.float16, jnp.bfloat16]))))
  def testPadSymmetricAndReflect(self, shape, dtype, mode, pad_width, reflect_type):
    rng = jtu.rand_default(self.rng())
    args_maker = lambda: [rng(shape, dtype)]

    np_fun = partial(np.pad, pad_width=pad_width, mode=mode, reflect_type=reflect_type)
    jnp_fun = partial(jnp.pad, pad_width=pad_width, mode=mode, reflect_type=reflect_type)

    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker,
                            check_dtypes=shape is not jtu.PYTHON_SCALAR_SHAPE,
                            tol={np.float32: 1e-3, np.complex64: 1e-3})
    self._CompileAndCheck(jnp_fun, args_maker)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_shape={}_mode={}_pad_width={}_end_values={}".format(
          jtu.format_shape_dtype_string(shape, dtype), "linear_ramp", pad_width, end_values),
       "shape": shape, "dtype": dtype, "pad_width": pad_width,
       "end_values": end_values}
      for shape, dtype in _shape_and_dtypes(nonempty_shapes, default_dtypes + complex_dtypes)
      for pad_width in [
        # ((before_1, after_1), ..., (before_N, after_N))
        tuple((i % 3, (i + 1) % 3) for i in range(len(shape))),
        # ((before, after),)
        ((1, 2),), ((2, 0),),
        # (before, after)  (not in the docstring but works in numpy)
        (2, 0), (0, 0),
        # (pad,)
        (1,), (2,),
        # pad
        0, 1,
      ]
      for end_values in [
        # ((before_1, after_1), ..., (before_N, after_N))
        tuple((i % 3, (i + 1) % 3) for i in range(len(shape))),
        # ((before, after),)
        ((1, 2),), ((2.0, 3.14),),
        # (before, after)  (not in the docstring but works in numpy)
        (0, 0), (-8.0, 2.0),
        # (end_values,)
        (1,), (2,),
        # end_values
        0, 1, 100, 10.0, 3.5, 4.2, -5, -3
      ]
      if (pad_width != () and end_values != () and
          # following types lack precision
          dtype not in [np.int8, np.int16, np.float16, jnp.bfloat16])))
  def testPadLinearRamp(self, shape, dtype, pad_width, end_values):
    if numpy_version < (1, 20) and np.issubdtype(dtype, np.integer):
      raise unittest.SkipTest("NumPy 1.20 changed the semantics of np.linspace")
    rng = jtu.rand_default(self.rng())
    args_maker = lambda: [rng(shape, dtype)]

    np_fun = partial(np.pad, pad_width=pad_width, mode="linear_ramp",
                     end_values=end_values)
    jnp_fun = partial(jnp.pad, pad_width=pad_width, mode="linear_ramp",
                      end_values=end_values)

    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker,
                            check_dtypes=shape is not jtu.PYTHON_SCALAR_SHAPE)
    self._CompileAndCheck(jnp_fun, args_maker)

  def testPadEmpty(self):
    arr = np.arange(6).reshape(2, 3)

    pad_width = ((2, 3), (3, 1))
    np_res = np.pad(arr, pad_width=pad_width, mode="empty")
    jnp_res = jnp.pad(arr, pad_width=pad_width, mode="empty")

    np.testing.assert_equal(np_res.shape, jnp_res.shape)
    np.testing.assert_equal(arr, np_res[2:-3, 3:-1])
    np.testing.assert_equal(arr, jnp_res[2:-3, 3:-1])
    np.testing.assert_equal(np_res[2:-3, 3:-1], jnp_res[2:-3, 3:-1])

  def testPadKwargs(self):
    modes = {
        'constant': {'constant_values': 0},
        'edge': {},
        'linear_ramp': {'end_values': 0},
        'maximum': {'stat_length': None},
        'mean': {'stat_length': None},
        'median': {'stat_length': None},
        'minimum': {'stat_length': None},
        'reflect': {'reflect_type': 'even'},
        'symmetric': {'reflect_type': 'even'},
        'wrap': {},
        'empty': {}
    }
    arr = jnp.array([1, 2, 3])
    pad_width = 1

    for mode in modes.keys():
      allowed = modes[mode]
      not_allowed = {}
      for kwargs in modes.values():
        if kwargs != allowed:
          not_allowed.update(kwargs)

      # Test if allowed keyword arguments pass
      jnp.pad(arr, pad_width, mode, **allowed)
      # Test if prohibited keyword arguments of other modes raise an error
      match = f"unsupported keyword arguments for mode '{mode}'"
      for key, value in not_allowed.items():
        with self.assertRaisesRegex(ValueError, match):
          jnp.pad(arr, pad_width, mode, **{key: value})

    # Test if unsupported mode raise error.
    unsupported_modes = [1, None, "foo"]
    for mode in unsupported_modes:
      match = f"Unimplemented padding mode '{mode}' for np.pad."
      with self.assertRaisesRegex(NotImplementedError, match):
        jnp.pad(arr, pad_width, mode)

  def testPadFunction(self):
    def np_pad_with(vector, pad_width, iaxis, kwargs):
      pad_value = kwargs.get('padder', 10)
      vector[:pad_width[0]] = pad_value
      vector[-pad_width[1]:] = pad_value

    def jnp_pad_with(vector, pad_width, iaxis, kwargs):
      pad_value = kwargs.get('padder', 10)
      vector = vector.at[:pad_width[0]].set(pad_value)
      vector = vector.at[-pad_width[1]:].set(pad_value)
      return vector

    arr = np.arange(6).reshape(2, 3)
    np_res = np.pad(arr, 2, np_pad_with)
    jnp_res = jnp.pad(arr, 2, jnp_pad_with)
    np.testing.assert_equal(np_res, jnp_res)

    arr = np.arange(24).reshape(2, 3, 4)
    np_res = np.pad(arr, 1, np_pad_with, padder=100)
    jnp_res = jnp.pad(arr, 1, jnp_pad_with, padder=100)
    np.testing.assert_equal(np_res, jnp_res)

    rng = jtu.rand_default(self.rng())
    args_maker = lambda: [rng(arr.shape, arr.dtype)]
    jnp_fun = partial(jnp.pad, pad_width=1, mode=jnp_pad_with)
    self._CompileAndCheck(jnp_fun, args_maker)

  def testPadWithNumpyPadWidth(self):
    a = jnp.array([1, 2, 3, 4, 5])
    f = jax.jit(
        partial(
            jnp.pad,
            pad_width=np.asarray((2, 3)),
            mode="constant",
            constant_values=(4, 6)))

    np.testing.assert_array_equal(
        f(a),
        np.pad(
            a,
            pad_width=np.asarray((2, 3)),
            mode="constant",
            constant_values=(4, 6)))

  def testPadWeakType(self):
    x = jnp.array(1.0)[None]
    for mode in ['constant', 'edge', 'linear_ramp', 'maximum', 'mean', 'median',
                 'minimum', 'reflect', 'symmetric', 'wrap', 'empty']:
      y = jnp.pad(x, 0, mode=mode)
      self.assertTrue(dtypes.is_weakly_typed(y))

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_shape=[{}]_reps={}".format(
          jtu.format_shape_dtype_string(shape, dtype), reps),
       "shape": shape, "dtype": dtype, "reps": reps}
      for reps in [(), (2,), (3, 4), (2, 3, 4), (1, 0, 2)]
      for shape, dtype in _shape_and_dtypes(all_shapes, default_dtypes)
      ))
  def testTile(self, shape, dtype, reps):
    rng = jtu.rand_default(self.rng())
    np_fun = lambda arg: np.tile(arg, reps)
    jnp_fun = lambda arg: jnp.tile(arg, reps)

    args_maker = lambda: [rng(shape, dtype)]

    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker,
                            check_dtypes=shape is not jtu.PYTHON_SCALAR_SHAPE)
    self._CompileAndCheck(jnp_fun, args_maker)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_shape={}".format(
          jtu.format_shape_dtype_string(shape, dtype)),
       "shape": shape, "dtype": dtype}
      for shape in all_shapes
      for dtype in all_dtypes))
  def testExtract(self, shape, dtype):
    rng = jtu.rand_some_zero(self.rng())
    args_maker = lambda: [rng(shape, jnp.float32), rng(shape, dtype)]
    self._CheckAgainstNumpy(np.extract, jnp.extract, args_maker)


  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_{}_ncond={}_nfunc={}".format(
          jtu.format_shape_dtype_string(shape, dtype), ncond, nfunc),
       "shape": shape, "dtype": dtype, "ncond": ncond, "nfunc": nfunc}
      for ncond in [1, 2, 3]
      for nfunc in [ncond, ncond + 1]
      for shape in all_shapes
      for dtype in all_dtypes))
  def testPiecewise(self, shape, dtype, ncond, nfunc):
    rng = jtu.rand_default(self.rng())
    rng_bool = jtu.rand_int(self.rng(), 0, 2)
    funclist = [lambda x: x - 1, 1, lambda x: x, 0][:nfunc]
    args_maker = lambda: (rng(shape, dtype), [rng_bool(shape, bool) for i in range(ncond)])
    np_fun = partial(np.piecewise, funclist=funclist)
    jnp_fun = partial(jnp.piecewise, funclist=funclist)
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker, check_dtypes=True)
    # This is a higher-order function, so the cache miss check will fail.
    self._CompileAndCheck(jnp_fun, args_maker, check_dtypes=True, check_cache_misses=False)

  def testPiecewiseRecompile(self):
    def g(x):
      g.num_traces += 1
      return x
    g.num_traces = 0
    x = jnp.arange(10.0)
    for i in range(5):
      jnp.piecewise(x, [x < 0], [g, 0.])
    self.assertEqual(g.num_traces, 1)

  @parameterized.named_parameters(jtu.cases_from_list(
    {"testcase_name": "{}_perm={}_{}".format(
      jtu.format_shape_dtype_string(shape, dtype), perm, arg_type),
     "dtype": dtype, "shape": shape, "perm": perm, "arg_type": arg_type}
    for dtype in default_dtypes
    for shape in array_shapes
    for arg_type in ["splat", "value"]
    for perm in [None, tuple(np.random.RandomState(0).permutation(np.zeros(shape).ndim))]))
  def testTransposeTuple(self, shape, dtype, perm, arg_type):
    rng = jtu.rand_some_zero(self.rng())
    args_maker = lambda: [rng(shape, dtype)]
    if arg_type == "value":
      np_fun = lambda x: x.transpose(perm)
      jnp_fun = lambda x: jnp.array(x).transpose(perm)
    else:
      np_fun = lambda x: x.transpose(*(perm or ()))
      jnp_fun = lambda x: jnp.array(x).transpose(*(perm or ()))

    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker, check_dtypes=True)
    self._CompileAndCheck(jnp_fun, args_maker, check_dtypes=True)


  @parameterized.named_parameters(jtu.cases_from_list(
    {"testcase_name": "{}_trim={}".format(
      jtu.format_shape_dtype_string(a_shape, dtype), trim),
     "dtype": dtype, "a_shape": a_shape, "trim": trim}
    for dtype in default_dtypes
    for a_shape in one_dim_array_shapes
    for trim in ["f", "b", "fb"]))
  def testTrimZeros(self, a_shape, dtype, trim):
    rng = jtu.rand_some_zero(self.rng())
    args_maker = lambda: [rng(a_shape, dtype)]
    np_fun = lambda arg1: np.trim_zeros(arg1, trim)
    jnp_fun = lambda arg1: jnp.trim_zeros(arg1, trim)
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker, check_dtypes=True)


  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_shape={}_rank{}".format(
        jtu.format_shape_dtype_string(a_shape, dtype), rank),
       "dtype": dtype, "a_shape": a_shape, "rank": rank}
      for rank in (1, 2)
      for dtype in default_dtypes
      for a_shape in one_dim_array_shapes))
  def testPoly(self, a_shape, dtype, rank):
    if dtype in (np.float16, jnp.bfloat16, np.int16):
      self.skipTest(f"{dtype} gets promoted to {np.float16}, which is not supported.")
    elif rank == 2 and jtu.device_under_test() in ("tpu", "gpu"):
      self.skipTest("Nonsymmetric eigendecomposition is only implemented on the CPU backend.")
    rng = jtu.rand_default(self.rng())
    tol = { np.int8: 1e-3, np.int32: 1e-3, np.float32: 1e-3, np.float64: 1e-6 }
    if jtu.device_under_test() == "tpu":
      tol[np.int32] = tol[np.float32] = 1e-1
    tol = jtu.tolerance(dtype, tol)
    args_maker = lambda: [rng(a_shape * rank, dtype)]
    self._CheckAgainstNumpy(np.poly, jnp.poly, args_maker, check_dtypes=False, tol=tol)
    self._CompileAndCheck(jnp.poly, args_maker, check_dtypes=True, rtol=tol, atol=tol)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "a_shape={} , b_shape={}".format(
          jtu.format_shape_dtype_string(a_shape, dtype),
          jtu.format_shape_dtype_string(b_shape, dtype)),
       "dtype": dtype, "a_shape": a_shape, "b_shape" : b_shape}
      for dtype in default_dtypes
      for a_shape in one_dim_array_shapes
      for b_shape in one_dim_array_shapes))
  def testPolyAdd(self, a_shape, b_shape, dtype):
    rng = jtu.rand_default(self.rng())
    np_fun = lambda arg1, arg2: np.polyadd(arg1, arg2)
    jnp_fun = lambda arg1, arg2: jnp.polyadd(arg1, arg2)
    args_maker = lambda: [rng(a_shape, dtype), rng(b_shape, dtype)]
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker, check_dtypes=True)
    self._CompileAndCheck(jnp_fun, args_maker, check_dtypes=True)


  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "a_shape={} , b_shape={}".format(
          jtu.format_shape_dtype_string(a_shape, dtype),
          jtu.format_shape_dtype_string(b_shape, dtype)),
       "dtype": dtype, "a_shape": a_shape, "b_shape" : b_shape}
      for dtype in default_dtypes
      for a_shape in one_dim_array_shapes
      for b_shape in one_dim_array_shapes))
  def testPolySub(self, a_shape, b_shape, dtype):
    rng = jtu.rand_default(self.rng())
    np_fun = lambda arg1, arg2: np.polysub(arg1, arg2)
    jnp_fun = lambda arg1, arg2: jnp.polysub(arg1, arg2)
    args_maker = lambda: [rng(a_shape, dtype), rng(b_shape, dtype)]
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker, check_dtypes=True)
    self._CompileAndCheck(jnp_fun, args_maker, check_dtypes=True)


  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_shape={}_order={}_k={}".format(
          jtu.format_shape_dtype_string(a_shape, dtype),
          order, k),
       "dtype": dtype, "a_shape": a_shape, "order" : order, "k": k}
      for dtype in default_dtypes
      for a_shape in one_dim_array_shapes
      for order in range(5)
      for k in [np.arange(order, dtype=dtype), np.ones(1, dtype), None]))
  def testPolyInt(self, a_shape, order, k, dtype):
    rng = jtu.rand_default(self.rng())
    np_fun = lambda arg1: np.polyint(arg1, m=order, k=k)
    jnp_fun = lambda arg1: jnp.polyint(arg1, m=order, k=k)
    args_maker = lambda: [rng(a_shape, dtype)]
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker, check_dtypes=False)
    self._CompileAndCheck(jnp_fun, args_maker, check_dtypes=True)


  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_shape={}_order={}".format(
          jtu.format_shape_dtype_string(a_shape, dtype),
          order),
       "dtype": dtype, "a_shape": a_shape, "order" : order}
      for dtype in default_dtypes
      for a_shape in one_dim_array_shapes
      for order in range(5)))
  def testPolyDer(self, a_shape, order, dtype):
    rng = jtu.rand_default(self.rng())
    np_fun = lambda arg1: np.polyder(arg1, m=order)
    jnp_fun = lambda arg1: jnp.polyder(arg1, m=order)
    args_maker = lambda: [rng(a_shape, dtype)]
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker, check_dtypes=False)
    self._CompileAndCheck(jnp_fun, args_maker, check_dtypes=True)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": f"_ptype={ptype}", "ptype": ptype}
      for ptype in ['int', 'np.int', 'jnp.int']))
  def testIntegerPower(self, ptype):
    p = {'int': 2, 'np.int': np.int32(2), 'jnp.int': jnp.int32(2)}[ptype]
    jaxpr = jax.make_jaxpr(partial(jnp.power, x2=p))(1)
    eqns = jaxpr.jaxpr.eqns
    self.assertLen(eqns, 1)
    self.assertEqual(eqns[0].primitive, lax.integer_pow_p)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": f"_x={x}_y={y}", "x": x, "y": y}
      for x in [-1, 0, 1]
      for y in [0, 32, 64, 128]))
  def testIntegerPowerOverflow(self, x, y):
    # Regression test for https://github.com/google/jax/issues/5987
    args_maker = lambda: [x, y]
    self._CheckAgainstNumpy(np.power, jnp.power, args_maker)
    self._CompileAndCheck(jnp.power, args_maker)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_shape={}_axis={}".format(
          jtu.format_shape_dtype_string(shape, dtype), axis),
       "shape": shape, "dtype": dtype, "axis": axis}
      for shape in all_shapes
      for dtype in all_dtypes
      for axis in [None] + list(range(len(shape)))))
  def testCompress(self, shape, dtype, axis):
    rng = jtu.rand_some_zero(self.rng())
    if shape in scalar_shapes or len(shape) == 0:
      cond_shape = (0,)
    elif axis is None:
      cond_shape = (prod(shape),)
    else:
      cond_shape = (shape[axis],)

    args_maker = lambda: [rng(cond_shape, jnp.float32), rng(shape, dtype)]

    np_fun = partial(np.compress, axis=axis)
    jnp_fun = partial(jnp.compress, axis=axis)
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_shape={}_condition=array[{}]_axis={}".format(
          jtu.format_shape_dtype_string(shape, dtype), len(condition), axis),
       "shape": shape, "dtype": dtype, "condition": condition, "axis": axis}
      for shape in [(2, 3)]
      for dtype in int_dtypes
      # condition entries beyond axis size must be zero.
      for condition in [[1], [1, 0, 0, 0, 0, 0, 0]]
      for axis in [None, 0, 1]))
  def testCompressMismatchedShapes(self, shape, dtype, condition, axis):
    rng = jtu.rand_default(self.rng())
    args_maker = lambda: [np.array(condition), rng(shape, dtype)]
    np_fun = partial(np.compress, axis=axis)
    jnp_fun = partial(jnp.compress, axis=axis)
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_shape={}_axis={}".format(
          jtu.format_shape_dtype_string(shape, dtype), axis),
       "shape": shape, "dtype": dtype, "axis": axis}
      for shape in array_shapes
      for dtype in all_dtypes
      for axis in [None] + list(range(len(shape)))))
  def testCompressMethod(self, shape, dtype, axis):
    rng = jtu.rand_some_zero(self.rng())
    if shape in scalar_shapes or len(shape) == 0:
      cond_shape = (0,)
    elif axis is None:
      cond_shape = (prod(shape),)
    else:
      cond_shape = (shape[axis],)

    args_maker = lambda: [rng(cond_shape, jnp.float32), rng(shape, dtype)]

    np_fun = lambda condition, x: np.compress(condition, x, axis=axis)
    jnp_fun = lambda condition, x: x.compress(condition, axis=axis)
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_axis={}_baseshape=[{}]_dtypes=[{}]".format(
          axis, ",".join(str(d) for d in base_shape),
          ",".join(np.dtype(dtype).name for dtype in arg_dtypes)),
       "axis": axis, "base_shape": base_shape, "arg_dtypes": arg_dtypes}
      for num_arrs in [3]
      for arg_dtypes in itertools.combinations_with_replacement(default_dtypes, num_arrs)
      for base_shape in [(4,), (3, 4), (2, 3, 4)]
      for axis in range(-len(base_shape)+1, len(base_shape))))
  def testConcatenate(self, axis, base_shape, arg_dtypes):
    rng = jtu.rand_default(self.rng())
    wrapped_axis = axis % len(base_shape)
    shapes = [base_shape[:wrapped_axis] + (size,) + base_shape[wrapped_axis+1:]
              for size, _ in zip(itertools.cycle([3, 1, 4]), arg_dtypes)]
    def np_fun(*args):
      args = [x if x.dtype != jnp.bfloat16 else x.astype(np.float32)
              for x in args]
      dtype = functools.reduce(jnp.promote_types, arg_dtypes)
      return np.concatenate(args, axis=axis).astype(dtype)
    jnp_fun = lambda *args: jnp.concatenate(args, axis=axis)

    def args_maker():
      return [rng(shape, dtype) for shape, dtype in zip(shapes, arg_dtypes)]

    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)
    self._CompileAndCheck(jnp_fun, args_maker)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_{}_axis={}".format(
          jtu.format_shape_dtype_string(shape, dtype), axis),
        "shape": shape, "dtype": dtype, "axis": axis}
      for shape in [(4, 1), (4, 3), (4, 5, 6)]
      for dtype in all_dtypes
      for axis in [None] + list(range(1 - len(shape), len(shape) - 1))))
  def testConcatenateArray(self, shape, dtype, axis):
    rng = jtu.rand_default(self.rng())
    args_maker = lambda: [rng(shape, dtype)]
    np_fun = lambda x: np.concatenate(x, axis=axis)
    jnp_fun = lambda x: jnp.concatenate(x, axis=axis)
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)
    self._CompileAndCheck(jnp_fun, args_maker)

  def testConcatenateAxisNone(self):
    # https://github.com/google/jax/issues/3419
    a = jnp.array([[1, 2], [3, 4]])
    b = jnp.array([[5]])
    jnp.concatenate((a, b), axis=None)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_axis={}_baseshape=[{}]_dtypes=[{}]".format(
          axis, ",".join(str(d) for d in base_shape),
          ",".join(np.dtype(dtype).name for dtype in arg_dtypes)),
       "axis": axis, "base_shape": base_shape, "arg_dtypes": arg_dtypes}
      for arg_dtypes in itertools.combinations_with_replacement(default_dtypes, 2)
      for base_shape in [(4,), (3, 4), (2, 3, 4)]
      for axis in range(-len(base_shape)+1, len(base_shape))))
  def testAppend(self, axis, base_shape, arg_dtypes):
    rng = jtu.rand_default(self.rng())
    wrapped_axis = axis % len(base_shape)
    shapes = [base_shape[:wrapped_axis] + (size,) + base_shape[wrapped_axis+1:]
              for size, _ in zip(itertools.cycle([3, 1, 4]), arg_dtypes)]
    def np_fun(arr, values):
      arr = arr.astype(np.float32) if arr.dtype == jnp.bfloat16 else arr
      values = (values.astype(np.float32) if values.dtype == jnp.bfloat16
                else values)
      out = np.append(arr, values, axis=axis)
      return out.astype(jnp.promote_types(*arg_dtypes))
    jnp_fun = lambda arr, values: jnp.append(arr, values, axis=axis)

    def args_maker():
      return [rng(shape, dtype) for shape, dtype in zip(shapes, arg_dtypes)]

    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)
    self._CompileAndCheck(jnp_fun, args_maker)


  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_{}_axis={}_idx={}".format(
          jtu.format_shape_dtype_string(shape, dtype), axis, idx),
       "dtype": dtype, "shape": shape, "axis": axis, "idx": idx}
      for shape in nonempty_nonscalar_array_shapes
      for dtype in all_dtypes
      for axis in [None] + list(range(-len(shape), len(shape)))
      for idx in (range(-prod(shape), prod(shape))
                  if axis is None else
                  range(-shape[axis], shape[axis]))))
  def testDeleteInteger(self, shape, dtype, idx, axis):
    rng = jtu.rand_default(self.rng())
    args_maker = lambda: [rng(shape, dtype)]
    np_fun = lambda arg: np.delete(arg, idx, axis=axis)
    jnp_fun = lambda arg: jnp.delete(arg, idx, axis=axis)
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)
    self._CompileAndCheck(jnp_fun, args_maker)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_{}_axis={}_slc={}".format(
          jtu.format_shape_dtype_string(shape, dtype), axis, slc),
       "dtype": dtype, "shape": shape, "axis": axis, "slc": slc}
      for shape in nonempty_nonscalar_array_shapes
      for dtype in all_dtypes
      for axis in [None] + list(range(-len(shape), len(shape)))
      for slc in [slice(None), slice(1, 3), slice(1, 5, 2)]))
  def testDeleteSlice(self, shape, dtype, axis, slc):
    rng = jtu.rand_default(self.rng())
    args_maker = lambda: [rng(shape, dtype)]
    np_fun = lambda arg: np.delete(arg, slc, axis=axis)
    jnp_fun = lambda arg: jnp.delete(arg, slc, axis=axis)
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)
    self._CompileAndCheck(jnp_fun, args_maker)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_{}_axis={}_idx={}".format(
          jtu.format_shape_dtype_string(shape, dtype), axis,
          jtu.format_shape_dtype_string(idx_shape, int)),
       "dtype": dtype, "shape": shape, "axis": axis, "idx_shape": idx_shape}
      for shape in nonempty_nonscalar_array_shapes
      for dtype in all_dtypes
      for axis in [None] + list(range(-len(shape), len(shape)))
      for idx_shape in all_shapes))
  def testDeleteIndexArray(self, shape, dtype, axis, idx_shape):
    rng = jtu.rand_default(self.rng())
    max_idx = np.zeros(shape).size if axis is None else np.zeros(shape).shape[axis]
    # Previous to numpy 1.19, negative indices were ignored so we don't test this.
    low = 0 if numpy_version < (1, 19, 0) else -max_idx
    idx = jtu.rand_int(self.rng(), low=low, high=max_idx)(idx_shape, int)
    args_maker = lambda: [rng(shape, dtype)]
    np_fun = lambda arg: np.delete(arg, idx, axis=axis)
    jnp_fun = lambda arg: jnp.delete(arg, idx, axis=axis)
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)
    self._CompileAndCheck(jnp_fun, args_maker)

  @unittest.skipIf(numpy_version < (1, 19), "boolean mask not supported in numpy < 1.19.0")
  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_{}_axis={}".format(
          jtu.format_shape_dtype_string(shape, dtype), axis),
       "dtype": dtype, "shape": shape, "axis": axis}
      for shape in nonempty_nonscalar_array_shapes
      for dtype in all_dtypes
      for axis in [None] + list(range(-len(shape), len(shape)))))
  def testDeleteMaskArray(self, shape, dtype, axis):
    rng = jtu.rand_default(self.rng())
    mask_size = np.zeros(shape).size if axis is None else np.zeros(shape).shape[axis]
    mask = jtu.rand_int(self.rng(), low=0, high=2)(mask_size, bool)
    args_maker = lambda: [rng(shape, dtype)]
    np_fun = lambda arg: np.delete(arg, mask, axis=axis)
    jnp_fun = lambda arg: jnp.delete(arg, mask, axis=axis)
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)
    self._CompileAndCheck(jnp_fun, args_maker)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_{}_axis={}".format(
          jtu.format_shape_dtype_string(shape, dtype), axis),
       "dtype": dtype, "shape": shape, "axis": axis}
      for shape in nonempty_nonscalar_array_shapes
      for dtype in all_dtypes
      for axis in [None] + list(range(-len(shape), len(shape)))))
  def testInsertInteger(self, shape, dtype, axis):
    x = jnp.empty(shape)
    max_ind = x.size if axis is None else x.shape[axis]
    rng = jtu.rand_default(self.rng())
    i_rng = jtu.rand_int(self.rng(), -max_ind, max_ind)
    args_maker = lambda: [rng(shape, dtype), i_rng((), np.int32), rng((), dtype)]
    np_fun = lambda *args: np.insert(*args, axis=axis)
    jnp_fun = lambda *args: jnp.insert(*args, axis=axis)
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)
    self._CompileAndCheck(jnp_fun, args_maker)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_{}_axis={}".format(
          jtu.format_shape_dtype_string(shape, dtype), axis),
       "dtype": dtype, "shape": shape, "axis": axis}
      for shape in nonempty_nonscalar_array_shapes
      for dtype in all_dtypes
      for axis in [None] + list(range(-len(shape), len(shape)))))
  def testInsertSlice(self, shape, dtype, axis):
    x = jnp.empty(shape)
    max_ind = x.size if axis is None else x.shape[axis]
    rng = jtu.rand_default(self.rng())
    i_rng = jtu.rand_int(self.rng(), -max_ind, max_ind)
    slc = slice(i_rng((), jnp.int32).item(), i_rng((), jnp.int32).item())
    args_maker = lambda: [rng(shape, dtype), rng((), dtype)]
    np_fun = lambda x, val: np.insert(x, slc, val, axis=axis)
    jnp_fun = lambda x, val: jnp.insert(x, slc, val, axis=axis)
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)
    self._CompileAndCheck(jnp_fun, args_maker)

  @parameterized.parameters([
    [[[1, 1], [2, 2], [3, 3]], 1, 5, None],
    [[[1, 1], [2, 2], [3, 3]], 1, 5, 1],
    [[[1, 1], [2, 2], [3, 3]], 1, [1, 2, 3], 1],
    [[[1, 1], [2, 2], [3, 3]], [1], [[1],[2],[3]], 1],
    [[1, 1, 2, 2, 3, 3], [2, 2], [5, 6], None],
    [[1, 1, 2, 2, 3, 3], slice(2, 4), [5, 6], None],
    [[1, 1, 2, 2, 3, 3], [2, 2], [7.13, False], None],
    [[[0, 1, 2, 3], [4, 5, 6, 7]], (1, 3), 999, 1]
  ])
  def testInsertExamples(self, arr, index, values, axis):
    # Test examples from the np.insert docstring
    args_maker = lambda: (
      np.asarray(arr), index if isinstance(index, slice) else np.array(index),
      np.asarray(values), axis)
    self._CheckAgainstNumpy(np.insert, jnp.insert, args_maker)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_{}_axis={}_out_dims={}".format(
          jtu.format_shape_dtype_string(shape, dtype),
          axis, out_dims),
        "shape": shape, "dtype": dtype, "axis": axis, "out_dims": out_dims}
      for shape in nonempty_array_shapes
      for dtype in default_dtypes
      for axis in range(-len(shape), len(shape))
      for out_dims in [0, 1, 2]))
  def testApplyAlongAxis(self, shape, dtype, axis, out_dims):
    def func(x, out_dims):
      if out_dims == 0:
        return x.sum()
      elif out_dims == 1:
        return x * x[0]
      elif out_dims == 2:
        return x[:, None] + x[None, :]
      else:
        raise NotImplementedError(f"out_dims={out_dims}")
    rng = jtu.rand_default(self.rng())
    args_maker = lambda: [rng(shape, dtype)]
    np_fun = lambda arr: np.apply_along_axis(func, axis, arr, out_dims=out_dims)
    jnp_fun = lambda arr: jnp.apply_along_axis(func, axis, arr, out_dims=out_dims)
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)
    self._CompileAndCheck(jnp_fun, args_maker)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_{}_func={}_keepdims={}_axes={}".format(
          jtu.format_shape_dtype_string(shape, dtype),
          func, keepdims, axes),
        "shape": shape, "dtype": dtype, "func": func, "keepdims": keepdims, "axes": axes}
      for shape in nonempty_shapes
      for func in ["sum"]
      for keepdims in [True, False]
      for axes in itertools.combinations(range(len(shape)), 2)
      # Avoid low-precision types in sum()
      for dtype in default_dtypes if dtype not in [np.float16, jnp.bfloat16]))
  def testApplyOverAxes(self, shape, dtype, func, keepdims, axes):
    f = lambda x, axis: getattr(x, func)(axis=axis, keepdims=keepdims)
    rng = jtu.rand_default(self.rng())
    args_maker = lambda: (rng(shape, dtype),)
    np_fun = lambda a: np.apply_over_axes(f, a, axes)
    jnp_fun = lambda a: jnp.apply_over_axes(f, a, axes)
    self._CompileAndCheck(jnp_fun, args_maker)
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_shape=[{}]_axis={}_repeats={}_fixed_size={}".format(
          jtu.format_shape_dtype_string(shape, dtype),
          axis, repeats, fixed_size),
       "axis": axis, "shape": shape, "dtype": dtype, "repeats": repeats,
       'fixed_size': fixed_size}
      for repeats in [0, 1, 2]
      for shape, dtype in _shape_and_dtypes(all_shapes, default_dtypes)
      for axis in [None] + list(range(-len(shape), max(1, len(shape))))
      for fixed_size in [True, False]))
  def testRepeat(self, axis, shape, dtype, repeats, fixed_size):
    rng = jtu.rand_default(self.rng())
    np_fun = lambda arg: np.repeat(arg, repeats=repeats, axis=axis)
    np_fun = _promote_like_jnp(np_fun)
    if fixed_size:
      total_repeat_length = np.repeat(np.zeros(shape), repeats, axis).shape[axis or 0]
      jnp_fun = lambda arg, rep: jnp.repeat(arg, repeats=rep, axis=axis,
                                     total_repeat_length=total_repeat_length)
      jnp_args_maker = lambda: [rng(shape, dtype), repeats]
      clo_fun = lambda arg: jnp.repeat(arg, repeats=repeats, axis=axis,
                                       total_repeat_length=total_repeat_length)
      clo_fun_args_maker = lambda: [rng(shape, dtype)]
      self._CompileAndCheck(jnp_fun, jnp_args_maker)
      self._CheckAgainstNumpy(np_fun, clo_fun, clo_fun_args_maker)
    else:
      # Now repeats is in a closure, so a constant.
      jnp_fun = lambda arg: jnp.repeat(arg, repeats=repeats, axis=axis)
      args_maker = lambda: [rng(shape, dtype)]
      self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)
      self._CompileAndCheck(jnp_fun, args_maker)

  def testRepeatScalarFastPath(self):
    a = jnp.array([1,2,3,4])
    f = lambda a: jnp.repeat(a, repeats=2)
    jaxpr = jax.make_jaxpr(f)(a)
    self.assertLessEqual(len(jaxpr.jaxpr.eqns), 6)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_{}_axis={}_ind={}_inv={}_count={}".format(
          jtu.format_shape_dtype_string(shape, dtype), axis,
          return_index, return_inverse, return_counts),
       "shape": shape, "dtype": dtype, "axis": axis,
       "return_index": return_index, "return_inverse": return_inverse,
       "return_counts": return_counts}
      for dtype in number_dtypes
      for shape in all_shapes
      for axis in [None] + list(range(len(shape)))
      for return_index in [False, True]
      for return_inverse in [False, True]
      for return_counts in [False, True]))
  def testUnique(self, shape, dtype, axis, return_index, return_inverse, return_counts):
    if axis is not None and numpy_version < (1, 19) and np.empty(shape).size == 0:
      self.skipTest("zero-sized axis in unique leads to error in older numpy.")
    rng = jtu.rand_some_equal(self.rng())
    args_maker = lambda: [rng(shape, dtype)]
    extra_args = (return_index, return_inverse, return_counts)
    use_defaults =  (False, *(True for arg in extra_args if arg)) if any(extra_args) else False
    np_fun = jtu.with_jax_dtype_defaults(lambda x: np.unique(x, *extra_args, axis=axis), use_defaults)
    jnp_fun = lambda x: jnp.unique(x, *extra_args, axis=axis)
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_{}_axis={}_size={}_fill_value={}".format(
         jtu.format_shape_dtype_string(shape, dtype), axis, size, fill_value),
       "shape": shape, "dtype": dtype, "axis": axis,
       "size": size, "fill_value": fill_value}
      for dtype in number_dtypes
      for size in [1, 5, 10]
      for fill_value in [None, -1.0, "slice"]
      for shape in nonempty_array_shapes
      for axis in [None] + list(range(len(shape)))))
  def testUniqueSize(self, shape, dtype, axis, size, fill_value):
    rng = jtu.rand_some_equal(self.rng())
    args_maker = lambda: [rng(shape, dtype)]
    kwds = dict(axis=axis, return_index=True, return_inverse=True, return_counts=True)

    if fill_value == "slice":
      if axis is None:
        fill_value = rng((), dtype)
      else:
        fill_value = rng(shape[:axis] + shape[axis + 1:], dtype)

    @partial(jtu.with_jax_dtype_defaults, use_defaults=(False, True, True, True))
    def np_fun(x, fill_value=fill_value):
      u, ind, inv, counts = np.unique(x, **kwds)
      axis = kwds['axis']
      if axis is None:
        x = x.ravel()
        axis = 0

      n_unique = u.shape[axis]
      if size <= u.shape[axis]:
        slc = (slice(None),) * axis + (slice(size),)
        u, ind, counts = u[slc], ind[:size], counts[:size]
      else:
        extra = (0, size - n_unique)
        pads = [(0, 0)] * u.ndim
        pads[axis] = extra
        u = np.pad(u, pads, constant_values=0)
        slices = [slice(None)] * u.ndim
        slices[axis] = slice(1)
        if fill_value is None:
          fill_value = u[tuple(slices)]
        elif np.ndim(fill_value):
          fill_value = lax.expand_dims(fill_value, (axis,))
        slices[axis] = slice(n_unique, None)
        u[tuple(slices)] = fill_value
        ind = np.pad(ind, extra, constant_values=ind[0])
        counts = np.pad(counts, extra, constant_values=0)
      return u, ind, inv, counts

    jnp_fun = lambda x: jnp.unique(x, size=size, fill_value=fill_value, **kwds)

    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)
    self._CompileAndCheck(jnp_fun, args_maker)

  @unittest.skipIf(numpy_version < (1, 21), "Numpy < 1.21 does not properly handle NaN values in unique.")
  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": f"_{dtype.__name__}", "dtype": dtype}
      for dtype in inexact_dtypes))
  def testUniqueNans(self, dtype):
    def args_maker():
      x = [-0.0, 0.0, 1.0, 1.0, np.nan, -np.nan]
      if np.issubdtype(dtype, np.complexfloating):
        x = [complex(i, j) for i, j in itertools.product(x, repeat=2)]
      return [np.array(x, dtype=dtype)]

    kwds = dict(return_index=True, return_inverse=True, return_counts=True)
    jnp_fun = partial(jnp.unique, **kwds)
    def np_fun(x):
      dtype = x.dtype
      # numpy unique fails for bfloat16 NaNs, so we cast to float64
      if x.dtype == jnp.bfloat16:
        x = x.astype('float64')
      u, *rest = np.unique(x, **kwds)
      return (u.astype(dtype), *rest)
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)


  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": f"_fixed_size={fixed_size}",
      "fixed_size": fixed_size}
      for fixed_size in [True, False]))
  def testNonScalarRepeats(self, fixed_size):
    '''
    Following numpy test suite from `test_repeat` at
    https://github.com/numpy/numpy/blob/main/numpy/core/tests/test_multiarray.py
    '''
    tol = 1e-5

    def test_single(m, args_maker, repeats, axis):
      lax_ans = jnp.repeat(m, repeats, axis)
      numpy_ans = np.repeat(m, repeats, axis)

      self.assertAllClose(lax_ans, numpy_ans, rtol=tol, atol=tol)
      if fixed_size:

        # Calculate expected size of the repeated axis.
        rep_length = np.repeat(np.zeros_like(m), repeats, axis).shape[axis or 0]
        jnp_fun = lambda arg, rep: jnp.repeat(
            arg, repeats=rep, axis=axis, total_repeat_length=rep_length)
      else:
        jnp_fun = lambda arg: jnp.repeat(arg, repeats = repeats, axis=axis)
      self._CompileAndCheck(jnp_fun, args_maker)

    m = jnp.array([1,2,3,4,5,6])
    if fixed_size:
      args_maker = lambda: [m, repeats]
    else:
      args_maker = lambda: [m]

    for repeats in [2, jnp.array([1,3,0,1,1,2]), jnp.array([1,3,2,1,1,2]), jnp.array([2])]:
      test_single(m, args_maker, repeats, axis=None)
      test_single(m, args_maker, repeats, axis=0)

    m_rect = m.reshape((2,3))
    if fixed_size:
      args_maker = lambda: [m_rect, repeats]
    else:
      args_maker = lambda: [m_rect]

    for repeats in [2, jnp.array([2,1]), jnp.array([2])]:
      test_single(m_rect, args_maker, repeats, axis=0)

    for repeats in [2, jnp.array([1,3,2]), jnp.array([2])]:
      test_single(m_rect, args_maker, repeats, axis=1)

  def testIssue2330(self):
    '''
    Make sure return value of jnp.concatenate is a jax.ndarray and is side-effect save
    '''
    def attempt_sideeffect(x):
      x = [x]
      x = jnp.concatenate(x)
      x -= 1.
      return x

    np_input = np.ones(1)
    jnp_input = jnp.ones(1)
    expected_np_input_after_call = np.ones(1)
    expected_jnp_input_after_call = jnp.ones(1)

    self.assertTrue(device_array.type_is_device_array(jnp.concatenate([np_input])))

    attempt_sideeffect(np_input)
    attempt_sideeffect(jnp_input)

    self.assertAllClose(np_input, expected_np_input_after_call)
    self.assertAllClose(jnp_input, expected_jnp_input_after_call)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "op={}_xshape=[{}]_yshape=[{}]_mode={}".format(
          op,
          jtu.format_shape_dtype_string(xshape, dtype),
          jtu.format_shape_dtype_string(yshape, dtype),
          mode),
       "xshape": xshape, "yshape": yshape, "dtype": dtype, "mode": mode,
       "jnp_op": getattr(jnp, op),
       "np_op": getattr(np, op)}
      for mode in ['full', 'same', 'valid']
      for op in ['convolve', 'correlate']
      for dtype in number_dtypes
      for xshape in one_dim_array_shapes
      for yshape in one_dim_array_shapes))
  def testConvolutions(self, xshape, yshape, dtype, mode, jnp_op, np_op):
    rng = jtu.rand_default(self.rng())
    args_maker = lambda: [rng(xshape, dtype), rng(yshape, dtype)]
    precision = lax.Precision.HIGHEST if jtu.device_under_test() == "tpu" else None
    np_fun = partial(np_op, mode=mode)
    jnp_fun = partial(jnp_op, mode=mode, precision=precision)
    tol = {np.float16: 2e-1, np.float32: 1e-2, np.float64: 1e-14,
           np.complex128: 1e-14}
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker, check_dtypes=False,
                            tol=tol)
    self._CompileAndCheck(jnp_fun, args_maker)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "op={}_shape=[{}]_axis={}_out_dtype={}".format(
          op, jtu.format_shape_dtype_string(shape, dtype), axis,
          out_dtype.__name__),
       "axis": axis, "shape": shape, "dtype": dtype, "out_dtype": out_dtype,
       "jnp_op": getattr(jnp, op), "np_op": getattr(np, op)}
      for op in ["cumsum", "cumprod"]
      for dtype in all_dtypes
      for out_dtype in default_dtypes
      for shape in all_shapes
      for axis in [None] + list(range(-len(shape), len(shape)))))
  def testCumSumProd(self, axis, shape, dtype, out_dtype, np_op, jnp_op):
    rng = jtu.rand_default(self.rng())
    np_fun = lambda arg: np_op(arg, axis=axis, dtype=out_dtype)
    np_fun = jtu.ignore_warning(category=np.ComplexWarning)(np_fun)
    jnp_fun = lambda arg: jnp_op(arg, axis=axis, dtype=out_dtype)
    jnp_fun = jtu.ignore_warning(category=jnp.ComplexWarning)(jnp_fun)

    args_maker = lambda: [rng(shape, dtype)]

    tol_thresholds = {dtypes.bfloat16: 4e-2}
    tol = max(jtu.tolerance(dtype, tol_thresholds),
              jtu.tolerance(out_dtype, tol_thresholds))
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker,
                            tol=tol)
    self._CompileAndCheck(jnp_fun, args_maker)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "op={}_shape=[{}]_axis={}_out_dtype={}".format(
          op, jtu.format_shape_dtype_string(shape, dtype), axis,
          out_dtype.__name__),
       "axis": axis, "shape": shape, "dtype": dtype, "out_dtype": out_dtype,
       "jnp_op": getattr(jnp, op), "np_op": getattr(np, op)}
      for op in ["nancumsum", "nancumprod"]
      for dtype in all_dtypes
      for out_dtype in default_dtypes
      for shape in all_shapes
      for axis in [None] + list(range(-len(shape), len(shape)))))
  def testNanCumSumProd(self, axis, shape, dtype, out_dtype, np_op, jnp_op):
    rng = jtu.rand_some_nan(self.rng())
    np_fun = partial(np_op, axis=axis, dtype=out_dtype)
    np_fun = jtu.ignore_warning(category=np.ComplexWarning)(np_fun)
    jnp_fun = partial(jnp_op, axis=axis, dtype=out_dtype)
    jnp_fun = jtu.ignore_warning(category=jnp.ComplexWarning)(jnp_fun)

    args_maker = lambda: [rng(shape, dtype)]

    tol_thresholds = {dtypes.bfloat16: 4e-2}
    tol = max(jtu.tolerance(dtype, tol_thresholds),
              jtu.tolerance(out_dtype, tol_thresholds))
    if dtype != jnp.bfloat16:
      # numpy functions do not properly handle bfloat16
      self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker, check_dtypes=True,
                              tol=tol)
    self._CompileAndCheck(jnp_fun, args_maker, check_dtypes=True)


  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_yshape={}_xshape={}_dx={}_axis={}".format(
        jtu.format_shape_dtype_string(yshape, dtype),
        jtu.format_shape_dtype_string(xshape, dtype) if xshape is not None else None,
        dx, axis),
        "yshape": yshape, "xshape": xshape, "dtype": dtype, "dx": dx, "axis": axis}
        for dtype in default_dtypes
        for yshape, xshape, dx, axis in [
          ((10,), None, 1.0, -1),
          ((3, 10), None, 2.0, -1),
          ((3, 10), None, 3.0, -0),
          ((10, 3), (10,), 1.0, -2),
          ((3, 10), (10,), 1.0, -1),
          ((3, 10), (3, 10), 1.0, -1),
          ((2, 3, 10), (3, 10), 1.0, -2),
        ]))
  @jtu.skip_on_devices("tpu")  # TODO(jakevdp): fix and reenable this test.
  @jax.numpy_rank_promotion('allow')  # This test explicitly exercises implicit rank promotion.
  def testTrapz(self, yshape, xshape, dtype, dx, axis):
    rng = jtu.rand_default(self.rng())
    args_maker = lambda: [rng(yshape, dtype), rng(xshape, dtype) if xshape is not None else None]
    np_fun = partial(np.trapz, dx=dx, axis=axis)
    jnp_fun = partial(jnp.trapz, dx=dx, axis=axis)
    tol = jtu.tolerance(dtype, {np.float64: 1e-12,
                                dtypes.bfloat16: 4e-2})
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker, tol=tol,
                            check_dtypes=False)
    self._CompileAndCheck(jnp_fun, args_maker, atol=tol, rtol=tol,
                          check_dtypes=False)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_dtype={}_m={}_n={}_k={}".format(
          np.dtype(dtype).name, m, n, k),
       "m": m, "n": n, "k": k, "dtype": dtype}
      for dtype in default_dtypes
      for n in [0, 4]
      for m in [None, 0, 1, 3, 4]
      for k in list(range(-4, 4))))
  def testTri(self, m, n, k, dtype):
    np_fun = lambda: np.tri(n, M=m, k=k, dtype=dtype)
    jnp_fun = lambda: jnp.tri(n, M=m, k=k, dtype=dtype)
    args_maker = lambda: []
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)
    self._CompileAndCheck(jnp_fun, args_maker)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_op={}_shape={}_k={}".format(
          op, jtu.format_shape_dtype_string(shape, dtype), k),
       "dtype": dtype, "shape": shape, "op": op, "k": k}
      for dtype in default_dtypes
      for shape in [shape for shape in all_shapes if len(shape) >= 2]
      for op in ["tril", "triu"]
      for k in list(range(-3, 3))))
  def testTriLU(self, dtype, shape, op, k):
    rng = jtu.rand_default(self.rng())
    np_fun = lambda arg: getattr(np, op)(arg, k=k)
    jnp_fun = lambda arg: getattr(jnp, op)(arg, k=k)
    args_maker = lambda: [rng(shape, dtype)]
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)
    self._CompileAndCheck(jnp_fun, args_maker)

  @parameterized.named_parameters(jtu.cases_from_list(
    {"testcase_name": f"n={n}_k={k}_m={m}",
      "n": n, "k": k, "m": m}
    for n in range(1, 5)
    for k in [-1, 0, 1]
    for m in range(1, 5)))
  def testTrilIndices(self, n, k, m):
    np_fun = lambda n, k, m: np.tril_indices(n, k=k, m=m)
    jnp_fun = lambda n, k, m: jnp.tril_indices(n, k=k, m=m)
    args_maker = lambda: [n, k, m]
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)

  @parameterized.named_parameters(jtu.cases_from_list(
    {"testcase_name": f"n={n}_k={k}_m={m}",
      "n": n, "k": k, "m": m}
    for n in range(1, 5)
    for k in [-1, 0, 1]
    for m in range(1, 5)))
  def testTriuIndices(self, n, k, m):
    np_fun = lambda n, k, m: np.triu_indices(n, k=k, m=m)
    jnp_fun = lambda n, k, m: jnp.triu_indices(n, k=k, m=m)
    args_maker = lambda: [n, k, m]
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)

  @parameterized.named_parameters(jtu.cases_from_list(
    {"testcase_name": "_shape={}_k={}".format(
      jtu.format_shape_dtype_string(shape, dtype), k),
      "dtype": dtype, "shape": shape, "k": k}
    for dtype in default_dtypes
    for shape in [(1,1), (1,2), (2,2), (2,3), (3,2), (3,3), (4,4)]
    for k in [-1, 0, 1]))
  def testTriuIndicesFrom(self, shape, dtype, k):
    rng = jtu.rand_default(self.rng())
    np_fun = lambda arr, k: np.triu_indices_from(arr, k=k)
    jnp_fun = lambda arr, k: jnp.triu_indices_from(arr, k=k)
    args_maker = lambda: [rng(shape, dtype), k]
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)

  @parameterized.named_parameters(jtu.cases_from_list(
    {"testcase_name": "_shape={}_k={}".format(
      jtu.format_shape_dtype_string(shape, dtype), k),
      "dtype": dtype, "shape": shape, "k": k}
    for dtype in default_dtypes
    for shape in [(1,1), (1,2), (2,2), (2,3), (3,2), (3,3), (4,4)]
    for k in [-1, 0, 1]))
  def testTrilIndicesFrom(self, shape, dtype, k):
    rng = jtu.rand_default(self.rng())
    np_fun = lambda arr, k: np.tril_indices_from(arr, k=k)
    jnp_fun = lambda arr, k: jnp.tril_indices_from(arr, k=k)
    args_maker = lambda: [rng(shape, dtype), k]
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": f"_ndim={ndim}_n={n}",
       "ndim": ndim, "n": n}
      for ndim in [0, 1, 4]
      for n in [0, 1, 7]))
  def testDiagIndices(self, ndim, n):
    np.testing.assert_equal(jtu.with_jax_dtype_defaults(np.diag_indices)(n, ndim),
                            jnp.diag_indices(n, ndim))

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "arr_shape={}".format(
        jtu.format_shape_dtype_string(shape, dtype)
      ),
       "dtype": dtype, "shape": shape}
      for dtype in default_dtypes
      for shape in [(1,1), (2,2), (3,3), (4,4), (5,5)]))
  def testDiagIndicesFrom(self, dtype, shape):
    rng = jtu.rand_default(self.rng())
    np_fun = jtu.with_jax_dtype_defaults(np.diag_indices_from)
    jnp_fun = jnp.diag_indices_from
    args_maker = lambda : [rng(shape, dtype)]
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)
    self._CompileAndCheck(jnp_fun, args_maker)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_shape={}_k={}".format(
          jtu.format_shape_dtype_string(shape, dtype), k),
       "dtype": dtype, "shape": shape, "k": k}
      for dtype in default_dtypes
      for shape in [shape for shape in all_shapes if len(shape) in (1, 2)]
      for k in list(range(-4, 4))))
  def testDiag(self, shape, dtype, k):
    rng = jtu.rand_default(self.rng())
    np_fun = lambda arg: np.diag(arg, k)
    jnp_fun = lambda arg: jnp.diag(arg, k)
    args_maker = lambda: [rng(shape, dtype)]
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)
    self._CompileAndCheck(jnp_fun, args_maker)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_shape={}_k={}".format(
          jtu.format_shape_dtype_string(shape, dtype), k),
       "dtype": dtype, "shape": shape, "k": k}
      for dtype in default_dtypes
      for shape in all_shapes
      for k in range(-4, 4)))
  def testDiagFlat(self, shape, dtype, k):
    rng = jtu.rand_default(self.rng())
    # numpy has inconsistencies for scalar values
    # https://github.com/numpy/numpy/issues/16477
    # jax differs in that it treats scalars values as length-1 arrays
    np_fun = lambda arg: np.diagflat(np.atleast_1d(arg), k)
    jnp_fun = lambda arg: jnp.diagflat(arg, k)
    args_maker = lambda: [rng(shape, dtype)]
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker, check_dtypes=True)
    self._CompileAndCheck(jnp_fun, args_maker, check_dtypes=True)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_a1_shape={}_a2_shape2={}".format(
          jtu.format_shape_dtype_string(a1_shape, dtype),
          jtu.format_shape_dtype_string(a2_shape, dtype)),
       "dtype": dtype, "a1_shape": a1_shape, "a2_shape": a2_shape}
      for dtype in default_dtypes
      for a1_shape in one_dim_array_shapes
      for a2_shape in one_dim_array_shapes))
  def testPolyMul(self, a1_shape, a2_shape, dtype):
    rng = jtu.rand_default(self.rng())
    np_fun = lambda arg1, arg2: np.polymul(arg1, arg2)
    jnp_fun_np = lambda arg1, arg2: jnp.polymul(arg1, arg2, trim_leading_zeros=True)
    jnp_fun_co = lambda arg1, arg2: jnp.polymul(arg1, arg2)
    args_maker = lambda: [rng(a1_shape, dtype), rng(a2_shape, dtype)]
    tol = {np.float16: 2e-1, np.float32: 5e-2, np.float64: 1e-13}
    self._CheckAgainstNumpy(np_fun, jnp_fun_np, args_maker, check_dtypes=False, tol=tol)
    self._CompileAndCheck(jnp_fun_co, args_maker, check_dtypes=False)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "a_shape={} , b_shape={}".format(
          jtu.format_shape_dtype_string(a_shape, dtype),
          jtu.format_shape_dtype_string(b_shape, dtype)),
       "dtype": dtype, "a_shape": a_shape, "b_shape" : b_shape}
      for dtype in default_dtypes
      for a_shape in one_dim_array_shapes
      for b_shape in one_dim_array_shapes))
  def testPolyDiv(self, a_shape, b_shape, dtype):
    rng = jtu.rand_default(self.rng())

    @jtu.ignore_warning(category=RuntimeWarning, message="divide by zero.*")
    @jtu.ignore_warning(category=RuntimeWarning, message="invalid value.*")
    def np_fun(arg1, arg2):
      q, r = np.polydiv(arg1, arg2)
      while r.size < max(arg1.size, arg2.size): # Pad residual to same size
        r = np.pad(r, (1, 0), 'constant')
      return q, r

    def jnp_fun(arg1, arg2):
      q, r = jnp.polydiv(arg1, arg2, trim_leading_zeros=True)
      while r.size < max(arg1.size, arg2.size): # Pad residual to same size
        r = jnp.pad(r, (1, 0), 'constant')
      return q, r

    args_maker = lambda: [rng(a_shape, dtype), rng(b_shape, dtype)]
    tol = {np.float16: 2e-1, np.float32: 5e-2, np.float64: 1e-13}

    jnp_compile = jnp.polydiv # Without trim_leading_zeros (trim_zeros make it unable to be compiled by XLA)
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker, check_dtypes=False, tol=tol)
    self._CompileAndCheck(jnp_compile, args_maker, check_dtypes=True, atol=tol, rtol=tol)


  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_shape={}_offset={}_axis1={}_axis2={}".format(
          jtu.format_shape_dtype_string(shape, dtype), offset, axis1, axis2),
       "dtype": dtype, "shape": shape, "offset": offset, "axis1": axis1,
       "axis2": axis2}
      for dtype in default_dtypes
      for shape in [shape for shape in all_shapes if len(shape) >= 2]
      for axis1 in range(-len(shape), len(shape))
      for axis2 in [a for a in range(-len(shape), len(shape))
                    if a % len(shape) != axis1 % len(shape)]
      for offset in list(range(-4, 4))))
  def testDiagonal(self, shape, dtype, offset, axis1, axis2):
    rng = jtu.rand_default(self.rng())
    np_fun = lambda arg: np.diagonal(arg, offset, axis1, axis2)
    jnp_fun = lambda arg: jnp.diagonal(arg, offset, axis1, axis2)
    args_maker = lambda: [rng(shape, dtype)]
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)
    self._CompileAndCheck(jnp_fun, args_maker)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": f"_shape={np.dtype(dtype).name}_n={n}",
       "dtype": dtype, "n": n}
      for dtype in default_dtypes
      for n in list(range(4))))
  def testIdentity(self, n, dtype):
    np_fun = lambda: np.identity(n, dtype)
    jnp_fun = lambda: jnp.identity(n, dtype)
    args_maker = lambda: []
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)
    self._CompileAndCheck(jnp_fun, args_maker)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_{}_period={}_left={}_right={}".format(
       jtu.format_shape_dtype_string(shape, dtype), period, left, right),
       "shape": shape, "dtype": dtype,
       "period": period, "left": left, "right": right}
      for shape in nonempty_shapes
      for period in [None, 0.59]
      for left in [None, 0]
      for right in [None, 1]
      for dtype in default_dtypes
      # following types lack precision for meaningful tests
      if dtype not in [np.int8, np.int16, np.float16, jnp.bfloat16]
  ))
  def testInterp(self, shape, dtype, period, left, right):
    rng = jtu.rand_default(self.rng(), scale=10)
    kwds = dict(period=period, left=left, right=right)
    np_fun = partial(np.interp, **kwds)
    jnp_fun = partial(jnp.interp, **kwds)
    args_maker = lambda: [rng(shape, dtype), np.sort(rng((20,), dtype)), np.linspace(0, 1, 20)]

    # skip numpy comparison for integer types with period specified, because numpy
    # uses an unstable sort and so results differ for duplicate values.
    if not (period and np.issubdtype(dtype, np.integer)):
      self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker, tol={np.float32: 2E-4})
    self._CompileAndCheck(jnp_fun, args_maker)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_x1={}_x2={}_x1_rng={}".format(
          jtu.format_shape_dtype_string(x1_shape, x1_dtype),
          jtu.format_shape_dtype_string(x2_shape, np.int32),
          x1_rng_factory_id),
       "x1_shape": x1_shape, "x1_dtype": x1_dtype,
       "x2_shape": x2_shape, "x1_rng_factory": x1_rng_factory,
       "x2_rng_factory": x2_rng_factory}
      for x1_rng_factory_id, x1_rng_factory in
        enumerate([jtu.rand_some_inf_and_nan, jtu.rand_some_zero])
      for x2_rng_factory in [partial(jtu.rand_int, low=-1075, high=1024)]
      for x1_shape, x2_shape in filter(_shapes_are_broadcast_compatible,
                                       itertools.combinations_with_replacement(array_shapes, 2))
      for x1_dtype in default_dtypes))
  @jax.numpy_rank_promotion('allow')  # This test explicitly exercises implicit rank promotion.
  def testLdexp(self, x1_shape, x1_dtype, x2_shape, x1_rng_factory, x2_rng_factory):
    # integer types are converted to float64 in numpy's implementation
    if (x1_dtype not in [jnp.bfloat16, np.float16, np.float32]
        and not config.x64_enabled):
      self.skipTest("Only run float64 testcase when float64 is enabled.")
    x1_rng = x1_rng_factory(self.rng())
    x2_rng = x2_rng_factory(self.rng())
    np_fun = lambda x1, x2: np.ldexp(x1, x2)
    np_fun = jtu.ignore_warning(category=RuntimeWarning,
                                 message="overflow.*")(np_fun)
    jnp_fun = lambda x1, x2: jnp.ldexp(x1, x2)
    args_maker = lambda: [x1_rng(x1_shape, x1_dtype),
                          x2_rng(x2_shape, np.int32)]
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)
    self._CompileAndCheck(jnp_fun, args_maker)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_x={}_rng_factory={}".format(
          jtu.format_shape_dtype_string(shape, dtype), rng_factory_id),
       "shape": shape, "dtype": dtype, "rng_factory": rng_factory}
      for rng_factory_id, rng_factory in enumerate([
          jtu.rand_some_inf_and_nan,
          jtu.rand_some_zero,
          partial(jtu.rand_not_small, offset=1e8),
      ])
      for shape in all_shapes
      for dtype in default_dtypes))
  def testFrexp(self, shape, dtype, rng_factory):
    # integer types are converted to float64 in numpy's implementation
    if (dtype not in [jnp.bfloat16, np.float16, np.float32]
        and not config.x64_enabled):
      self.skipTest("Only run float64 testcase when float64 is enabled.")
    rng = rng_factory(self.rng())
    np_fun = lambda x: np.frexp(x)
    jnp_fun = lambda x: jnp.frexp(x)
    args_maker = lambda: [rng(shape, dtype)]
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker,
                            check_dtypes=np.issubdtype(dtype, np.inexact))
    self._CompileAndCheck(jnp_fun, args_maker)


  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_shape={}_dtype_{}_offset={}_axis1={}_axis2={}".format(
          jtu.format_shape_dtype_string(shape, dtype),
          out_dtype, offset, axis1, axis2),
       "dtype": dtype, "out_dtype": out_dtype, "shape": shape, "offset": offset,
       "axis1": axis1, "axis2": axis2}
      for dtype in default_dtypes
      for out_dtype in [None] + number_dtypes
      for shape in [shape for shape in all_shapes if len(shape) >= 2]
      for axis1 in range(-len(shape), len(shape))
      for axis2 in range(-len(shape), len(shape))
      if (axis1 % len(shape)) != (axis2 % len(shape))
      for offset in list(range(-4, 4))))
  def testTrace(self, shape, dtype, out_dtype, offset, axis1, axis2):
    rng = jtu.rand_default(self.rng())
    def np_fun(arg):
      if out_dtype == jnp.bfloat16:
        return np.trace(arg, offset, axis1, axis2, np.float32).astype(jnp.bfloat16)
      else:
        return np.trace(arg, offset, axis1, axis2, out_dtype)
    jnp_fun = lambda arg: jnp.trace(arg, offset, axis1, axis2, out_dtype)
    args_maker = lambda: [rng(shape, dtype)]
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)
    self._CompileAndCheck(jnp_fun, args_maker)

  @parameterized.named_parameters(jtu.cases_from_list(
    {"testcase_name": "_a={}_v={}_side={}".format(
      jtu.format_shape_dtype_string(ashape, dtype),
      jtu.format_shape_dtype_string(vshape, dtype),
      side), "ashape": ashape, "vshape": vshape, "side": side,
     "dtype": dtype}
    for ashape in [(15,), (16,), (17,)]
    for vshape in [(), (5,), (5, 5)]
    for side in ['left', 'right']
    for dtype in number_dtypes
  ))
  def testSearchsorted(self, ashape, vshape, side, dtype):
    rng = jtu.rand_default(self.rng())
    args_maker = lambda: [np.sort(rng(ashape, dtype)), rng(vshape, dtype)]
    def np_fun(a, v):
      return np.searchsorted(a, v, side=side).astype('int32')
    jnp_fun = lambda a, v: jnp.searchsorted(a, v, side=side)
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)
    self._CompileAndCheck(jnp_fun, args_maker)

  def testSearchsortedDtype(self):
    # Test that for large arrays, int64 indices are used. We test this
    # via abstract evaluation to avoid allocating a large array in tests.
    a_int32 = jax.core.ShapedArray((np.iinfo(np.int32).max,), np.float32)
    a_int64 = jax.core.ShapedArray((np.iinfo(np.int32).max + 1,), np.float32)
    v = jax.core.ShapedArray((), np.float32)

    out_int32 = jax.eval_shape(jnp.searchsorted, a_int32, v)
    self.assertEqual(out_int32.dtype, np.int32)

    if config.x64_enabled:
      out_int64 = jax.eval_shape(jnp.searchsorted, a_int64, v)
      self.assertEqual(out_int64.dtype, np.int64)
    else:
      with self.assertWarnsRegex(UserWarning, "Explicitly requested dtype int64"):
        out_int64 = jax.eval_shape(jnp.searchsorted, a_int64, v)

  @parameterized.named_parameters(jtu.cases_from_list(
    {"testcase_name": f"_dtype={dtype.__name__}_side={side}", "dtype": dtype, "side": side}
    for dtype in inexact_dtypes
    for side in ['left', 'right']))
  def testSearchsortedNans(self, dtype, side):
    if np.issubdtype(dtype, np.complexfloating):
      raise SkipTest("Known failure for complex inputs; see #9107")
    x = np.array([-np.inf, -1.0, 0.0, -0.0, 1.0, np.inf, np.nan, -np.nan], dtype=dtype)
    # The sign bit should not matter for 0.0 or NaN, so argsorting the above should be
    # equivalent to argsorting the following:
    x_equiv = np.array([0, 1, 2, 2, 3, 4, 5, 5])

    if jnp.issubdtype(dtype, jnp.complexfloating):
      x = np.array([complex(r, c) for r, c in itertools.product(x, repeat=2)])
      x_equiv = np.array([complex(r, c) for r, c in itertools.product(x_equiv, repeat=2)])

    fun = partial(jnp.searchsorted, side=side)
    self.assertArraysEqual(fun(x, x), fun(x_equiv, x_equiv))
    self.assertArraysEqual(jax.jit(fun)(x, x), fun(x_equiv, x_equiv))

  @parameterized.named_parameters(jtu.cases_from_list(
    {"testcase_name": "_x={}_bins={}_right={}_reverse={}".format(
      jtu.format_shape_dtype_string(xshape, dtype),
      jtu.format_shape_dtype_string(binshape, dtype),
      right, reverse), "xshape": xshape, "binshape": binshape,
      "right": right, "reverse": reverse, "dtype": dtype}
    for xshape in [(20,), (5, 4)]
    for binshape in [(1,), (5,)]
    for right in [True, False]
    for reverse in [True, False]
    for dtype in default_dtypes
  ))
  def testDigitize(self, xshape, binshape, right, reverse, dtype):
    order = jnp.index_exp[::-1] if reverse else jnp.index_exp[:]
    rng = jtu.rand_default(self.rng())
    args_maker = lambda: [rng(xshape, dtype), jnp.sort(rng(binshape, dtype))[order]]
    np_fun = lambda x, bins: np.digitize(x, bins, right=right).astype('int32')
    jnp_fun = lambda x, bins: jnp.digitize(x, bins, right=right)
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)
    self._CompileAndCheck(jnp_fun, args_maker)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_{}_array={}".format(
          jtu.format_test_name_suffix("", [shape] * len(dtypes), dtypes), array_input),
       "shape": shape, "dtypes": dtypes, "array_input": array_input}
      for dtypes in [
        [np.float32],
        [np.float32, np.float32],
        [np.float32, np.int32, np.float32],
        [np.float32, np.int64, np.float32],
        [np.float32, np.int32, np.float64],
      ]
      for shape in [(), (2,), (3, 4), (1, 5)]
      for array_input in [True, False]))
  def testColumnStack(self, shape, dtypes, array_input):
    rng = jtu.rand_default(self.rng())
    if array_input:
      args_maker = lambda: [np.array([rng(shape, dtype) for dtype in dtypes])]
    else:
      args_maker = lambda: [[rng(shape, dtype) for dtype in dtypes]]
    np_fun = _promote_like_jnp(np.column_stack)
    jnp_fun = jnp.column_stack
    self._CheckAgainstNumpy(jnp_fun, np_fun, args_maker)
    self._CompileAndCheck(jnp_fun, args_maker)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_{}_axis={}_array={}".format(
          jtu.format_test_name_suffix("", [shape] * len(dtypes), dtypes), axis, array_input),
       "shape": shape, "axis": axis, "dtypes": dtypes, "array_input": array_input}
      for dtypes in [
        [np.float32],
        [np.float32, np.float32],
        [np.float32, np.int32, np.float32],
        [np.float32, np.int64, np.float32],
        [np.float32, np.int32, np.float64],
      ]
      for shape in [(), (2,), (3, 4), (1, 100)]
      for axis in range(-len(shape), len(shape) + 1)
      for array_input in [True, False]))
  def testStack(self, shape, axis, dtypes, array_input):
    rng = jtu.rand_default(self.rng())
    if array_input:
      args_maker = lambda: [np.array([rng(shape, dtype) for dtype in dtypes])]
    else:
      args_maker = lambda: [[rng(shape, dtype) for dtype in dtypes]]
    np_fun = _promote_like_jnp(partial(np.stack, axis=axis))
    jnp_fun = partial(jnp.stack, axis=axis)
    self._CheckAgainstNumpy(jnp_fun, np_fun, args_maker)
    self._CompileAndCheck(jnp_fun, args_maker)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_op={}_{}_array={}".format(
          op, jtu.format_test_name_suffix("", [shape] * len(dtypes), dtypes), array_input),
       "shape": shape, "op": op, "dtypes": dtypes, "array_input": array_input}
      for op in ["hstack", "vstack", "dstack"]
      for dtypes in [
        [np.float32],
        [np.float32, np.float32],
        [np.float32, np.int32, np.float32],
        [np.float32, np.int64, np.float32],
        [np.float32, np.int32, np.float64],
      ]
      for shape in [(), (2,), (3, 4), (1, 100), (2, 3, 4)]
      for array_input in [True, False]))
  def testHVDStack(self, shape, op, dtypes, array_input):
    rng = jtu.rand_default(self.rng())
    if array_input:
      args_maker = lambda: [np.array([rng(shape, dtype) for dtype in dtypes])]
    else:
      args_maker = lambda: [[rng(shape, dtype) for dtype in dtypes]]
    np_fun = _promote_like_jnp(getattr(np, op))
    jnp_fun = getattr(jnp, op)
    self._CheckAgainstNumpy(jnp_fun, np_fun, args_maker)
    self._CompileAndCheck(jnp_fun, args_maker)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_inshape={}_outdtype={}_fillshape={}".format(
          jtu.format_shape_dtype_string(shape, fill_value_dtype),
          np.dtype(out_dtype).name if out_dtype else "None",
          fill_value_shape),
       "fill_value_dtype": fill_value_dtype, "fill_value_shape": fill_value_shape,
       "shape": shape, "out_dtype": out_dtype}
      for shape in array_shapes + [3, np.array(7, dtype=np.int32)]
      for fill_value_dtype in default_dtypes
      for fill_value_shape in _compatible_shapes(shape)
      for out_dtype in [None] + default_dtypes))
  def testFull(self, shape, fill_value_dtype, fill_value_shape, out_dtype):
    rng = jtu.rand_default(self.rng())
    np_fun = lambda fill_value: np.full(shape, fill_value, dtype=out_dtype)
    jnp_fun = lambda fill_value: jnp.full(shape, fill_value, dtype=out_dtype)
    args_maker = lambda: [rng(fill_value_shape, fill_value_dtype)]
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)
    self._CompileAndCheck(jnp_fun, args_maker)

  @parameterized.named_parameters(jtu.named_cases_from_sampler(lambda s: ({
      "testcase_name": "_shape={}_n={}_axis={}_prepend={}_append={}".format(
           jtu.format_shape_dtype_string(shape, dtype),
           n, axis, prepend, append),
        "shape": shape, "dtype": dtype, "n": n, "axis": axis,
        "prepend": prepend, "append": append
    } for shape, dtype in s(_shape_and_dtypes(nonempty_nonscalar_array_shapes, default_dtypes))
      for n in s([0, 1, 2])
      for axis in s(list(range(-len(shape), max(1, len(shape)))))
      for prepend in s([None, 1, np.zeros(shape, dtype=dtype)])
      for append in s([None, 1, np.zeros(shape, dtype=dtype)])
      )))
  def testDiff(self, shape, dtype, n, axis, prepend, append):
    rng = jtu.rand_default(self.rng())
    args_maker = lambda: [rng(shape, dtype)]

    def np_fun(x, n=n, axis=axis, prepend=prepend, append=append):
      if prepend is None:
        prepend = np._NoValue
      elif not np.isscalar(prepend) and prepend.dtype == jnp.bfloat16:
        prepend = prepend.astype(np.float32)

      if append is None:
        append = np._NoValue
      elif not np.isscalar(append) and append.dtype == jnp.bfloat16:
        append = append.astype(np.float32)

      if x.dtype == jnp.bfloat16:
        return np.diff(x.astype(np.float32), n=n, axis=axis, prepend=prepend, append=append).astype(jnp.bfloat16)
      else:
        return np.diff(x, n=n, axis=axis, prepend=prepend, append=append)

    jnp_fun = lambda x: jnp.diff(x, n=n, axis=axis, prepend=prepend, append=append)
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker, check_dtypes=False)
    self._CompileAndCheck(jnp_fun, args_maker)

  @parameterized.named_parameters(
    jtu.cases_from_list(
      {"testcase_name": ("_op={}_shape={}_dtype={}").format(op, shape, dtype),
       "np_op": getattr(np, op), "jnp_op": getattr(jnp, op),
       "shape": shape, "dtype": dtype}
      for op in ["zeros", "ones"]
      for shape in [2, (), (2,), (3, 0), np.array((4, 5, 6), dtype=np.int32),
                    np.array(4, dtype=np.int32)]
      for dtype in all_dtypes))
  def testZerosOnes(self, np_op, jnp_op, shape, dtype):
    args_maker = lambda: []
    np_op = partial(np_op, shape, dtype)
    jnp_op = partial(jnp_op, shape, dtype)
    self._CheckAgainstNumpy(np_op, jnp_op, args_maker)
    self._CompileAndCheck(jnp_op, args_maker)

  def testOnesWithInvalidShape(self):
    with self.assertRaises(TypeError):
      jnp.ones((-1, 1))

  @parameterized.named_parameters(jtu.named_cases_from_sampler(lambda s: ({
       "testcase_name": "_inshape={}_filldtype={}_fillshape={}_outdtype={}_outshape={}".format(
          jtu.format_shape_dtype_string(shape, in_dtype),
          np.dtype(fill_value_dtype).name, fill_value_shape,
          np.dtype(out_dtype).name, out_shape),
       "shape": shape, "in_dtype": in_dtype,
       "fill_value_dtype": fill_value_dtype, "fill_value_shape": fill_value_shape,
       "out_dtype": out_dtype, "out_shape": out_shape
    } for shape in s(array_shapes)
      for out_shape in s([None] + array_shapes)
      for in_dtype in s(default_dtypes)
      for fill_value_dtype in s(default_dtypes)
      for fill_value_shape in s(_compatible_shapes(shape if out_shape is None else out_shape))
      for out_dtype in s(default_dtypes))))
  def testFullLike(self, shape, in_dtype, fill_value_dtype, fill_value_shape, out_dtype, out_shape):
    if numpy_version < (1, 19) and out_shape == ():
      raise SkipTest("Numpy < 1.19 treats out_shape=() like out_shape=None")
    rng = jtu.rand_default(self.rng())
    np_fun = lambda x, fill_value: np.full_like(
      x, fill_value, dtype=out_dtype, shape=out_shape)
    jnp_fun = lambda x, fill_value: jnp.full_like(
      x, fill_value, dtype=out_dtype, shape=out_shape)
    args_maker = lambda: [rng(shape, in_dtype), rng(fill_value_shape, fill_value_dtype)]
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)
    self._CompileAndCheck(jnp_fun, args_maker)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_func={}_inshape={}_outshape={}_outdtype={}".format(
          func, jtu.format_shape_dtype_string(shape, in_dtype),
          out_shape, out_dtype),
       "func": func, "shape": shape, "in_dtype": in_dtype,
       "out_shape": out_shape, "out_dtype": out_dtype}
      for shape in array_shapes
      for out_shape in [None] + array_shapes
      for in_dtype in default_dtypes
      for func in ["ones_like", "zeros_like"]
      for out_dtype in default_dtypes))
  def testZerosOnesLike(self, func, shape, in_dtype, out_shape, out_dtype):
    if numpy_version < (1, 19) and out_shape == ():
      raise SkipTest("Numpy < 1.19 treats out_shape=() like out_shape=None")
    rng = jtu.rand_default(self.rng())
    np_fun = lambda x: getattr(np, func)(x, dtype=out_dtype, shape=out_shape)
    jnp_fun = lambda x: getattr(jnp, func)(x, dtype=out_dtype, shape=out_shape)
    args_maker = lambda: [rng(shape, in_dtype)]
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)
    self._CompileAndCheck(jnp_fun, args_maker)


  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_func={}_inshape={}_weak_type={}_outshape={}_outdtype={}".format(
          func, jtu.format_shape_dtype_string(shape, in_dtype),
          weak_type, out_shape, out_dtype),
       "func": func, "args": args,
       "shape": shape, "in_dtype": in_dtype, "weak_type": weak_type,
       "out_shape": out_shape, "out_dtype": out_dtype}
      for shape in array_shapes
      for in_dtype in [np.int32, np.float32, np.complex64]
      for weak_type in [True, False]
      for out_shape in [None, (), (10,)]
      for func, args in [("full_like", (-100,)), ("ones_like", ()), ("zeros_like", ())]
      for out_dtype in [None, float]))
  def testZerosOnesFullLikeWeakType(self, func, args, shape, in_dtype, weak_type, out_shape, out_dtype):
    if numpy_version < (1, 19) and out_shape == ():
      raise SkipTest("Numpy < 1.19 treats out_shape=() like out_shape=None")
    rng = jtu.rand_default(self.rng())
    x = lax_internal._convert_element_type(rng(shape, in_dtype),
                                           weak_type=weak_type)
    fun = lambda x: getattr(jnp, func)(x, *args, dtype=out_dtype, shape=out_shape)
    expected_weak_type = weak_type and (out_dtype is None)
    self.assertEqual(dtypes.is_weakly_typed(fun(x)), expected_weak_type)
    self.assertEqual(dtypes.is_weakly_typed(jax.jit(fun)(x)), expected_weak_type)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_funcname={}_input_type={}_val={}_dtype={}".format(
          funcname, input_type, val, dtype),
       "funcname": funcname, "input_type": input_type, "val": val, "dtype": dtype}
      for funcname in ["array", "asarray"]
      for dtype in [int, float, None]
      for val in [0, 1]
      for input_type in [int, float, np.int32, np.float32]))
  def testArrayWeakType(self, funcname, input_type, val, dtype):
    func = lambda x: getattr(jnp, funcname)(x, dtype=dtype)
    fjit = jax.jit(func)
    val = input_type(val)
    expected_weak_type = dtype is None and input_type in set(dtypes._weak_types)
    self.assertEqual(dtypes.is_weakly_typed(func(val)), expected_weak_type)
    self.assertEqual(dtypes.is_weakly_typed(fjit(val)), expected_weak_type)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_{}_weak_type={}_slc={}".format(
        jtu.format_shape_dtype_string(shape, dtype), weak_type, slc),
       "shape": shape, "dtype": dtype, "weak_type": weak_type, "slc": slc}
      for shape in nonempty_nonscalar_array_shapes
      for dtype in [int, float, complex]
      for weak_type in [True, False]
      for slc in [slice(None), slice(0), slice(3), 0, ...]))
  def testSliceWeakTypes(self, shape, dtype, weak_type, slc):
    rng = jtu.rand_default(self.rng())
    x = lax_internal._convert_element_type(rng(shape, dtype),
                                           weak_type=weak_type)
    op = lambda x: x[slc]
    self.assertEqual(op(x).aval.weak_type, weak_type)
    self.assertEqual(jax.jit(op)(x).aval.weak_type, weak_type)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_{}_axis={}_{}sections".format(
          jtu.format_shape_dtype_string(shape, dtype), axis, num_sections),
       "shape": shape, "num_sections": num_sections, "axis": axis,
       "dtype": dtype}
      for shape, axis, num_sections in [
          ((3,), 0, 3), ((12,), 0, 3), ((12, 4), 0, 4), ((12, 4), 1, 2),
          ((2, 3, 4), -1, 2), ((2, 3, 4), -2, 3)]
      for dtype in default_dtypes))
  def testSplitStaticInt(self, shape, num_sections, axis, dtype):
    rng = jtu.rand_default(self.rng())
    np_fun = lambda x: np.split(x, num_sections, axis=axis)
    jnp_fun = lambda x: jnp.split(x, num_sections, axis=axis)
    args_maker = lambda: [rng(shape, dtype)]
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)
    self._CompileAndCheck(jnp_fun, args_maker)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_{}_axis={}_{}sections".format(
          jtu.format_shape_dtype_string(shape, dtype), axis, num_sections),
       "shape": shape, "num_sections": num_sections, "axis": axis, "dtype": dtype}
      # All testcases split the specified axis unequally
      for shape, axis, num_sections in [
          ((3,), 0, 2), ((12,), 0, 5), ((12, 4), 0, 7), ((12, 4), 1, 3),
          ((2, 3, 5), -1, 2), ((2, 4, 4), -2, 3), ((7, 2, 2), 0, 3)]
      for dtype in default_dtypes))
  def testArraySplitStaticInt(self, shape, num_sections, axis, dtype):
    rng = jtu.rand_default(self.rng())
    np_fun = lambda x: np.array_split(x, num_sections, axis=axis)
    jnp_fun = lambda x: jnp.array_split(x, num_sections, axis=axis)
    args_maker = lambda: [rng(shape, dtype)]
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)
    self._CompileAndCheck(jnp_fun, args_maker)

  def testSplitTypeError(self):
    # If we pass an ndarray for indices_or_sections -> no error
    self.assertEqual(3, len(jnp.split(jnp.zeros(3), jnp.array([1, 2]))))

    CONCRETIZATION_MSG = "Abstract tracer value encountered where concrete value is expected."
    with self.assertRaisesRegex(TypeError, CONCRETIZATION_MSG):
      # An abstract tracer for idx
      jax.jit(lambda idx: jnp.split(jnp.zeros((12, 2)), idx))(2.)
    with self.assertRaisesRegex(TypeError, CONCRETIZATION_MSG):
      # A list including an abstract tracer
      jax.jit(lambda idx: jnp.split(jnp.zeros((12, 2)), [2, idx]))(2.)

    # A concrete tracer -> no error
    jax.jvp(lambda idx: jnp.split(jnp.zeros((12, 2)), idx),
            (2.,), (1.,))
    # A tuple including a concrete tracer -> no error
    jax.jvp(lambda idx: jnp.split(jnp.zeros((12, 2)), (1, idx)),
            (2.,), (1.,))

  @parameterized.named_parameters(jtu.cases_from_list(
    {"testcase_name": "_{}_bins={}_range={}_weights={}".format(
      jtu.format_shape_dtype_string(shape, dtype), bins, range, weights),
      "shape": shape,
      "dtype": dtype,
      "bins": bins,
      "range": range,
      "weights": weights,
    }
    for shape in [(5,), (5, 5)]
    for dtype in number_dtypes
    for bins in [10, np.arange(-5, 6), np.array([-5, 0, 3])]
    for range in [None, (0, 0), (0, 10)]
    for weights in [True, False]
  ))
  def testHistogramBinEdges(self, shape, dtype, bins, range, weights):
    rng = jtu.rand_default(self.rng())
    _weights = lambda w: abs(w) if weights else None
    np_fun = lambda a, w, r: np.histogram_bin_edges(a, bins=bins, range=r,
                                                    weights=_weights(w))
    jnp_fun = lambda a, w, r: jnp.histogram_bin_edges(a, bins=bins, range=r,
                                                      weights=_weights(w))
    args_maker = lambda: [rng(shape, dtype), rng(shape, dtype), range]
    tol = {jnp.bfloat16: 2E-2, np.float16: 1E-2}
    # linspace() compares poorly to numpy when using bfloat16
    if dtype != jnp.bfloat16:
      self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker, check_dtypes=False, tol=tol)
    self._CompileAndCheck(jnp_fun, args_maker,
                          atol=tol, rtol=tol)

  @parameterized.named_parameters(jtu.cases_from_list(
    {"testcase_name": "_{}_bins={}_density={}_weights={}".format(
      jtu.format_shape_dtype_string(shape, dtype), bins, density, weights),
      "shape": shape,
      "dtype": dtype,
      "bins": bins,
      "density": density,
      "weights": weights,
    }
    for shape in [(5,), (5, 5)]
    for dtype in default_dtypes
    # We only test explicit integer-valued bin edges because in other cases
    # rounding errors lead to flaky tests.
    for bins in [np.arange(-5, 6), np.array([-5, 0, 3])]
    for density in [True, False]
    for weights in [True, False]
  ))
  def testHistogram(self, shape, dtype, bins, density, weights):
    rng = jtu.rand_default(self.rng())
    _weights = lambda w: abs(w) if weights else None
    np_fun = lambda a, w: np.histogram(a, bins=bins, density=density,
                                         weights=_weights(w))
    jnp_fun = lambda a, w: jnp.histogram(a, bins=bins, density=density,
                                         weights=_weights(w))
    args_maker = lambda: [rng(shape, dtype), rng(shape, dtype)]
    tol = {jnp.bfloat16: 2E-2, np.float16: 1E-1}
    # np.searchsorted errors on bfloat16 with
    # "TypeError: invalid type promotion with custom data type"
    if dtype != jnp.bfloat16:
      self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker, check_dtypes=False,
                              tol=tol)
    self._CompileAndCheck(jnp_fun, args_maker)

  @parameterized.named_parameters(jtu.cases_from_list(
    {"testcase_name": "_{}_bins={}_weights={}_density={}_range={}".format(
      jtu.format_shape_dtype_string(shape, dtype), bins, weights, density, range),
      "shape": shape, "dtype": dtype, "bins": bins, "weights": weights, "density": density, "range": range,
    }
    for shape in [(5,), (12,)]
    for dtype in int_dtypes
    for bins in [2, [2, 2], [np.array([0, 1, 3, 5]), np.array([0, 2, 3, 4, 6])]]
    for weights in [False, True]
    for density in [False, True]
    for range in [None, [(-1, 1), None], [(-1, 1), (-2, 2)]]
  ))
  def testHistogram2d(self, shape, dtype, bins, weights, density, range):
    rng = jtu.rand_default(self.rng())
    _weights = lambda w: abs(w) if weights else None
    np_fun = jtu.ignore_warning(category=RuntimeWarning, message="invalid value.*")(
        lambda a, b, w: np.histogram2d(a, b, bins=bins, weights=_weights(w), density=density, range=range))
    jnp_fun = lambda a, b, w: jnp.histogram2d(a, b, bins=bins, weights=_weights(w), density=density, range=range)
    args_maker = lambda: [rng(shape, dtype), rng(shape, dtype), rng(shape, dtype)]
    tol = {jnp.bfloat16: 2E-2, np.float16: 1E-1}
    # np.searchsorted errors on bfloat16 with
    # "TypeError: invalid type promotion with custom data type"
    with np.errstate(divide='ignore', invalid='ignore'):
      if dtype != jnp.bfloat16:
        self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker, check_dtypes=False,
                          tol=tol)
    self._CompileAndCheck(jnp_fun, args_maker)

  @parameterized.named_parameters(jtu.cases_from_list(
    {"testcase_name": "_{}_bins={}_weights={}_density={}_range={}".format(
      jtu.format_shape_dtype_string(shape, dtype), bins, weights, density, range),
      "shape": shape, "dtype": dtype, "bins": bins, "weights": weights, "density": density, "range": range,
    }
    for shape in [(5, 3), (10, 3)]
    for dtype in int_dtypes
    for bins in [(2, 2, 2), [np.array([-5, 0, 4]), np.array([-4, -1, 2]), np.array([-6, -1, 4])]]
    for weights in [False, True]
    for density in [False, True]
    for range in [None, [(-1, 1), None, None], [(-1, 1), (-2, 2), (-3, 3)]]
  ))
  def testHistogramdd(self, shape, dtype, bins, weights, density, range):
    rng = jtu.rand_default(self.rng())
    _weights = lambda w: abs(w) if weights else None
    np_fun = jtu.ignore_warning(category=RuntimeWarning, message="invalid value.*")(
        lambda a, w: np.histogramdd(a, bins=bins, weights=_weights(w), density=density, range=range))
    jnp_fun = lambda a, w: jnp.histogramdd(a, bins=bins, weights=_weights(w), density=density, range=range)
    args_maker = lambda: [rng(shape, dtype), rng((shape[0],), dtype)]
    tol = {jnp.bfloat16: 2E-2, np.float16: 1E-1}
    # np.searchsorted errors on bfloat16 with
    # "TypeError: invalid type promotion with custom data type"
    if dtype != jnp.bfloat16:
      self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker, check_dtypes=False,
                            tol=tol)
    self._CompileAndCheck(jnp_fun, args_maker)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_{}_axis={}_{}sections".format(
          jtu.format_shape_dtype_string(shape, dtype), axis, num_sections),
       "shape": shape, "num_sections": num_sections, "axis": axis,
       "dtype": dtype}
      for shape, axis, num_sections in [
          ((12, 4), 0, 4), ((12, 4), 1, 2),
          ((2, 3, 4), 2, 2), ((4, 3, 4), 0, 2)]
      for dtype in default_dtypes))
  def testHVDSplit(self, shape, num_sections, axis, dtype):
    rng = jtu.rand_default(self.rng())
    def fn(module, axis):
      if axis == 0:
        return module.vsplit
      elif axis == 1:
        return module.hsplit
      else:
        assert axis == 2
        return module.dsplit

    np_fun = lambda x: fn(np, axis)(x, num_sections)
    jnp_fun = lambda x: fn(jnp, axis)(x, num_sections)
    args_maker = lambda: [rng(shape, dtype)]
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)
    self._CompileAndCheck(jnp_fun, args_maker)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_inshape={}_outshape={}_order={}".format(
          jtu.format_shape_dtype_string(arg_shape, dtype),
          jtu.format_shape_dtype_string(out_shape, dtype),
          order),
       "arg_shape": arg_shape, "out_shape": out_shape, "dtype": dtype,
       "order": order}
      for dtype in default_dtypes
      for order in ["C", "F"]
      for arg_shape, out_shape in [
          (jtu.NUMPY_SCALAR_SHAPE, (1, 1, 1)),
          ((), (1, 1, 1)),
          ((7, 0), (0, 42, 101)),
          ((3, 4), 12),
          ((3, 4), (12,)),
          ((3, 4), -1),
          ((2, 1, 4), (-1,)),
          ((2, 2, 4), (2, 8))
      ]))
  def testReshape(self, arg_shape, out_shape, dtype, order):
    rng = jtu.rand_default(self.rng())
    np_fun = lambda x: np.reshape(x, out_shape, order=order)
    jnp_fun = lambda x: jnp.reshape(x, out_shape, order=order)
    args_maker = lambda: [rng(arg_shape, dtype)]
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)
    self._CompileAndCheck(jnp_fun, args_maker)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_inshape={}_outshape={}".format(
          jtu.format_shape_dtype_string(arg_shape, dtype),
          jtu.format_shape_dtype_string(out_shape, dtype)),
       "arg_shape": arg_shape, "out_shape": out_shape, "dtype": dtype}
      for dtype in default_dtypes
      for arg_shape, out_shape in [
          ((7, 0), (0, 42, 101)),
          ((2, 1, 4), (-1,)),
          ((2, 2, 4), (2, 8))
      ]))
  def testReshapeMethod(self, arg_shape, out_shape, dtype):
    rng = jtu.rand_default(self.rng())
    np_fun = lambda x: np.reshape(x, out_shape)
    jnp_fun = lambda x: x.reshape(*out_shape)
    args_maker = lambda: [rng(arg_shape, dtype)]
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)
    self._CompileAndCheck(jnp_fun, args_maker)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_inshape={}_outshape={}".format(
          jtu.format_shape_dtype_string(arg_shape, dtype),
          jtu.format_shape_dtype_string(out_shape, dtype)),
       "arg_shape": arg_shape, "out_shape": out_shape, "dtype": dtype}
      for dtype in default_dtypes
      for arg_shape, out_shape in itertools.product(all_shapes, array_shapes)))
  def testResize(self, arg_shape, out_shape, dtype):
    rng = jtu.rand_default(self.rng())
    np_fun = lambda x: np.resize(x, out_shape)
    jnp_fun = lambda x: jnp.resize(x, out_shape)
    args_maker = lambda: [rng(arg_shape, dtype)]
    if len(out_shape) > 0 or numpy_version >= (1, 20, 0):
      self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)
    self._CompileAndCheck(jnp_fun, args_maker)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_inshape={}_expanddim={!r}".format(
          jtu.format_shape_dtype_string(arg_shape, dtype), dim),
       "arg_shape": arg_shape, "dtype": dtype, "dim": dim}
      for arg_shape in [(), (3,), (3, 4)]
      for dtype in default_dtypes
      for dim in (list(range(-len(arg_shape)+1, len(arg_shape)))
                  + [np.array(0), np.array(-1), (0,), [np.array(0)],
                     (len(arg_shape), len(arg_shape) + 1)])))
  def testExpandDimsStaticDim(self, arg_shape, dtype, dim):
    rng = jtu.rand_default(self.rng())
    np_fun = lambda x: np.expand_dims(x, dim)
    jnp_fun = lambda x: jnp.expand_dims(x, dim)
    args_maker = lambda: [rng(arg_shape, dtype)]
    self._CompileAndCheck(jnp_fun, args_maker)
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)

  def testExpandDimsRepeatedAxisError(self):
    x = jnp.ones((2, 3))
    self.assertRaisesRegex(
        ValueError, 'repeated axis.*',
        lambda: jnp.expand_dims(x, [1, 1]))
    self.assertRaisesRegex(
        ValueError, 'repeated axis.*',
        lambda: jnp.expand_dims(x, [3, -1]))

    # ensure this is numpy's behavior too, so that we remain consistent
    x = np.ones((2, 3))
    self.assertRaisesRegex(
        ValueError, 'repeated axis.*',
        lambda: np.expand_dims(x, [1, 1]))
    self.assertRaisesRegex(
        ValueError, 'repeated axis.*',
        lambda: np.expand_dims(x, [3, -1]))

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_inshape={}_axes=({},{})".format(
          jtu.format_shape_dtype_string(arg_shape, dtype), ax1, ax2),
       "arg_shape": arg_shape, "dtype": dtype, "ax1": ax1, "ax2": ax2}
      for arg_shape, ax1, ax2 in [
          ((3, 4), 0, 1), ((3, 4), 1, 0), ((3, 4, 5), 1, 2),
          ((3, 4, 5), -1, -2), ((3, 4, 5), 0, 1)]
      for dtype in default_dtypes))
  def testSwapAxesStaticAxes(self, arg_shape, dtype, ax1, ax2):
    rng = jtu.rand_default(self.rng())
    np_fun = lambda x: np.swapaxes(x, ax1, ax2)
    jnp_fun = lambda x: jnp.swapaxes(x, ax1, ax2)
    args_maker = lambda: [rng(arg_shape, dtype)]
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)
    self._CompileAndCheck(jnp_fun, args_maker)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_inshape={}_axis={!r}".format(
          jtu.format_shape_dtype_string(arg_shape, dtype), ax),
       "arg_shape": arg_shape, "dtype": dtype, "ax": ax}
      for arg_shape, ax in [
          ((3, 1), None),
          ((3, 1), 1),
          ((3, 1), -1),
          ((3, 1), np.array(1)),
          ((1, 3, 1), (0, 2)),
          ((1, 3, 1), (0,)),
          ((1, 4, 1), (np.array(0),))]
      for dtype in default_dtypes))
  def testSqueeze(self, arg_shape, dtype, ax):
    rng = jtu.rand_default(self.rng())
    np_fun = lambda x: np.squeeze(x, ax)
    jnp_fun = lambda x: jnp.squeeze(x, ax)
    args_maker = lambda: [rng(arg_shape, dtype)]
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)
    self._CompileAndCheck(jnp_fun, args_maker)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_shape={}_axis={}_weights={}_returned={}".format(
          jtu.format_shape_dtype_string(shape, dtype),
          axis,
          (None if weights_shape is None else jtu.format_shape_dtype_string(weights_shape, dtype)),
          returned),
       "shape": shape, "dtype": dtype, "axis": axis,
       "weights_shape": weights_shape, "returned": returned}
      for shape, dtype in _shape_and_dtypes(nonempty_shapes, number_dtypes)
      for axis in list(range(-len(shape), len(shape))) + [None]
      # `weights_shape` is either `None`, same as the averaged axis, or same as
      # that of the input
      for weights_shape in ([None, shape] if axis is None or len(shape) == 1
                            else [None, (shape[axis],), shape])
      for returned in [False, True]))
  def testAverage(self, shape, dtype, axis, weights_shape, returned):
    rng = jtu.rand_default(self.rng())
    if weights_shape is None:
      np_fun = lambda x: np.average(x, axis, returned=returned)
      jnp_fun = lambda x: jnp.average(x, axis, returned=returned)
      args_maker = lambda: [rng(shape, dtype)]
    else:
      np_fun = lambda x, weights: np.average(x, axis, weights, returned)
      jnp_fun = lambda x, weights: jnp.average(x, axis, weights, returned)
      args_maker = lambda: [rng(shape, dtype), rng(weights_shape, dtype)]
    np_fun = _promote_like_jnp(np_fun, inexact=True)
    tol = {dtypes.bfloat16: 2e-1, np.float16: 1e-2, np.float32: 1e-5,
           np.float64: 1e-12, np.complex64: 1e-5}
    check_dtypes = shape is not jtu.PYTHON_SCALAR_SHAPE and numpy_version >= (1, 22)
    try:
      self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker,
                              check_dtypes=check_dtypes, tol=tol)
    except ZeroDivisionError:
      self.skipTest("don't support checking for ZeroDivisionError")
    self._CompileAndCheck(jnp_fun, args_maker, check_dtypes=check_dtypes,
                          rtol=tol, atol=tol)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name":
       f"_arg{i}_ndmin={ndmin}_dtype={np.dtype(dtype) if dtype else None}",
       "arg": arg, "ndmin": ndmin, "dtype": dtype}
      for i, (arg, dtypes) in enumerate([
          ([True, False, True], all_dtypes),
          (3., all_dtypes),
          ([1, 2, 3], all_dtypes),
          (np.array([1, 2, 3], dtype=np.int64), all_dtypes),
          ([1., 2., 3.], all_dtypes),
          ([[1, 2], [3, 4], [5, 6]], all_dtypes),
          ([[1, 2.], [3, 4], [5, 6]], all_dtypes),
          ([[1., 2j], [3., 4.], [5., 6.]], complex_dtypes),
          ([[3, np.array(2, dtype=jnp.float_), 1],
           np.arange(3., dtype=jnp.float_)], all_dtypes),
      ])
      for dtype in [None] + dtypes
      for ndmin in [None, np.ndim(arg), np.ndim(arg) + 1, np.ndim(arg) + 2]))
  def testArray(self, arg, ndmin, dtype):
    args_maker = lambda: [arg]
    canonical_dtype = dtypes.canonicalize_dtype(dtype or np.array(arg).dtype)
    if ndmin is not None:
      np_fun = partial(np.array, ndmin=ndmin, dtype=canonical_dtype)
      jnp_fun = partial(jnp.array, ndmin=ndmin, dtype=dtype)
    else:
      np_fun = partial(np.array, dtype=canonical_dtype)
      jnp_fun = partial(jnp.array, dtype=dtype)

    # We are testing correct canonicalization behavior here, so we turn off the
    # permissive canonicalization logic in the test harness.
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker,
                            canonicalize_dtypes=False)
    self._CompileAndCheck(jnp_fun, args_maker)

  @jtu.ignore_warning(category=UserWarning, message="Explicitly requested dtype.*")
  def testArrayDtypeInference(self):
    def _check(obj, out_dtype, weak_type):
      dtype_reference = np.array(obj, dtype=out_dtype)

      out = jnp.array(obj)
      self.assertDtypesMatch(out, dtype_reference)
      self.assertEqual(dtypes.is_weakly_typed(out), weak_type)

      out_jit = jax.jit(jnp.array)(obj)
      self.assertDtypesMatch(out_jit, dtype_reference)
      self.assertEqual(dtypes.is_weakly_typed(out_jit), weak_type)

    # Python scalars become 64-bit weak types.
    _check(1, np.int64, True)
    _check(1.0, np.float64, True)
    _check(1.0j, np.complex128, True)

    # Lists become strongly-typed defaults.
    _check([1], jnp.int_, False)
    _check([1.0], jnp.float_, False)
    _check([1.0j], jnp.complex_, False)

    # Lists of weakly-typed objects become strongly-typed defaults.
    _check([jnp.array(1)], jnp.int_, False)
    _check([jnp.array(1.0)], jnp.float_, False)
    _check([jnp.array(1.0j)], jnp.complex_, False)

    # Lists of strongly-typed objects maintain their strong type.
    _check([jnp.int64(1)], np.int64, False)
    _check([jnp.float64(1)], np.float64, False)
    _check([jnp.complex128(1)], np.complex128, False)

    # Mixed inputs use JAX-style promotion.
    # (regression test for https://github.com/google/jax/issues/8945)
    _check([0, np.int16(1)], np.int16, False)
    _check([0.0, np.float16(1)], np.float16, False)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": f"_dtype={np.dtype(dtype)}_func={func}",
       "dtype": dtype, "func": func}
      for dtype in all_dtypes
      for func in ["array", "copy", "copy.copy", "copy.deepcopy"]))
  def testArrayCopy(self, dtype, func):
    x = jnp.ones(10, dtype=dtype)
    if func == "copy.deepcopy":
      copy_func = copy.deepcopy
    elif func == "copy.copy":
      copy_func = copy.copy
    else:
      copy_func = getattr(jnp, func)

    x_view = jnp.asarray(x)
    x_view_jit = jax.jit(jnp.asarray)(x)
    x_copy = copy_func(x)
    x_copy_jit = jax.jit(copy_func)(x)

    _ptr = lambda x: x.device_buffer.unsafe_buffer_pointer()

    self.assertEqual(_ptr(x), _ptr(x_view))
    self.assertEqual(_ptr(x), _ptr(x_view_jit))
    self.assertNotEqual(_ptr(x), _ptr(x_copy))
    self.assertNotEqual(_ptr(x), _ptr(x_copy_jit))

    x.delete()

    self.assertTrue(x_view.is_deleted())
    self.assertTrue(x_view_jit.is_deleted())

    self.assertFalse(x_copy.is_deleted())
    self.assertFalse(x_copy_jit.is_deleted())

  def testArrayCopyAutodiff(self):
    f = lambda x: jnp.array(x, copy=True)

    x = jnp.ones(10)
    xdot = jnp.ones(10)
    y, ydot = jax.jvp(f, (x,), (xdot,))
    self.assertIsNot(x, y)
    self.assertIsNot(xdot, ydot)

    ybar = jnp.ones(10)
    y, f_vjp = jax.vjp(f, x)
    xbar, = f_vjp(ybar)
    self.assertIsNot(x, y)
    self.assertIsNot(xbar, ybar)

  def testArrayCopyVmap(self):
    f = lambda x: jnp.array(x, copy=True)
    x = jnp.ones(10)
    y = jax.vmap(f)(x)
    self.assertIsNot(x, y)

  def testArrayUnsupportedDtypeError(self):
    with self.assertRaisesRegex(TypeError,
                                "JAX only supports number and bool dtypes.*"):
      jnp.array(3, [('a','<i4'),('b','<i4')])

  def testArrayFromInteger(self):
    int_dtype = dtypes.canonicalize_dtype(jnp.int64)
    int_max = jnp.iinfo(int_dtype).max
    int_min = jnp.iinfo(int_dtype).min

    # Values at extremes are converted correctly.
    for val in [int_min, 0, int_max]:
      self.assertEqual(jnp.array(val).dtype, int_dtype)

    # out of bounds leads to an OverflowError
    val = int_max + 1
    with self.assertRaisesRegex(OverflowError, f"Python int {val} too large to convert to {int_dtype.name}"):
      jnp.array(val)

    # explicit uint64 should work
    if config.x64_enabled:
      self.assertEqual(np.uint64(val), jnp.array(val, dtype='uint64'))

  def testArrayFromList(self):
    int_max = jnp.iinfo(jnp.int64).max
    int_min = jnp.iinfo(jnp.int64).min

    # Values at extremes are converted correctly.
    for val in [int_min, 0, int_max]:
      self.assertEqual(jnp.array([val]).dtype, dtypes.canonicalize_dtype('int64'))

    # list of values results in promoted type.
    self.assertEqual(jnp.array([0, np.float16(1)]).dtype, jnp.result_type('int64', 'float16'))

    # out of bounds leads to an OverflowError
    val = int_min - 1
    with self.assertRaisesRegex(OverflowError, "Python int too large.*"):
      jnp.array([0, val])

  def testIssue121(self):
    assert not np.isscalar(jnp.array(3))

  def testArrayOutputsDeviceArrays(self):
    assert device_array.type_is_device_array(jnp.array([]))
    assert device_array.type_is_device_array(jnp.array(np.array([])))

    class NDArrayLike:
      def __array__(self, dtype=None):
        return np.array([], dtype=dtype)
    assert device_array.type_is_device_array(jnp.array(NDArrayLike()))

    # NOTE(mattjj): disabled b/c __array__ must produce ndarrays
    # class DeviceArrayLike:
    #     def __array__(self, dtype=None):
    #         return jnp.array([], dtype=dtype)
    # assert  xla.type_is_device_array(jnp.array(DeviceArrayLike()))

  def testArrayMethod(self):
    class arraylike:
      dtype = np.dtype('float32')
      def __array__(self, dtype=None):
        return np.array(3., dtype=dtype)
    a = arraylike()
    ans = jnp.array(a)
    self.assertEqual(ans, 3.)

  def testJaxArrayOps(self):
    class arraylike:
      def __jax_array__(self):
        return jnp.array(3.)
    self.assertArraysEqual(arraylike() * jnp.arange(10), jnp.array(3.) * jnp.arange(10))

  def testMemoryView(self):
    self.assertAllClose(
        jnp.array(bytearray(b'\x2a')),
        np.array(bytearray(b'\x2a'))
    )
    self.assertAllClose(
        jnp.array(bytearray(b'\x2a\xf3'), ndmin=2),
        np.array(bytearray(b'\x2a\xf3'), ndmin=2)
    )

  def testIsClose(self):
    c_isclose = jax.jit(jnp.isclose)
    c_isclose_nan = jax.jit(partial(jnp.isclose, equal_nan=True))
    n = 2

    rng = self.rng()
    x = rng.randn(n, 1)
    y = rng.randn(n, 1)
    inf = np.asarray(n * [np.inf]).reshape([n, 1])
    nan = np.asarray(n * [np.nan]).reshape([n, 1])
    args = [x, y, inf, -inf, nan]

    for arg0 in args:
      for arg1 in args:
        result_np = np.isclose(arg0, arg1)
        result_jax = jnp.isclose(arg0, arg1)
        result_jit = c_isclose(arg0, arg1)
        self.assertTrue(jnp.all(jnp.equal(result_np, result_jax)))
        self.assertTrue(jnp.all(jnp.equal(result_np, result_jit)))
        result_np = np.isclose(arg0, arg1, equal_nan=True)
        result_jax = jnp.isclose(arg0, arg1, equal_nan=True)
        result_jit = c_isclose_nan(arg0, arg1)
        self.assertTrue(jnp.all(jnp.equal(result_np, result_jax)))
        self.assertTrue(jnp.all(jnp.equal(result_np, result_jit)))

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": f"_x={x}_y={y}_equal_nan={equal_nan}",
       "x": x, "y": y, "equal_nan": equal_nan}
      for x, y in itertools.product([
         1, [1], [1, 1 + 1E-4], [1, np.nan]], repeat=2)
      for equal_nan in [True, False]))
  def testAllClose(self, x, y, equal_nan):
    jnp_fun = partial(jnp.allclose, equal_nan=equal_nan, rtol=1E-3)
    np_fun = partial(np.allclose, equal_nan=equal_nan, rtol=1E-3)
    args_maker = lambda: [np.array(x), np.array(y)]
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)
    self._CompileAndCheck(jnp_fun, args_maker)

  def testZeroStridesConstantHandler(self):
    raw_const = self.rng().randn(1, 2, 1, 1, 5, 1)
    const = np.broadcast_to(raw_const, (3, 2, 3, 4, 5, 6))

    def fun(x):
      return x * const

    fun = jax.jit(fun)
    out_val = fun(3.)
    self.assertAllClose(out_val, 3. * const, check_dtypes=False)

  def testIsInstanceNdarrayDuringTracing(self):
    arr = np.ones(3)

    @jax.jit
    def f(x):
      self.assertIsInstance(x, jnp.ndarray)
      return jnp.sum(x)

    f(arr)

  def testNonArrayErrorMessage(self):
    x = [1., 2.]
    y = np.array([3., 4.])

    def g(x, y):
      return jnp.add(x, y)

    def f(x, y):
      return jnp.dot(x, y)

    self.assertRaises(TypeError, lambda: g(x, y))
    self.assertRaises(TypeError, lambda: f(x, y))
    self.assertRaises(TypeError, lambda: jax.jit(g)(x, y))
    self.assertRaises(TypeError, lambda: jax.jit(f)(x, y))

  def testAbstractionErrorMessage(self):

    @jax.jit
    def f(x, n):
      for _ in range(n):
        x = x * x
      return x

    self.assertRaises(jax.errors.TracerIntegerConversionError, lambda: f(3., 3))

    @jax.jit
    def g(x):
      if x > 0.:
        return x * 2
      else:
        return x + 2

    self.assertRaises(jax.errors.ConcretizationTypeError, lambda: g(3.))

  def testTracingPrimitiveWithNoTranslationErrorMessage(self):
    # TODO(mattjj): update this for jax3
    self.skipTest("test needs jax3 update")
    foo = jnp._not_implemented(lambda x: x)

    # No error if there's no tracing.
    foo(np.arange(3))

    cfoo = jax.jit(foo)
    self.assertRaises(NotImplementedError, lambda: cfoo(np.arange(3)))

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_{}_axis={}".format(
          jtu.format_shape_dtype_string(shape, dtype), axis),
       "shape": shape, "dtype": dtype, "axis": axis}
      for shape in [(3,), (2, 3)]
      for dtype in default_dtypes
      for axis in list(range(-len(shape), len(shape))) + [None] + [tuple(range(len(shape)))]  # Test negative axes and tuples
    ))
  def testFlip(self, shape, dtype, axis):
    rng = jtu.rand_default(self.rng())
    args_maker = self._GetArgsMaker(rng, [shape], [dtype])
    jnp_op = lambda x: jnp.flip(x, axis)
    np_op = lambda x: np.flip(x, axis)
    self._CheckAgainstNumpy(np_op, jnp_op, args_maker)
    self._CompileAndCheck(jnp_op, args_maker)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_{}".format(
          jtu.format_shape_dtype_string(shape, dtype)),
       "shape": shape, "dtype": dtype}
      for shape in [(3,), (2, 3), (3, 2, 4)]
      for dtype in default_dtypes))
  def testFlipud(self, shape, dtype):
    rng = jtu.rand_default(self.rng())
    args_maker = self._GetArgsMaker(rng, [shape], [dtype])
    jnp_op = lambda x: jnp.flipud(x)
    np_op = lambda x: np.flipud(x)
    self._CheckAgainstNumpy(np_op, jnp_op, args_maker)
    self._CompileAndCheck(jnp_op, args_maker)


  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_{}".format(
          jtu.format_shape_dtype_string(shape, dtype)),
       "shape": shape, "dtype": dtype}
      for shape in [(3, 2), (2, 3), (3, 2, 4)]
      for dtype in default_dtypes))
  def testFliplr(self, shape, dtype):
    rng = jtu.rand_default(self.rng())
    args_maker = self._GetArgsMaker(rng, [shape], [dtype])
    jnp_op = lambda x: jnp.fliplr(x)
    np_op = lambda x: np.fliplr(x)
    self._CheckAgainstNumpy(np_op, jnp_op, args_maker)
    self._CompileAndCheck(jnp_op, args_maker)


  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_{}_k={}_axes={}".format(
          jtu.format_shape_dtype_string(shape, dtype), k, axes),
       "shape": shape, "dtype": dtype, "k": k, "axes": axes}
      for shape, axes in [
          [(2, 3), (0, 1)],
          [(2, 3), (1, 0)],
          [(4, 3, 2), (0, 2)],
          [(4, 3, 2), (2, 1)],
      ]
      for k in range(-3, 4)
      for dtype in default_dtypes))
  def testRot90(self, shape, dtype, k, axes):
    rng = jtu.rand_default(self.rng())
    args_maker = self._GetArgsMaker(rng, [shape], [dtype])
    jnp_op = lambda x: jnp.rot90(x, k, axes)
    np_op = lambda x: np.rot90(x, k, axes)
    self._CheckAgainstNumpy(np_op, jnp_op, args_maker)
    self._CompileAndCheck(jnp_op, args_maker)

  # TODO(mattjj): test infix operator overrides

  def testRavel(self):
    rng = self.rng()
    args_maker = lambda: [rng.randn(3, 4).astype("float32")]
    self._CompileAndCheck(lambda x: x.ravel(), args_maker)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_shape={}_order={}_mode={}".format(
          shape, order, mode),
        "shape": shape, "order": order, "mode": mode}
      for shape in nonempty_nonscalar_array_shapes
      for order in ['C', 'F']
      for mode in ['wrap', 'clip', 'raise']))
  @jax.numpy_rank_promotion('allow')  # This test explicitly exercises implicit rank promotion.
  def testRavelMultiIndex(self, shape, order, mode):
    # generate indices in each dimension with a few out of bounds.
    rngs = [jtu.rand_int(self.rng(), low=-1, high=dim + 1)
            for dim in shape]
    # generate multi_indices of different dimensions that broadcast.
    args_maker = lambda: [tuple(rng(ndim * (3,), jnp.int_)
                                for ndim, rng in enumerate(rngs))]
    def np_fun(x):
      try:
        return np.ravel_multi_index(x, shape, order=order, mode=mode)
      except ValueError as err:
        if str(err).startswith('invalid entry'):
          # sentinel indicating expected error.
          return -999
        else:
          raise
    def jnp_fun(x):
      try:
        return jnp.ravel_multi_index(x, shape, order=order, mode=mode)
      except ValueError as err:
        if str(err).startswith('invalid entry'):
          # sentinel indicating expected error.
          return -999
        else:
          raise
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker, check_dtypes=False)
    if mode == 'raise':
      msg = ("The error occurred because ravel_multi_index was jit-compiled "
             "with mode='raise'. Use mode='wrap' or mode='clip' instead.")
      with self.assertRaisesRegex(jax.core.ConcretizationTypeError, msg):
        jax.jit(jnp_fun)(*args_maker())
    else:
      self._CompileAndCheck(jnp_fun, args_maker)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_ashape={}{}_cshapes={}{}_mode={}".format(
          adtype.__name__, ashape, cdtype.__name__, cshapes, mode),
        "ashape": ashape, "adtype": adtype, "cshapes": cshapes, "cdtype": cdtype, "mode": mode}
      for ashape in ((), (4,), (3, 4))
      for cshapes in [
        [(), (4,)],
        [(3, 4), (4,), (3, 1)]
      ]
      for adtype in int_dtypes
      for cdtype in default_dtypes
      for mode in ['wrap', 'clip', 'raise']))
  def testChoose(self, ashape, adtype, cshapes, cdtype, mode):
    rng = jtu.rand_default(self.rng())
    args_maker = lambda: [rng(ashape, adtype), [rng(s, cdtype) for s in cshapes]]
    def np_fun(a, c):
      try:
        return np.choose(a, c, mode=mode)
      except ValueError as err:
        if mode == 'raise' and str(err).startswith('invalid entry'):
          return -999  # sentinel indicating expected error.
        else:
          raise
    def jnp_fun(a, c):
      try:
        return jnp.choose(a, c, mode=mode)
      except ValueError as err:
        if mode == 'raise' and str(err).startswith('invalid entry'):
          return -999  # sentinel indicating expected error.
        else:
          raise
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker, check_dtypes=False)
    if mode == 'raise':
      msg = ("The error occurred because jnp.choose was jit-compiled"
             " with mode='raise'. Use mode='wrap' or mode='clip' instead.")
      with self.assertRaisesRegex(jax.core.ConcretizationTypeError, msg):
        jax.jit(jnp_fun)(*args_maker())
    else:
      self._CompileAndCheck(jnp_fun, args_maker)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_shape={}_idx={}".format(shape,
        jtu.format_shape_dtype_string(idx_shape, dtype)),
       "shape": shape, "idx_shape": idx_shape, "dtype": dtype}
      for shape in nonempty_nonscalar_array_shapes
      for dtype in int_dtypes
      for idx_shape in all_shapes))
  def testUnravelIndex(self, shape, idx_shape, dtype):
    size = prod(shape)
    rng = jtu.rand_int(self.rng(), low=-((2 * size) // 3), high=(2 * size) // 3)

    def np_fun(index, shape):
      # JAX's version outputs the same dtype as the input in the typical case
      # where shape is weakly-typed.
      out_dtype = index.dtype
      # Adjust out-of-bounds behavior to match jax's documented behavior.
      index = np.clip(index, -size, size - 1)
      index = np.where(index < 0, index + size, index)
      return [i.astype(out_dtype) for i in np.unravel_index(index, shape)]

    jnp_fun = jnp.unravel_index
    args_maker = lambda: [rng(idx_shape, dtype), shape]
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)
    self._CompileAndCheck(jnp_fun, args_maker)

  def testAstype(self):
    rng = self.rng()
    args_maker = lambda: [rng.randn(3, 4).astype("float32")]
    np_op = lambda x: np.asarray(x).astype(jnp.int32)
    jnp_op = lambda x: jnp.asarray(x).astype(jnp.int32)
    self._CheckAgainstNumpy(np_op, jnp_op, args_maker)
    self._CompileAndCheck(jnp_op, args_maker)

  def testAstypeNone(self):
    rng = self.rng()
    args_maker = lambda: [rng.randn(3, 4).astype("int32")]
    np_op = jtu.with_jax_dtype_defaults(lambda x: np.asarray(x).astype(None))
    jnp_op = lambda x: jnp.asarray(x).astype(None)
    self._CheckAgainstNumpy(np_op, jnp_op, args_maker)
    self._CompileAndCheck(jnp_op, args_maker)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_{}".format(
          jtu.format_shape_dtype_string(shape, dtype)),
      "shape": shape, "dtype": dtype}
      for shape in array_shapes
      for dtype in all_dtypes))
  def testNbytes(self, shape, dtype):
    rng = jtu.rand_default(self.rng())
    np_op = lambda x: np.asarray(x).nbytes
    jnp_op = lambda x: jnp.asarray(x).nbytes
    args_maker = lambda: [rng(shape, dtype)]
    self._CheckAgainstNumpy(np_op, jnp_op, args_maker)
    self._CompileAndCheck(jnp_op, args_maker)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_{}".format(
          jtu.format_shape_dtype_string(shape, dtype)),
      "shape": shape, "dtype": dtype}
      for shape in array_shapes
      for dtype in all_dtypes))
  def testItemsize(self, shape, dtype):
    rng = jtu.rand_default(self.rng())
    np_op = lambda x: np.asarray(x).itemsize
    jnp_op = lambda x: jnp.asarray(x).itemsize
    args_maker = lambda: [rng(shape, dtype)]
    self._CheckAgainstNumpy(np_op, jnp_op, args_maker)
    self._CompileAndCheck(jnp_op, args_maker)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_{}_dtype={}".format(
          jtu.format_shape_dtype_string(shape, a_dtype), dtype),
      "shape": shape, "a_dtype": a_dtype, "dtype": dtype}
      for shape in [(8,), (3, 8)]  # last dim = 8 to ensure shape compatibility
      for a_dtype in (default_dtypes + unsigned_dtypes + bool_dtypes)
      for dtype in (default_dtypes + unsigned_dtypes + bool_dtypes)))
  def testView(self, shape, a_dtype, dtype):
    if jtu.device_under_test() == 'tpu':
      if jnp.dtype(a_dtype).itemsize in [1, 2] or jnp.dtype(dtype).itemsize in [1, 2]:
        self.skipTest("arr.view() not supported on TPU for 8- or 16-bit types.")
    if not config.x64_enabled:
      if jnp.dtype(a_dtype).itemsize == 8 or jnp.dtype(dtype).itemsize == 8:
        self.skipTest("x64 types are disabled by jax_enable_x64")
    rng = jtu.rand_fullrange(self.rng())
    args_maker = lambda: [rng(shape, a_dtype)]
    np_op = lambda x: np.asarray(x).view(dtype)
    jnp_op = lambda x: jnp.asarray(x).view(dtype)
    # Above may produce signaling nans; ignore warnings from invalid values.
    with np.errstate(invalid='ignore'):
      self._CheckAgainstNumpy(np_op, jnp_op, args_maker)
      self._CompileAndCheck(jnp_op, args_maker)

  def testPathologicalFloats(self):
    args_maker = lambda: [np.array([
      0b_0111_1111_1000_0000_0000_0000_0000_0000, # inf
      0b_1111_1111_1000_0000_0000_0000_0000_0000, # -inf
      0b_0111_1111_1100_0000_0000_0000_0000_0000, # qnan
      0b_1111_1111_1100_0000_0000_0000_0000_0000, # -qnan
      0b_0111_1111_1000_0000_0000_0000_0000_0001, # snan
      0b_1111_1111_1000_0000_0000_0000_0000_0001, # -snan
      0b_0111_1111_1000_0000_0000_1100_0000_0000, # nonstandard nan
      0b_1111_1111_1000_0000_0000_1100_0000_0000, # -nonstandard nan
      0b_0000_0000_0000_0000_0000_0000_0000_0000, # zero
      0b_1000_0000_0000_0000_0000_0000_0000_0000, # -zero
    ], dtype='uint32')]

    np_op = lambda x: np.asarray(x).view('float32').view('uint32')
    jnp_op = lambda x: jnp.asarray(x).view('float32').view('uint32')

    self._CheckAgainstNumpy(np_op, jnp_op, args_maker)
    self._CompileAndCheck(jnp_op, args_maker)

  # TODO(mattjj): test other ndarray-like method overrides

  def testNpMean(self):
    # from https://github.com/google/jax/issues/125
    x = jnp.eye(3, dtype=float) + 0.
    ans = np.mean(x)
    self.assertAllClose(ans, np.array(1./3), check_dtypes=False)

  def testArangeOnFloats(self):
    np_arange = jtu.with_jax_dtype_defaults(np.arange)
    # from https://github.com/google/jax/issues/145
    self.assertAllClose(np_arange(0.0, 1.0, 0.1),
                        jnp.arange(0.0, 1.0, 0.1))
    # from https://github.com/google/jax/issues/3450
    self.assertAllClose(np_arange(2.5),
                        jnp.arange(2.5))
    self.assertAllClose(np_arange(0., 2.5),
                        jnp.arange(0., 2.5))

  def testArangeTypes(self):
    # Test that arange() output type is equal to the default types.
    int_ = dtypes.canonicalize_dtype(jnp.int_)
    float_ = dtypes.canonicalize_dtype(jnp.float_)

    self.assertEqual(jnp.arange(10).dtype, int_)
    self.assertEqual(jnp.arange(10.).dtype, float_)
    self.assertEqual(jnp.arange(10, dtype='uint16').dtype, np.uint16)
    self.assertEqual(jnp.arange(10, dtype='bfloat16').dtype, jnp.bfloat16)

    self.assertEqual(jnp.arange(0, 10, 1).dtype, int_)
    self.assertEqual(jnp.arange(0, 10, 1.).dtype, float_)
    self.assertEqual(jnp.arange(0., 10, 1).dtype, float_)


  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_{}_axis={}".format(
          jtu.format_shape_dtype_string(shape, dtype), axis),
       "shape": shape, "dtype": dtype, "axis": axis}
      for dtype in all_dtypes
      for shape in nonzerodim_shapes
      for axis in (None, *range(len(shape)))))
  def testSort(self, dtype, shape, axis):
    rng = jtu.rand_some_equal(self.rng())
    args_maker = lambda: [rng(shape, dtype)]
    jnp_fun = jnp.sort
    np_fun = np.sort
    if axis is not None:
      jnp_fun = partial(jnp_fun, axis=axis)
      np_fun = partial(np_fun, axis=axis)
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)
    self._CompileAndCheck(jnp_fun, args_maker)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_{}_axis={}".format(
          jtu.format_shape_dtype_string(shape, dtype), axis),
        "shape": shape, "dtype": dtype, "axis": axis}
      for dtype in all_dtypes
      for shape in one_dim_array_shapes
      for axis in [None]))
  def testSortComplex(self, dtype, shape, axis):
    rng = jtu.rand_some_equal(self.rng())
    args_maker = lambda: [rng(shape, dtype)]
    self._CheckAgainstNumpy(np.sort_complex, jnp.sort_complex, args_maker, check_dtypes=False)
    self._CompileAndCheck(jnp.sort_complex, args_maker)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_{}_input_type={}_axis={}".format(
          jtu.format_shape_dtype_string(shape, dtype),
          input_type.__name__, axis),
       "shape": shape, "dtype": dtype, "input_type": input_type, "axis": axis}
      for dtype in all_dtypes
      for shape in nonempty_nonscalar_array_shapes
      for input_type in [np.array, tuple]
      for axis in (-1, *range(len(shape) - 1))))
  def testLexsort(self, dtype, shape, input_type, axis):
    rng = jtu.rand_some_equal(self.rng())
    args_maker = lambda: [input_type(rng(shape, dtype))]
    jnp_op = lambda x: jnp.lexsort(x, axis=axis)
    np_op = jtu.with_jax_dtype_defaults(lambda x: np.lexsort(x, axis=axis))
    self._CheckAgainstNumpy(np_op, jnp_op, args_maker)
    self._CompileAndCheck(jnp_op, args_maker)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_{}_axis={}".format(
          jtu.format_shape_dtype_string(shape, dtype), axis),
       "shape": shape, "dtype": dtype, "axis": axis}
      for dtype in all_dtypes
      for shape in nonzerodim_shapes
      for axis in (None, *range(len(shape)))))
  def testArgsort(self, dtype, shape, axis):
    rng = jtu.rand_some_equal(self.rng())
    args_maker = lambda: [rng(shape, dtype)]
    jnp_fun = jnp.argsort
    np_fun = jtu.with_jax_dtype_defaults(np.argsort)
    if axis is not None:
      jnp_fun = partial(jnp_fun, axis=axis)
      np_fun = partial(np_fun, axis=axis)
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)
    self._CompileAndCheck(jnp_fun, args_maker)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_{}".format(
          jtu.format_shape_dtype_string(shape, dtype)),
       "shape": shape, "dtype": dtype}
      for dtype in all_dtypes
      for shape in nonzerodim_shapes))
  def testMsort(self, dtype, shape):
    rng = jtu.rand_some_equal(self.rng())
    args_maker = lambda: [rng(shape, dtype)]
    self._CheckAgainstNumpy(np.msort, jnp.msort, args_maker)
    self._CompileAndCheck(jnp.msort, args_maker)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_{}_shifts={}_axis={}".format(
          jtu.format_shape_dtype_string(shape, dtype),
          shifts, axis),
       "shape": shape, "dtype": dtype, "shifts": shifts, "axis": axis}
      for dtype in all_dtypes
      for shape in [(3, 4), (3, 4, 5), (7, 4, 0)]
      for shifts, axis in [
        (3, None),
        (1, 1),
        ((3,), (0,)),
        ((-2,), (-2,)),
        ((1, 2), (0, -1)),
        ((4, 2, 5, 5, 2, 4), None),
        (100, None),
      ]))
  def testRoll(self, shape, dtype, shifts, axis):
    rng = jtu.rand_default(self.rng())
    args_maker = lambda: [rng(shape, dtype), np.array(shifts)]
    jnp_op = partial(jnp.roll, axis=axis)
    np_op = partial(np.roll, axis=axis)
    self._CheckAgainstNumpy(np_op, jnp_op, args_maker)
    self._CompileAndCheck(jnp_op, args_maker)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_{}_axis={}_start={}".format(
          jtu.format_shape_dtype_string(shape, dtype),
          axis, start),
       "shape": shape, "dtype": dtype, "axis": axis,
       "start": start}
      for dtype in all_dtypes
      for shape in [(1, 2, 3, 4)]
      for axis in [-3, 0, 2, 3]
      for start in [-4, -1, 2, 4]))
  def testRollaxis(self, shape, dtype, start, axis):
    rng = jtu.rand_default(self.rng())
    args_maker = lambda: [rng(shape, dtype)]
    jnp_op = partial(jnp.rollaxis, axis=axis, start=start)
    np_op = partial(np.rollaxis, axis=axis, start=start)
    self._CheckAgainstNumpy(np_op, jnp_op, args_maker)
    self._CompileAndCheck(jnp_op, args_maker)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_{}_axis={}_bitorder={}".format(
          jtu.format_shape_dtype_string(shape, dtype), axis, bitorder),
       "shape": shape, "dtype": dtype, "axis": axis,
       "bitorder": bitorder}
      for dtype in [np.uint8, np.bool_]
      for bitorder in ['big', 'little']
      for shape in [(1, 2, 3, 4)]
      for axis in [None, 0, 1, -2, -1]))
  def testPackbits(self, shape, dtype, axis, bitorder):
    rng = jtu.rand_some_zero(self.rng())
    args_maker = lambda: [rng(shape, dtype)]
    jnp_op = partial(jnp.packbits, axis=axis, bitorder=bitorder)
    np_op = partial(np.packbits, axis=axis, bitorder=bitorder)
    self._CheckAgainstNumpy(np_op, jnp_op, args_maker)
    self._CompileAndCheck(jnp_op, args_maker)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_{}_axis={}_bitorder={}_count={}".format(
          jtu.format_shape_dtype_string(shape, dtype), axis, bitorder, count),
      "shape": shape, "dtype": dtype, "axis": axis, "bitorder": bitorder,
      "count": count}
      for dtype in [np.uint8]
      for bitorder in ['big', 'little']
      for shape in [(1, 2, 3, 4)]
      for axis in [None, 0, 1, -2, -1]
      for count in [None, 20]))
  def testUnpackbits(self, shape, dtype, axis, bitorder, count):
    rng = jtu.rand_int(self.rng(), 0, 256)
    args_maker = lambda: [rng(shape, dtype)]
    jnp_op = partial(jnp.unpackbits, axis=axis, bitorder=bitorder)
    np_op = partial(np.unpackbits, axis=axis, bitorder=bitorder)
    self._CheckAgainstNumpy(np_op, jnp_op, args_maker)
    self._CompileAndCheck(jnp_op, args_maker)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_{}_index={}_axis={}_mode={}".format(
          jtu.format_shape_dtype_string(shape, dtype),
          jtu.format_shape_dtype_string(index_shape, index_dtype),
          axis, mode),
       "shape": shape, "index_shape": index_shape, "dtype": dtype,
       "index_dtype": index_dtype, "axis": axis, "mode": mode}
      for shape in [(3,), (3, 4), (3, 4, 5)]
      for index_shape in scalar_shapes + [(3,), (2, 1, 3)]
      for axis in itertools.chain(range(-len(shape), len(shape)),
                                  [cast(Optional[int], None)])
      for dtype in all_dtypes
      for index_dtype in int_dtypes
      for mode in [None, 'wrap', 'clip']))
  def testTake(self, shape, dtype, index_shape, index_dtype, axis, mode):
    def args_maker():
      x = rng(shape, dtype)
      i = rng_indices(index_shape, index_dtype)
      return x, i

    rng = jtu.rand_default(self.rng())
    if mode is None:
      rng_indices = jtu.rand_int(self.rng(), -shape[axis or 0], shape[axis or 0])
    else:
      rng_indices = jtu.rand_int(self.rng(), -5, 5)
    jnp_op = lambda x, i: jnp.take(x, i, axis=axis, mode=mode)
    np_op = lambda x, i: np.take(x, i, axis=axis, mode=mode)
    self._CheckAgainstNumpy(np_op, jnp_op, args_maker)
    self._CompileAndCheck(jnp_op, args_maker)

  def testTakeEmpty(self):
    np.testing.assert_array_equal(
      jnp.array([], dtype=jnp.float32),
      jnp.take(jnp.array([], jnp.float32), jnp.array([], jnp.int32)))

    np.testing.assert_array_equal(
      jnp.ones((2, 0, 4), dtype=jnp.float32),
      jnp.take(jnp.ones((2, 0, 4), dtype=jnp.float32), jnp.array([], jnp.int32),
               axis=1))

    with self.assertRaisesRegex(IndexError, "non-empty jnp.take"):
      jnp.take(jnp.ones((2, 0, 4), dtype=jnp.float32),
               jnp.array([0], jnp.int32), axis=1)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_{}_index={}_axis={}".format(
          jtu.format_shape_dtype_string(x_shape, dtype),
          jtu.format_shape_dtype_string(i_shape, index_dtype), axis),
       "x_shape": x_shape, "i_shape": i_shape, "dtype": dtype,
       "index_dtype": index_dtype, "axis": axis}
      for x_shape, i_shape in filter(
        _shapes_are_equal_length,
        filter(_shapes_are_broadcast_compatible,
               itertools.combinations_with_replacement(nonempty_nonscalar_array_shapes, 2)))
      for axis in itertools.chain(range(len(x_shape)), [-1],
                                  [cast(Optional[int], None)])
      for dtype in default_dtypes
      for index_dtype in int_dtypes))
  def testTakeAlongAxis(self, x_shape, i_shape, dtype, index_dtype, axis):
    rng = jtu.rand_default(self.rng())

    i_shape = np.array(i_shape)
    if axis is None:
      i_shape = [np.prod(i_shape, dtype=np.int64)]
    else:
      # Test the case where the size of the axis doesn't necessarily broadcast.
      i_shape[axis] *= 3
      i_shape = list(i_shape)
    def args_maker():
      x = rng(x_shape, dtype)
      n = np.prod(x_shape, dtype=np.int32) if axis is None else x_shape[axis]
      if np.issubdtype(index_dtype, np.unsignedinteger):
        index_rng = jtu.rand_int(self.rng(), 0, n)
      else:
        index_rng = jtu.rand_int(self.rng(), -n, n)
      i = index_rng(i_shape, index_dtype)
      return x, i

    jnp_op = lambda x, i: jnp.take_along_axis(x, i, axis=axis)

    if hasattr(np, "take_along_axis"):
      np_op = lambda x, i: np.take_along_axis(x, i, axis=axis)
      self._CheckAgainstNumpy(np_op, jnp_op, args_maker)
    self._CompileAndCheck(jnp_op, args_maker)

  def testTakeAlongAxisWithUint8IndicesDoesNotOverflow(self):
    # https://github.com/google/jax/issues/5088
    h = jtu.rand_default(self.rng())((256, 256, 100), np.float32)
    g = jtu.rand_int(self.rng(), 0, 100)((256, 256, 1), np.uint8)
    q0 = jnp.take_along_axis(h, g, axis=-1)
    q1 = np.take_along_axis( h, g, axis=-1)
    np.testing.assert_equal(q0, q1)

  def testTakeAlongAxisOutOfBounds(self):
    x = jnp.arange(10, dtype=jnp.float32)
    idx = jnp.array([-11, -10, -9, -5, -1, 0, 1, 5, 9, 10, 11])
    out = jnp.take_along_axis(x, idx, axis=0)
    expected_fill = np.array([jnp.nan, 0, 1, 5, 9, 0, 1, 5, 9, jnp.nan,
                              jnp.nan], np.float32)
    np.testing.assert_array_equal(expected_fill, out)
    out = jnp.take_along_axis(x, idx, axis=0, mode="fill")
    np.testing.assert_array_equal(expected_fill, out)

    expected_clip = np.array([0, 0, 1, 5, 9, 0, 1, 5, 9, 9, 9], np.float32)
    out = jnp.take_along_axis(x, idx, axis=0, mode="clip")
    np.testing.assert_array_equal(expected_clip, out)

  def testTakeAlongAxisRequiresIntIndices(self):
    x = jnp.arange(5)
    idx = jnp.array([3.], jnp.float32)
    with self.assertRaisesRegex(
        TypeError,
        "take_along_axis indices must be of integer type, got float32"):
      jnp.take_along_axis(x, idx, axis=0)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_shape={}_n={}_increasing={}".format(
          jtu.format_shape_dtype_string([shape], dtype),
          n, increasing),
       "dtype": dtype, "shape": shape, "n": n, "increasing": increasing}
      for dtype in inexact_dtypes
      for shape in [0, 5]
      for n in [2, 4]
      for increasing in [False, True]))
  def testVander(self, shape, dtype, n, increasing):
    rng = jtu.rand_default(self.rng())
    def np_fun(arg):
      arg = arg.astype(np.float32) if dtype == jnp.bfloat16 else arg
      return np.vander(arg, N=n, increasing=increasing)
    jnp_fun = lambda arg: jnp.vander(arg, N=n, increasing=increasing)
    args_maker = lambda: [rng([shape], dtype)]
    # np.vander seems to return float64 for all floating types. We could obey
    # those semantics, but they seem like a bug.
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker, check_dtypes=False,
                            tol={np.float32: 1e-3})
    self._CompileAndCheck(jnp_fun, args_maker, check_dtypes=False)

  @parameterized.named_parameters(jtu.cases_from_list(
        {"testcase_name": jtu.format_test_name_suffix(
            "nan_to_num", [shape], [dtype]),
         "shape": shape, "dtype": dtype}
        for shape in array_shapes
        for dtype in inexact_dtypes))
  def testNanToNum(self, shape, dtype):
    rng = jtu.rand_some_inf_and_nan(self.rng())
    dtype = np.dtype(dtypes.canonicalize_dtype(dtype)).type
    def np_fun(x):
      if dtype == jnp.bfloat16:
        x = np.where(np.isnan(x), dtype(0), x)
        x = np.where(np.isposinf(x), jnp.finfo(dtype).max, x)
        x = np.where(np.isneginf(x), jnp.finfo(dtype).min, x)
        return x
      else:
        return np.nan_to_num(x).astype(dtype)

    args_maker = lambda: [rng(shape, dtype)]
    check_dtypes = shape is not jtu.PYTHON_SCALAR_SHAPE
    self._CheckAgainstNumpy(np_fun, jnp.nan_to_num, args_maker,
                            check_dtypes=check_dtypes)
    self._CompileAndCheck(jnp.nan_to_num, args_maker,
                          check_dtypes=check_dtypes)

  @parameterized.named_parameters(jtu.cases_from_list(
        {"testcase_name": jtu.format_test_name_suffix("ix_", shapes, dtypes),
         "shapes": shapes, "dtypes": dtypes}
        for shapes, dtypes in (
          ((), ()),
          (((7,),), (np.int32,)),
          (((3,), (4,)), (np.int32, np.int32)),
          (((3,), (1,), (4,)), (np.int32, np.int32, np.int32)),
        )))
  def testIx_(self, shapes, dtypes):
    rng = jtu.rand_default(self.rng())
    args_maker = lambda: [rng(shape, dtype)
                          for shape, dtype in zip(shapes, dtypes)]
    self._CheckAgainstNumpy(np.ix_, jnp.ix_, args_maker)
    self._CompileAndCheck(jnp.ix_, args_maker)

  @parameterized.named_parameters(
      jtu.cases_from_list(
        {"testcase_name": "_dimensions={}_dtype={}_sparse={}".format(
            dimensions, dtype, sparse),
         "dimensions": dimensions, "dtype": dtype, "sparse": sparse}
        for dimensions in [(), (2,), (3, 0), (4, 5, 6)]
        for dtype in number_dtypes
        for sparse in [True, False]))
  def testIndices(self, dimensions, dtype, sparse):
    def args_maker(): return []
    np_fun = partial(np.indices, dimensions=dimensions,
                     dtype=dtype, sparse=sparse)
    jnp_fun = partial(jnp.indices, dimensions=dimensions,
                      dtype=dtype, sparse=sparse)
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)
    self._CompileAndCheck(jnp_fun, args_maker)

  @parameterized.named_parameters(jtu.cases_from_list(
        {"testcase_name":
           "_op={}_a_shape={}_q_shape={}_axis={}_keepdims={}_method={}".format(
             op,
             jtu.format_shape_dtype_string(a_shape, a_dtype),
             jtu.format_shape_dtype_string(q_shape, q_dtype),
             axis, keepdims, method),
         "a_rng": jtu.rand_some_nan,
         "q_rng": q_rng, "op": op,
         "a_shape": a_shape, "a_dtype": a_dtype,
         "q_shape": q_shape, "q_dtype": q_dtype, "axis": axis,
         "keepdims": keepdims,
         "method": method}
        for (op, q_rng) in (
          ("percentile", partial(jtu.rand_uniform, low=0., high=100.)),
          ("quantile", partial(jtu.rand_uniform, low=0., high=1.)),
          ("nanpercentile", partial(jtu.rand_uniform, low=0., high=100.)),
          ("nanquantile", partial(jtu.rand_uniform, low=0., high=1.)),
        )
        for a_dtype in default_dtypes
        for a_shape, axis in (
          ((7,), None),
          ((47, 7), 0),
          ((47, 7), ()),
          ((4, 101), 1),
          ((4, 47, 7), (1, 2)),
          ((4, 47, 7), (0, 2)),
          ((4, 47, 7), (1, 0, 2)),
        )
        for q_dtype in [np.float32]
        for q_shape in scalar_shapes + [(1,), (4,)]
        for keepdims in [False, True]
        for method in ['linear', 'lower', 'higher', 'nearest', 'midpoint']))
  def testQuantile(self, op, a_rng, q_rng, a_shape, a_dtype, q_shape, q_dtype,
                   axis, keepdims, method):
    a_rng = a_rng(self.rng())
    q_rng = q_rng(self.rng())
    if "median" in op:
      args_maker = lambda: [a_rng(a_shape, a_dtype)]
    else:
      args_maker = lambda: [a_rng(a_shape, a_dtype), q_rng(q_shape, q_dtype)]

    def np_fun(*args):
      args = [x if jnp.result_type(x) != jnp.bfloat16 else
              np.asarray(x, np.float32) for x in args]
      if numpy_version <= (1, 22):
        return getattr(np, op)(*args, axis=axis, keepdims=keepdims,
                               interpolation=method)
      else:
        return getattr(np, op)(*args, axis=axis, keepdims=keepdims,
                               method=method)
    jnp_fun = partial(getattr(jnp, op), axis=axis, keepdims=keepdims,
                      method=method)

    # TODO(phawkins): we currently set dtype=False because we aren't as
    # aggressive about promoting to float64. It's not clear we want to mimic
    # Numpy here.
    tol_spec = {np.float16: 1E-2, np.float32: 2e-4, np.float64: 5e-6}
    tol = max(jtu.tolerance(a_dtype, tol_spec),
              jtu.tolerance(q_dtype, tol_spec))
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker, check_dtypes=False,
                            tol=tol)
    self._CompileAndCheck(jnp_fun, args_maker, rtol=tol)

  @unittest.skipIf(not config.jax_enable_x64, "test requires X64")
  @unittest.skipIf(jtu.device_under_test() != 'cpu', "test is for CPU float64 precision")
  def testPercentilePrecision(self):
    # Regression test for https://github.com/google/jax/issues/8513
    x = jnp.float64([1, 2, 3, 4, 7, 10])
    self.assertEqual(jnp.percentile(x, 50), 3.5)

  @parameterized.named_parameters(jtu.cases_from_list(
        {"testcase_name":
           "_{}_a_shape={}_axis={}_keepdims={}".format(
             op, jtu.format_shape_dtype_string(a_shape, a_dtype),
             axis, keepdims),
         "op": op, "a_shape": a_shape, "a_dtype": a_dtype,
         "axis": axis,
         "keepdims": keepdims}
        for a_dtype in default_dtypes
        for a_shape, axis in (
          ((7,), None),
          ((47, 7), 0),
          ((4, 101), 1),
        )
        for keepdims in [False, True]
        for op in ["median", "nanmedian"]))
  def testMedian(self, op, a_shape, a_dtype, axis, keepdims):
    if op == "median":
      a_rng = jtu.rand_default(self.rng())
    else:
      a_rng = jtu.rand_some_nan(self.rng())
    args_maker = lambda: [a_rng(a_shape, a_dtype)]
    def np_fun(*args):
      args = [x if jnp.result_type(x) != jnp.bfloat16 else
              np.asarray(x, np.float32) for x in args]
      return getattr(np, op)(*args, axis=axis, keepdims=keepdims)
    jnp_fun = partial(getattr(jnp, op), axis=axis, keepdims=keepdims)
    # TODO(phawkins): we currently set dtype=False because we aren't as
    # aggressive about promoting to float64. It's not clear we want to mimic
    # Numpy here.
    tol_spec = {np.float32: 2e-4, np.float64: 5e-6}
    tol = jtu.tolerance(a_dtype, tol_spec)
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker, check_dtypes=False,
                            tol=tol)
    self._CompileAndCheck(jnp_fun, args_maker, rtol=tol)


  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_shape={}".format(
          jtu.format_shape_dtype_string(shape, dtype)),
       "shape": shape, "dtype": dtype}
      for shape in all_shapes for dtype in all_dtypes))
  def testWhereOneArgument(self, shape, dtype):
    rng = jtu.rand_some_zero(self.rng())
    np_fun = lambda x: np.where(x)
    np_fun = jtu.ignore_warning(
      category=DeprecationWarning,
      message="Calling nonzero on 0d arrays.*")(np_fun)
    jnp_fun = lambda x: jnp.where(x)
    args_maker = lambda: [rng(shape, dtype)]
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker, check_dtypes=False)

    # JIT compilation requires specifying a size statically. Full test of
    # this behavior is in testNonzeroSize().
    jnp_fun = lambda x: jnp.where(x, size=np.size(x) // 2)
    self._CompileAndCheck(jnp_fun, args_maker)

  @parameterized.named_parameters(jtu.named_cases_from_sampler(lambda s: ({
      "testcase_name": "_{}".format("_".join(
        jtu.format_shape_dtype_string(shape, dtype)
        for shape, dtype in zip(shapes, dtypes))),
      "shapes": shapes, "dtypes": dtypes
    } for shapes in s(filter(_shapes_are_broadcast_compatible,
                         itertools.combinations_with_replacement(all_shapes, 3)))
      for dtypes in s(itertools.combinations_with_replacement(all_dtypes, 3)))))
  def testWhereThreeArgument(self, shapes, dtypes):
    rng = jtu.rand_default(self.rng())
    args_maker = self._GetArgsMaker(rng, shapes, dtypes)
    def np_fun(cond, x, y):
      return _promote_like_jnp(partial(np.where, cond))(x, y)
    self._CheckAgainstNumpy(np_fun, jnp.where, args_maker)
    self._CompileAndCheck(jnp.where, args_maker)

  def testWhereScalarPromotion(self):
    x = jnp.where(jnp.array([True, False]), 3,
                  jnp.ones((2,), dtype=jnp.float32))
    self.assertEqual(x.dtype, np.dtype(np.float32))

  @parameterized.named_parameters(jtu.named_cases_from_sampler(lambda s: ({
      "testcase_name": jtu.format_test_name_suffix("", shapes, (np.bool_,) * n + dtypes),
      "shapes": shapes, "dtypes": dtypes
    } for n in s(range(1, 3))
      for shapes in s(filter(
          _shapes_are_broadcast_compatible,
          itertools.combinations_with_replacement(all_shapes, 2 * n + 1)))
      for dtypes in s(itertools.combinations_with_replacement(all_dtypes, n + 1)))))
  def testSelect(self, shapes, dtypes):
    rng = jtu.rand_default(self.rng())
    n = len(dtypes) - 1
    def args_maker():
      condlist = [rng(shape, np.bool_) for shape in shapes[:n]]
      choicelist = [rng(shape, dtype)
                    for shape, dtype in zip(shapes[n:-1], dtypes[:n])]
      default = rng(shapes[-1], dtypes[-1])
      return condlist, choicelist, default
    # TODO(phawkins): float32/float64 type mismatches
    def np_fun(condlist, choicelist, default):
      choicelist = [x if jnp.result_type(x) != jnp.bfloat16
                    else x.astype(np.float32) for x in choicelist]
      dtype = jnp.result_type(default, *choicelist)
      return np.select(condlist,
                        [np.asarray(x, dtype=dtype) for x in choicelist],
                        np.asarray(default, dtype=dtype))
    self._CheckAgainstNumpy(np_fun, jnp.select, args_maker,
                            check_dtypes=False)
    self._CompileAndCheck(jnp.select, args_maker,
                          rtol={np.float64: 1e-7, np.complex128: 1e-7})


  def testIssue330(self):
    x = jnp.full((1, 1), jnp.array([1])[0])  # doesn't crash
    self.assertEqual(x[0, 0], 1)

  def testScalarDtypePromotion(self):
    orig_numpy_result = (1 + np.eye(1, dtype=np.float32)).dtype
    jax_numpy_result = (1 + jnp.eye(1, dtype=jnp.float32)).dtype
    self.assertEqual(orig_numpy_result, jax_numpy_result)

  def testSymmetrizeDtypePromotion(self):
    x = np.eye(3, dtype=np.float32)
    orig_numpy_result = ((x + x.T) / 2).dtype

    x = jnp.eye(3, dtype=jnp.float32)
    jax_numpy_result = ((x + x.T) / 2).dtype
    self.assertEqual(orig_numpy_result, jax_numpy_result)

  # NOTE(mattjj): I disabled this test when removing lax._safe_mul because
  # introducing the convention 0 * inf = 0 leads to silently wrong results in
  # some cases. See this comment for details:
  # https://github.com/google/jax/issues/1052#issuecomment-514083352
  # def testIssue347(self):
  #   # https://github.com/google/jax/issues/347
  #   def test_fail(x):
  #     x = jnp.sqrt(jnp.sum(x ** 2, axis=1))
  #     ones = jnp.ones_like(x)
  #     x = jnp.where(x > 0.5, x, ones)
  #     return jnp.sum(x)
  #   x = jnp.array([[1, 2], [3, 4], [0, 0]], dtype=jnp.float64)
  #   result = jax.grad(test_fail)(x)
  #   assert not np.any(np.isnan(result))

  def testIssue453(self):
    # https://github.com/google/jax/issues/453
    a = np.arange(6) + 1
    ans = jnp.reshape(a, (3, 2), order='F')
    expected = np.reshape(a, (3, 2), order='F')
    self.assertAllClose(ans, expected)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": f"_op={op}_dtype={dtype.__name__}",
       "dtype": dtype, "op": op}
      for dtype in [int, float, bool, complex]
      for op in ["atleast_1d", "atleast_2d", "atleast_3d"]))
  def testAtLeastNdLiterals(self, dtype, op):
    # Fixes: https://github.com/google/jax/issues/634
    np_fun = lambda arg: getattr(np, op)(arg).astype(dtypes.python_scalar_dtypes[dtype])
    jnp_fun = lambda arg: getattr(jnp, op)(arg)
    args_maker = lambda: [dtype(2)]
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)
    self._CompileAndCheck(jnp_fun, args_maker)

  @parameterized.named_parameters(jtu.cases_from_list(
    {
      "testcase_name": "_shape={}_dtype={}_weights={}_minlength={}_length={}".format(
        shape, dtype, weights, minlength, length
      ),
      "shape": shape,
      "dtype": dtype,
      "weights": weights,
      "minlength": minlength,
      "length": length}
    for shape in [(0,), (5,), (10,)]
    for dtype in int_dtypes
    for weights in [True, False]
    for minlength in [0, 20]
    for length in [None, 8]
  ))
  def testBincount(self, shape, dtype, weights, minlength, length):
    rng = jtu.rand_default(self.rng())
    args_maker = lambda: (rng(shape, dtype), (rng(shape, 'float32') if weights else None))

    def np_fun(x, *args):
      x = np.clip(x, 0, None)  # jnp.bincount clips negative values to zero.
      out = np.bincount(x, *args, minlength=minlength)
      if length and length > out.size:
        return np.pad(out, (0, length - out.size))
      return out[:length]
    jnp_fun = partial(jnp.bincount, minlength=minlength, length=length)

    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker, check_dtypes=False)
    if length is not None:
      self._CompileAndCheck(jnp_fun, args_maker)

  def testBincountNegative(self):
    # Test that jnp.bincount ignores negative values.
    x_rng = jtu.rand_int(self.rng(), -100, 100)
    w_rng = jtu.rand_uniform(self.rng())
    shape = (1000,)
    x = x_rng(shape, 'int32')
    w = w_rng(shape, 'float32')

    xn = np.array(x)
    xn[xn < 0] = 0
    wn = np.array(w)
    np_result = np.bincount(xn[xn >= 0], wn[xn >= 0])
    jnp_result = jnp.bincount(x, w)
    self.assertAllClose(np_result, jnp_result, check_dtypes=False)


  @parameterized.named_parameters(*jtu.cases_from_list(
      {"testcase_name": f"_case={i}",
       "input": input}
       for i, input in enumerate([
         3,
         [3],
         [np.array(3)],
         [np.array([3])],
         [[np.array(3)]],
         [[np.array([3])]],
         [3, 4, 5],
         [
           [np.eye(2, dtype=np.int32) * 2, np.zeros((2, 3), dtype=np.int32)],
           [np.ones((3, 2), dtype=np.int32), np.eye(3, dtype=np.int32) * 3],
         ],
         [np.array([1, 2, 3]), np.array([2, 3, 4]), 10],
         [np.ones((2, 2), dtype=np.int32), np.zeros((2, 2), dtype=np.int32)],
         [[np.array([1, 2, 3])], [np.array([2, 3, 4])]],
       ])))
  def testBlock(self, input):
    args_maker = lambda: [input]
    self._CheckAgainstNumpy(np.block, jnp.block, args_maker)
    self._CompileAndCheck(jnp.block, args_maker)

  def testLongLong(self):
    self.assertAllClose(np.int64(7), jax.jit(lambda x: x)(np.longlong(7)))

  @jtu.ignore_warning(category=UserWarning,
                      message="Explicitly requested dtype.*")
  def testArange(self):
    # test cases inspired by dask tests at
    # https://github.com/dask/dask/blob/main/dask/array/tests/test_creation.py#L92
    np_arange = jtu.with_jax_dtype_defaults(np.arange)
    self.assertAllClose(jnp.arange(77),
                        np_arange(77))
    self.assertAllClose(jnp.arange(2, 13),
                        np_arange(2, 13))
    self.assertAllClose(jnp.arange(4, 21, 9),
                        np_arange(4, 21, 9))
    self.assertAllClose(jnp.arange(53, 5, -3),
                        np_arange(53, 5, -3))
    self.assertAllClose(jnp.arange(77, dtype=float),
                        np_arange(77, dtype=float))
    self.assertAllClose(jnp.arange(2, 13, dtype=int),
                        np_arange(2, 13, dtype=int))
    self.assertAllClose(jnp.arange(0, 1, -0.5),
                        np_arange(0, 1, -0.5))

    self.assertRaises(TypeError, lambda: jnp.arange())

    # test that jnp.arange(N) doesn't instantiate an ndarray
    self.assertNotEqual(type(jnp.arange(77)), type(np.arange(77)))
    self.assertEqual(type(jnp.arange(77)), type(lax.iota(np.int32, 77)))

    # test that jnp.arange(N, dtype=int32) doesn't instantiate an ndarray
    self.assertNotEqual(type(jnp.arange(77, dtype=jnp.int32)),
                        type(np.arange(77, dtype=np.int32)))
    self.assertEqual(type(jnp.arange(77, dtype=jnp.int32)),
                     type(lax.iota(np.int32, 77)))

  def testArangeJit(self):
    ans = jax.jit(lambda: jnp.arange(5))()
    expected = jtu.with_jax_dtype_defaults(np.arange)(5)
    self.assertAllClose(ans, expected)

  @parameterized.named_parameters(
    {"testcase_name": f"_{args}", "args": args} for args in [(5,), (0, 5)])
  def testArangeJaxpr(self, args):
    jaxpr = jax.make_jaxpr(lambda: jnp.arange(*args))()
    self.assertEqual(len(jaxpr.jaxpr.eqns), 1)
    self.assertEqual(jaxpr.jaxpr.eqns[0].primitive, lax.iota_p)

  def testIssue830(self):
    a = jnp.arange(4, dtype=jnp.complex64)
    self.assertEqual(a.dtype, jnp.complex64)

  def testIssue728(self):
    assert jnp.allclose(jnp.eye(5000), np.eye(5000))
    self.assertEqual(0, np.sum(jnp.eye(1050) - np.eye(1050)))

  def testIssue746(self):
    jnp.arange(12).reshape(3, 4)  # doesn't crash

  def testIssue764(self):
    x = jnp.linspace(190, 200, 4)
    f = jax.grad(lambda x: jnp.sum(jnp.tanh(x)))
    # Expected values computed with autograd in float64 precision.
    expected = np.array([3.71669453e-165, 4.72999108e-168, 6.01954653e-171,
                          7.66067839e-174], np.float64)
    self.assertAllClose(f(x), expected, check_dtypes=False)

  def testIssue776(self):
    """Tests that the scatter-add transpose rule instantiates symbolic zeros."""
    def f(u):
      y = jnp.ones(10).at[np.array([2, 4, 5])].add(u)
      # The transpose rule for lax.tie_in returns a symbolic zero for its first
      # argument.
      return lax.tie_in(y, 7.)

    self.assertAllClose(np.zeros(3,), jax.grad(f)(np.ones(3,)))

  # NOTE(mattjj): I disabled this test when removing lax._safe_mul because this
  # is a numerical stability issue that should be solved with a custom jvp rule
  # of the sigmoid function being differentiated here, not by safe_mul.
  # def testIssue777(self):
  #   x = jnp.linspace(-200, 0, 4, dtype=np.float32)
  #   f = jax.grad(lambda x: jnp.sum(1 / (1 + jnp.exp(-x))))
  #   self.assertAllClose(f(x), np.array([0., 0., 0., 0.25], dtype=np.float32))

  @parameterized.named_parameters(
      jtu.cases_from_list(
        {"testcase_name": jtu.format_test_name_suffix(op, [()], [dtype]),
         "dtype": dtype, "op": op}
      for dtype in float_dtypes
      for op in ("sqrt", "arccos", "arcsin", "arctan", "sin", "cos", "tan",
                 "sinh", "cosh", "tanh", "arccosh", "arcsinh", "arctanh", "exp",
                 "log", "expm1", "log1p")))
  def testMathSpecialFloatValues(self, op, dtype):
    np_op = getattr(np, op)
    np_op = jtu.ignore_warning(category=RuntimeWarning,
                                 message="invalid value.*")(np_op)
    np_op = jtu.ignore_warning(category=RuntimeWarning,
                                 message="divide by zero.*")(np_op)
    np_op = jtu.ignore_warning(category=RuntimeWarning,
                                 message="overflow.*")(np_op)

    jnp_op = getattr(jnp, op)
    dtype = np.dtype(dtypes.canonicalize_dtype(dtype)).type
    for x in (np.nan, -np.inf, -100., -2., -1., 0., 1., 2., 100., np.inf,
              jnp.finfo(dtype).max, np.sqrt(jnp.finfo(dtype).max),
              np.sqrt(jnp.finfo(dtype).max) * 2.):
      if (op in ("sin", "cos", "tan") and
          jtu.device_under_test() == "tpu"):
        continue  # TODO(b/132196789): fix and reenable.
      x = dtype(x)
      expected = np_op(x)
      actual = jnp_op(x)
      tol = jtu.tolerance(dtype, {np.float32: 1e-3, np.float64: 1e-7})
      self.assertAllClose(expected, actual, atol=tol,
                          rtol=tol)


  def testReductionOfOutOfBoundsAxis(self):  # Issue 888
    x = jnp.ones((3, 4))
    self.assertRaises(ValueError, lambda: jnp.sum(x, axis=2))

  def testIssue956(self):
    self.assertRaises(TypeError, lambda: jnp.ndarray((1, 1)))

  @parameterized.named_parameters(
      jtu.cases_from_list(
        {"testcase_name":
         "_shape={}_dtype={}_out_dtype={}_axis={}_ddof={}_keepdims={}"
         .format(shape, dtype.__name__, out_dtype.__name__, axis, ddof, keepdims),
         "shape": shape, "dtype": dtype, "out_dtype": out_dtype, "axis": axis,
         "ddof": ddof, "keepdims": keepdims}
        for shape in [(5,), (10, 5)]
        for dtype in all_dtypes
        for out_dtype in inexact_dtypes
        for axis in [None, 0, -1]
        for ddof in [0, 1, 2]
        for keepdims in [False, True]))
  def testVar(self, shape, dtype, out_dtype, axis, ddof, keepdims):
    rng = jtu.rand_default(self.rng())
    args_maker = self._GetArgsMaker(rng, [shape], [dtype])
    @jtu.ignore_warning(category=RuntimeWarning,
                        message="Degrees of freedom <= 0 for slice.")
    def np_fun(x):
      out = np.var(x.astype(jnp.promote_types(np.float32, dtype)),
                    axis=axis, ddof=ddof, keepdims=keepdims)
      return out.astype(out_dtype)
    jnp_fun = partial(jnp.var, dtype=out_dtype, axis=axis, ddof=ddof, keepdims=keepdims)
    tol = jtu.tolerance(out_dtype, {np.float16: 1e-1, np.float32: 1e-3,
                                    np.float64: 1e-3, np.complex128: 1e-6})
    if (jnp.issubdtype(dtype, jnp.complexfloating) and
        not jnp.issubdtype(out_dtype, jnp.complexfloating)):
      self.assertRaises(ValueError, lambda: jnp_fun(*args_maker()))
    else:
      self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker,
                              tol=tol)
      self._CompileAndCheck(jnp_fun, args_maker, rtol=tol,
                            atol=tol)

  @parameterized.named_parameters(
      jtu.cases_from_list(
        {"testcase_name":
         "_shape={}_dtype={}_out_dtype={}_axis={}_ddof={}_keepdims={}"
         .format(shape, dtype, out_dtype, axis, ddof, keepdims),
         "shape": shape, "dtype": dtype, "out_dtype": out_dtype, "axis": axis,
         "ddof": ddof, "keepdims": keepdims}
        for shape in [(5,), (10, 5)]
        for dtype in all_dtypes
        for out_dtype in inexact_dtypes
        for axis in [None, 0, -1]
        for ddof in [0, 1, 2]
        for keepdims in [False, True]))
  def testNanVar(self, shape, dtype, out_dtype, axis, ddof, keepdims):
    rng = jtu.rand_some_nan(self.rng())
    args_maker = self._GetArgsMaker(rng, [shape], [dtype])
    @jtu.ignore_warning(category=RuntimeWarning,
                        message="Degrees of freedom <= 0 for slice.")
    def np_fun(x):
      out = np.nanvar(x.astype(jnp.promote_types(np.float32, dtype)),
                    axis=axis, ddof=ddof, keepdims=keepdims)
      return out.astype(out_dtype)
    jnp_fun = partial(jnp.nanvar, dtype=out_dtype, axis=axis, ddof=ddof, keepdims=keepdims)
    tol = jtu.tolerance(out_dtype, {np.float16: 1e-1, np.float32: 1e-3,
                                    np.float64: 1e-3, np.complex128: 1e-6})
    if (jnp.issubdtype(dtype, jnp.complexfloating) and
        not jnp.issubdtype(out_dtype, jnp.complexfloating)):
      self.assertRaises(ValueError, lambda: jnp_fun(*args_maker()))
    else:
      self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker,
                              tol=tol)
      self._CompileAndCheck(jnp_fun, args_maker, rtol=tol,
                            atol=tol)

  def testNanStdGrad(self):
    # Regression test for https://github.com/google/jax/issues/8128
    x = jnp.arange(5.0).at[0].set(jnp.nan)
    y = jax.grad(jnp.nanvar)(x)
    self.assertAllClose(y, jnp.array([0.0, -0.75, -0.25, 0.25, 0.75]))

    z = jax.grad(jnp.nanstd)(x)
    self.assertEqual(jnp.isnan(z).sum(), 0)

  @parameterized.named_parameters(
      jtu.cases_from_list(
        {"testcase_name":
          "_shape={}_dtype={}_y_shape={}_y_dtype={}_rowvar={}_ddof={}_bias={}_fweights={}_aweights={}".format(
            shape, dtype, y_shape, y_dtype, rowvar, ddof, bias, fweights, aweights),
         "shape": shape, "y_shape": y_shape, "dtype": dtype, "y_dtype": y_dtype,"rowvar": rowvar, "ddof": ddof,
         "bias": bias, "fweights": fweights, "aweights": aweights}
        for shape in [(5,), (10, 5), (5, 10)]
        for dtype in all_dtypes
        for y_dtype in [None, dtype]
        for rowvar in [True, False]
        for y_shape in _get_y_shapes(y_dtype, shape, rowvar)
        for bias in [True, False]
        for ddof in [None, 2, 3]
        for fweights in [True, False]
        for aweights in [True, False]))
  def testCov(self, shape, dtype, y_shape, y_dtype, rowvar, ddof, bias, fweights, aweights):
    rng = jtu.rand_default(self.rng())
    wrng = jtu.rand_positive(self.rng())
    wdtype = np.real(dtype(0)).dtype
    wshape = shape[-1:] if rowvar or shape[0] == 1 else shape[:1]

    args_maker = lambda: [rng(shape, dtype),
                          rng(y_shape, y_dtype) if y_dtype else None,
                          wrng(wshape, int) if fweights else None,
                          wrng(wshape, wdtype) if aweights else None]
    kwargs = dict(rowvar=rowvar, ddof=ddof, bias=bias)
    np_fun = lambda m, y, f, a: np.cov(m, y, fweights=f, aweights=a, **kwargs)
    jnp_fun = lambda m, y, f, a: jnp.cov(m, y, fweights=f, aweights=a, **kwargs)
    tol = {jnp.bfloat16: 5E-2, np.float16: 1E-2, np.float32: 1e-5,
           np.float64: 1e-13, np.complex64: 1e-5, np.complex128: 1e-13}
    tol = 7e-2 if jtu.device_under_test() == "tpu" else tol
    tol = jtu.join_tolerance(tol, jtu.tolerance(dtype))
    self._CheckAgainstNumpy(
        np_fun, jnp_fun, args_maker, check_dtypes=False, tol=tol)
    self._CompileAndCheck(jnp_fun, args_maker, atol=tol,
                          rtol=tol)

  def testIssue967(self):
    self.assertRaises(TypeError, lambda: jnp.zeros(1.5))

  @parameterized.named_parameters(
      jtu.cases_from_list(
        {"testcase_name": "_shape={}_dtype={}_rowvar={}".format(
            shape, dtype.__name__, rowvar),
         "shape": shape, "dtype": dtype, "rowvar": rowvar}
        for shape in [(5,), (10, 5), (3, 10)]
        for dtype in number_dtypes
        for rowvar in [True, False]))
  def testCorrCoef(self, shape, dtype, rowvar):
    rng = jtu.rand_default(self.rng())
    def args_maker():
      ok = False
      while not ok:
        x = rng(shape, dtype)
        ok = not np.any(np.isclose(np.std(x), 0.0))
      return (x,)
    np_fun = partial(np.corrcoef, rowvar=rowvar)
    np_fun = jtu.ignore_warning(
      category=RuntimeWarning, message="invalid value encountered.*")(np_fun)
    jnp_fun = partial(jnp.corrcoef, rowvar=rowvar)
    tol = 1e-2 if jtu.device_under_test() == "tpu" else None
    self._CheckAgainstNumpy(
        np_fun, jnp_fun, args_maker, check_dtypes=False,
        tol=tol)
    self._CompileAndCheck(jnp_fun, args_maker, atol=tol, rtol=tol)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_{}_{}_{}".format(jtu.format_shape_dtype_string(shape, dtype),
       "None" if end_dtype is None else jtu.format_shape_dtype_string(end_shape, end_dtype),
       "None" if begin_dtype is None else jtu.format_shape_dtype_string(begin_shape, begin_dtype)),
       "shape": shape, "dtype": dtype, "end_shape": end_shape,
       "end_dtype": end_dtype, "begin_shape": begin_shape,
       "begin_dtype": begin_dtype}
      for dtype in number_dtypes
      for end_dtype in [None] + [dtype]
      for begin_dtype in [None] + [dtype]
      for shape in [s for s in all_shapes if s != jtu.PYTHON_SCALAR_SHAPE]
      for begin_shape in (
        [None] if begin_dtype is None
        else [s for s in all_shapes if s != jtu.PYTHON_SCALAR_SHAPE])
      for end_shape in (
        [None] if end_dtype is None
        else [s for s in all_shapes if s != jtu.PYTHON_SCALAR_SHAPE])))
  def testEDiff1d(self, shape, dtype, end_shape, end_dtype, begin_shape,
                  begin_dtype):
    rng = jtu.rand_default(self.rng())
    args_maker = lambda: [rng(shape, dtype),
            (None if end_dtype is None else rng(end_shape, end_dtype)),
            (None if begin_dtype is None else rng(begin_shape, begin_dtype))]
    np_fun = lambda x, to_end, to_begin: np.ediff1d(x, to_end, to_begin)
    jnp_fun = lambda x, to_end, to_begin: jnp.ediff1d(x, to_end, to_begin)
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)
    self._CompileAndCheck(jnp_fun, args_maker)

  def testEDiff1dWithDtypeCast(self):
    rng = jtu.rand_default(self.rng())
    shape = jtu.NUMPY_SCALAR_SHAPE
    dtype = jnp.float32
    end_dtype = jnp.int32
    args_maker = lambda: [rng(shape, dtype), rng(shape, end_dtype), rng(shape, dtype)]
    np_fun = lambda x, to_end, to_begin: np.ediff1d(x, to_end, to_begin)
    jnp_fun = lambda x, to_end, to_begin: jnp.ediff1d(x, to_end, to_begin)
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)
    self._CompileAndCheck(jnp_fun, args_maker)

  @parameterized.named_parameters(
      jtu.cases_from_list(
        {"testcase_name": "_shapes={}_dtype={}_indexing={}_sparse={}".format(
            shapes, dtype, indexing, sparse),
         "shapes": shapes, "dtype": dtype, "indexing": indexing,
         "sparse": sparse}
        for shapes in [(), (5,), (5, 3)]
        for dtype in number_dtypes
        for indexing in ['xy', 'ij']
        for sparse in [True, False]))
  def testMeshGrid(self, shapes, dtype, indexing, sparse):
    rng = jtu.rand_default(self.rng())
    args_maker = self._GetArgsMaker(rng, [(x,) for x in shapes],
                                    [dtype] * len(shapes))
    np_fun = partial(np.meshgrid, indexing=indexing, sparse=sparse)
    jnp_fun = partial(jnp.meshgrid, indexing=indexing, sparse=sparse)
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)
    self._CompileAndCheck(jnp_fun, args_maker)

  def testMgrid(self):
    # wrap indexer for appropriate dtype defaults.
    np_mgrid = _indexer_with_default_outputs(np.mgrid)
    assertAllEqual = partial(self.assertAllClose, atol=0, rtol=0)
    assertAllEqual(np_mgrid[:4], jnp.mgrid[:4])
    assertAllEqual(np_mgrid[:4,], jnp.mgrid[:4,])
    assertAllEqual(np_mgrid[:4], jax.jit(lambda: jnp.mgrid[:4])())
    assertAllEqual(np_mgrid[:5, :5], jnp.mgrid[:5, :5])
    assertAllEqual(np_mgrid[:3, :2], jnp.mgrid[:3, :2])
    assertAllEqual(np_mgrid[1:4:2], jnp.mgrid[1:4:2])
    assertAllEqual(np_mgrid[1:5:3, :5], jnp.mgrid[1:5:3, :5])
    assertAllEqual(np_mgrid[:3, :2, :5], jnp.mgrid[:3, :2, :5])
    assertAllEqual(np_mgrid[:3:2, :2, :5], jnp.mgrid[:3:2, :2, :5])
    # Corner cases
    assertAllEqual(np_mgrid[:], jnp.mgrid[:])
    # When the step length is a complex number, because of float calculation,
    # the values between jnp and np might slightly different.
    atol = 1e-6
    rtol = 1e-6
    self.assertAllClose(np_mgrid[-1:1:5j],
                        jnp.mgrid[-1:1:5j],
                        atol=atol,
                        rtol=rtol)
    self.assertAllClose(np_mgrid[3:4:7j],
                        jnp.mgrid[3:4:7j],
                        atol=atol,
                        rtol=rtol)
    self.assertAllClose(np_mgrid[1:6:8j, 2:4],
                        jnp.mgrid[1:6:8j, 2:4],
                        atol=atol,
                        rtol=rtol)
    # Non-integer steps
    self.assertAllClose(np_mgrid[0:3.5:0.5],
                        jnp.mgrid[0:3.5:0.5],
                        atol=atol,
                        rtol=rtol)
    self.assertAllClose(np_mgrid[1.3:4.2:0.3],
                        jnp.mgrid[1.3:4.2:0.3],
                        atol=atol,
                        rtol=rtol)
    # abstract tracer value for jnp.mgrid slice
    with self.assertRaisesRegex(jax.core.ConcretizationTypeError,
                                "slice start of jnp.mgrid"):
      jax.jit(lambda a, b: jnp.mgrid[a:b])(0, 2)

  def testOgrid(self):
    # wrap indexer for appropriate dtype defaults.
    np_ogrid = _indexer_with_default_outputs(np.ogrid)
    def assertListOfArraysEqual(xs, ys):
      self.assertIsInstance(xs, list)
      self.assertIsInstance(ys, list)
      self.assertEqual(len(xs), len(ys))
      for x, y in zip(xs, ys):
        self.assertArraysEqual(x, y)

    self.assertArraysEqual(np_ogrid[:5], jnp.ogrid[:5])
    self.assertArraysEqual(np_ogrid[:5], jax.jit(lambda: jnp.ogrid[:5])())
    self.assertArraysEqual(np_ogrid[1:7:2], jnp.ogrid[1:7:2])
    # List of arrays
    assertListOfArraysEqual(np_ogrid[:5,], jnp.ogrid[:5,])
    assertListOfArraysEqual(np_ogrid[0:5, 1:3], jnp.ogrid[0:5, 1:3])
    assertListOfArraysEqual(np_ogrid[1:3:2, 2:9:3], jnp.ogrid[1:3:2, 2:9:3])
    assertListOfArraysEqual(np_ogrid[:5, :9, :11], jnp.ogrid[:5, :9, :11])
    # Corner cases
    self.assertArraysEqual(np_ogrid[:], jnp.ogrid[:])
    # Complex number steps
    atol = 1e-6
    rtol = 1e-6
    self.assertAllClose(np_ogrid[-1:1:5j],
                        jnp.ogrid[-1:1:5j],
                        atol=atol,
                        rtol=rtol)
    # Non-integer steps
    self.assertAllClose(np_ogrid[0:3.5:0.3],
                        jnp.ogrid[0:3.5:0.3],
                        atol=atol,
                        rtol=rtol)
    self.assertAllClose(np_ogrid[1.2:4.8:0.24],
                        jnp.ogrid[1.2:4.8:0.24],
                        atol=atol,
                        rtol=rtol)
    # abstract tracer value for ogrid slice
    with self.assertRaisesRegex(jax.core.ConcretizationTypeError,
                                "slice start of jnp.ogrid"):
      jax.jit(lambda a, b: jnp.ogrid[a:b])(0, 2)

  def testR_(self):
    a = np.arange(6).reshape((2,3))
    self.assertArraysEqual(np.r_[np.array([1,2,3]), 0, 0, np.array([4,5,6])],
                           jnp.r_[np.array([1,2,3]), 0, 0, np.array([4,5,6])])
    self.assertArraysEqual(np.r_['-1', a, a], jnp.r_['-1', a, a])

    # wrap indexer for appropriate dtype defaults.
    np_r_ = _indexer_with_default_outputs(np.r_)
    self.assertArraysEqual(np_r_['0,2', [1,2,3], [4,5,6]], jnp.r_['0,2', [1,2,3], [4,5,6]])
    self.assertArraysEqual(np_r_['0,2,0', [1,2,3], [4,5,6]], jnp.r_['0,2,0', [1,2,3], [4,5,6]])
    self.assertArraysEqual(np_r_['1,2,0', [1,2,3], [4,5,6]], jnp.r_['1,2,0', [1,2,3], [4,5,6]])
    # negative 1d axis start
    self.assertArraysEqual(np_r_['0,4,-1', [1,2,3], [4,5,6]], jnp.r_['0,4,-1', [1,2,3], [4,5,6]])
    self.assertArraysEqual(np_r_['0,4,-2', [1,2,3], [4,5,6]], jnp.r_['0,4,-2', [1,2,3], [4,5,6]])

    # matrix directives
    with warnings.catch_warnings():
      warnings.filterwarnings("ignore", category=PendingDeprecationWarning)
      self.assertArraysEqual(np_r_['r',[1,2,3], [4,5,6]], jnp.r_['r',[1,2,3], [4,5,6]])
      self.assertArraysEqual(np_r_['c', [1, 2, 3], [4, 5, 6]], jnp.r_['c', [1, 2, 3], [4, 5, 6]])

    # bad directive
    with self.assertRaisesRegex(ValueError, "could not understand directive.*"):
      jnp.r_["asdfgh",[1,2,3]]
    # abstract tracer value for r_ slice
    with self.assertRaisesRegex(jax.core.ConcretizationTypeError,
                                "slice start of jnp.r_"):
      jax.jit(lambda a, b: jnp.r_[a:b])(0, 2)

    # Complex number steps
    atol = 1e-6
    rtol = 1e-6
    self.assertAllClose(np_r_[-1:1:6j],
                        jnp.r_[-1:1:6j],
                        atol=atol,
                        rtol=rtol)
    self.assertAllClose(np_r_[-1:1:6j, [0]*3, 5, 6],
                        jnp.r_[-1:1:6j, [0]*3, 5, 6],
                        atol=atol,
                        rtol=rtol)
    # Non-integer steps
    self.assertAllClose(np_r_[1.2:4.8:0.24],
                        jnp.r_[1.2:4.8:0.24],
                        atol=atol,
                        rtol=rtol)

  def testC_(self):
    a = np.arange(6).reshape((2, 3))
    self.assertArraysEqual(np.c_[np.array([1,2,3]), np.array([4,5,6])],
                           jnp.c_[np.array([1,2,3]), np.array([4,5,6])])
    self.assertArraysEqual(np.c_[np.array([[1,2,3]]), 0, 0, np.array([[4,5,6]])],
                           jnp.c_[np.array([[1,2,3]]), 0, 0, np.array([[4,5,6]])])
    self.assertArraysEqual(np.c_['-1', a, a], jnp.c_['-1', a, a])

    # wrap indexer for appropriate dtype defaults.
    np_c_ = _indexer_with_default_outputs(np.c_)
    self.assertArraysEqual(np_c_['0,2', [1,2,3], [4,5,6]], jnp.c_['0,2', [1,2,3], [4,5,6]])
    self.assertArraysEqual(np_c_['0,2,0', [1,2,3], [4,5,6]], jnp.c_['0,2,0', [1,2,3], [4,5,6]])
    self.assertArraysEqual(np_c_['1,2,0', [1,2,3], [4,5,6]], jnp.c_['1,2,0', [1,2,3], [4,5,6]])
    # negative 1d axis start
    self.assertArraysEqual(np_c_['0,4,-1', [1,2,3], [4,5,6]], jnp.c_['0,4,-1', [1,2,3], [4,5,6]])
    self.assertArraysEqual(np_c_['0,4,-2', [1,2,3], [4,5,6]], jnp.c_['0,4,-2', [1,2,3], [4,5,6]])
    # matrix directives, avoid numpy deprecation warning
    with warnings.catch_warnings():
      warnings.filterwarnings("ignore", category=PendingDeprecationWarning)
      self.assertArraysEqual(np_c_['r',[1,2,3], [4,5,6]], jnp.c_['r',[1,2,3], [4,5,6]])
      self.assertArraysEqual(np_c_['c', [1, 2, 3], [4, 5, 6]], jnp.c_['c', [1, 2, 3], [4, 5, 6]])

    # bad directive
    with self.assertRaisesRegex(ValueError, "could not understand directive.*"):
      jnp.c_["asdfgh",[1,2,3]]
    # abstract tracer value for c_ slice
    with self.assertRaisesRegex(jax.core.ConcretizationTypeError,
                                "slice start of jnp.c_"):
      jax.jit(lambda a, b: jnp.c_[a:b])(0, 2)

    # Complex number steps
    atol = 1e-6
    rtol = 1e-6
    self.assertAllClose(np_c_[-1:1:6j],
                        jnp.c_[-1:1:6j],
                        atol=atol,
                        rtol=rtol)

    # Non-integer steps
    self.assertAllClose(np_c_[1.2:4.8:0.24],
                        jnp.c_[1.2:4.8:0.24],
                        atol=atol,
                        rtol=rtol)

  def testS_(self):
    self.assertEqual(np.s_[1:2:20],jnp.s_[1:2:20])

  def testIndex_exp(self):
    self.assertEqual(np.index_exp[5:3:2j],jnp.index_exp[5:3:2j])

  @parameterized.named_parameters(
      jtu.cases_from_list(
        {"testcase_name": f"_start_shape={start_shape}_stop_shape={stop_shape}"
                          f"_num={num}_endpoint={endpoint}_retstep={retstep}"
                          f"_dtype={dtype.__name__ if dtype else 'None'}",
         "start_shape": start_shape, "stop_shape": stop_shape,
         "num": num, "endpoint": endpoint, "retstep": retstep,
         "dtype": dtype}
        for start_shape in [(), (2,), (2, 2)]
        for stop_shape in [(), (2,), (2, 2)]
        for num in [0, 1, 2, 5, 20]
        for endpoint in [True, False]
        for retstep in [True, False]
        # floating-point compute between jitted platforms and non-jit + rounding
        # cause unavoidable variation in integer truncation for some inputs, so
        # we currently only test inexact 'dtype' arguments.
        for dtype in inexact_dtypes + [None,]))
  @jax.numpy_rank_promotion('allow')  # This test explicitly exercises implicit rank promotion.
  def testLinspace(self, start_shape, stop_shape, num, endpoint, retstep, dtype):
    rng = jtu.rand_default(self.rng())
    # relax default tolerances slightly
    tol = jtu.tolerance(dtype if dtype else np.float32) * 10
    args_maker = self._GetArgsMaker(rng,
                                    [start_shape, stop_shape],
                                    [dtype, dtype])
    start, stop = args_maker()
    ndim = len(np.shape(start + stop))
    for axis in range(-ndim, ndim):
      jnp_op = lambda start, stop: jnp.linspace(
        start, stop, num,
        endpoint=endpoint, retstep=retstep, dtype=dtype, axis=axis)
      # NumPy 1.20.0 changed the semantics of linspace to floor for integer
      # dtypes.
      if numpy_version >= (1, 20) or not np.issubdtype(dtype, np.integer):
        np_op = lambda start, stop: np.linspace(
          start, stop, num,
          endpoint=endpoint, retstep=retstep, dtype=dtype, axis=axis)
      else:
        def np_op(start, stop):
          out = np.linspace(start, stop, num, endpoint=endpoint,
                            retstep=retstep, axis=axis)
          if retstep:
            return np.floor(out[0]).astype(dtype), out[1]
          else:
            return np.floor(out).astype(dtype)

      self._CheckAgainstNumpy(np_op, jnp_op, args_maker,
                              check_dtypes=False, tol=tol)
      self._CompileAndCheck(jnp_op, args_maker,
                            check_dtypes=False, atol=tol, rtol=tol)

  @parameterized.named_parameters(
      jtu.cases_from_list(
        {"testcase_name": f"_dtype={dtype.__name__}", "dtype": dtype}
        for dtype in number_dtypes))
  def testLinspaceEndpoints(self, dtype):
    """Regression test for Issue #3014."""
    rng = jtu.rand_default(self.rng())
    endpoints = rng((2,), dtype)
    out = jnp.linspace(*endpoints, 10, dtype=dtype)
    self.assertAllClose(out[np.array([0, -1])], endpoints, rtol=0, atol=0)

  @parameterized.named_parameters(
      jtu.cases_from_list(
        {"testcase_name": ("_start_shape={}_stop_shape={}_num={}_endpoint={}"
                           "_base={}_dtype={}").format(
            start_shape, stop_shape, num, endpoint, base,
            dtype.__name__ if dtype else "None"),
         "start_shape": start_shape,
         "stop_shape": stop_shape,
         "num": num, "endpoint": endpoint, "base": base,
         "dtype": dtype}
        for start_shape in [(), (2,), (2, 2)]
        for stop_shape in [(), (2,), (2, 2)]
        for num in [0, 1, 2, 5, 20]
        for endpoint in [True, False]
        for base in [10.0, 2, np.e]
        # skip 16-bit floats due to insufficient precision for the test.
        for dtype in jtu.dtypes.inexact + [None,]))
  @jax.numpy_rank_promotion('allow')  # This test explicitly exercises implicit rank promotion.
  def testLogspace(self, start_shape, stop_shape, num,
                   endpoint, base, dtype):
    if (dtype in int_dtypes and
        jtu.device_under_test() in ("gpu", "tpu") and
        not config.x64_enabled):
      raise unittest.SkipTest("GPUx32 truncated exponentiation"
                              " doesn't exactly match other platforms.")
    rng = jtu.rand_default(self.rng())
    # relax default tolerances slightly
    tol = {np.float32: 1e-2, np.float64: 1e-6, np.complex64: 1e-3, np.complex128: 1e-6}
    args_maker = self._GetArgsMaker(rng,
                                    [start_shape, stop_shape],
                                    [dtype, dtype])
    start, stop = args_maker()
    ndim = len(np.shape(start + stop))
    for axis in range(-ndim, ndim):
      jnp_op = lambda start, stop: jnp.logspace(
        start, stop, num, endpoint=endpoint, base=base, dtype=dtype, axis=axis)
      @jtu.ignore_warning(category=RuntimeWarning,
                          message="overflow encountered in power")
      def np_op(start, stop):
        return np.logspace(start, stop, num, endpoint=endpoint,
                           base=base, dtype=dtype, axis=axis)
      self._CheckAgainstNumpy(np_op, jnp_op, args_maker,
                              check_dtypes=False, tol=tol)
      if dtype in (inexact_dtypes + [None,]):
        # Why do compiled and op-by-op float16 np.power numbers differ
        # slightly more than expected?
        atol = {np.float16: 1e-2}
        self._CompileAndCheck(jnp_op, args_maker,
                              check_dtypes=False, atol=atol, rtol=tol)

  @parameterized.named_parameters(
      jtu.cases_from_list(
        {"testcase_name": ("_start_shape={}_stop_shape={}_num={}_endpoint={}"
                           "_dtype={}_axis={}").format(
            start_shape, stop_shape, num, endpoint,
            dtype.__name__ if dtype else "None", axis),
         "start_shape": start_shape,
         "stop_shape": stop_shape,
         "num": num, "endpoint": endpoint,
         "dtype": dtype, "axis": axis}
        for start_shape in [(), (2,), (2, 2)]
        for stop_shape in [(), (2,), (2, 2)]
        for num in [0, 1, 2, 5, 20]
        for endpoint in [True, False]
        # NB: numpy's geomspace gives nonsense results on integer types
        for dtype in inexact_dtypes + [None,]
        for axis in range(-max(len(start_shape), len(stop_shape)),
                          max(len(start_shape), len(stop_shape)))))
  @jax.numpy_rank_promotion('allow')  # This test explicitly exercises implicit rank promotion.
  def testGeomspace(self, start_shape, stop_shape, num,
                    endpoint, dtype, axis):
    rng = jtu.rand_default(self.rng())
    # relax default tolerances slightly
    tol = {np.float16: 4e-3, np.float32: 2e-3, np.float64: 1e-14,
           np.complex128: 1e-14}
    def args_maker():
      """Test the set of inputs np.geomspace is well-defined on."""
      start, stop = self._GetArgsMaker(rng,
                                [start_shape, stop_shape],
                                [dtype, dtype])()
      # np.geomspace can't handle differently ranked tensors
      # w. negative numbers!
      start, stop = jnp.broadcast_arrays(start, stop)
      if dtype in complex_dtypes:
        return start, stop
      # to avoid NaNs, non-complex start and stop cannot
      # differ in sign, elementwise
      start = start * jnp.sign(start) * jnp.sign(stop)
      return start, stop
    start, stop = args_maker()
    def jnp_op(start, stop):
      return jnp.geomspace(start, stop, num, endpoint=endpoint, dtype=dtype,
                           axis=axis)
    def np_op(start, stop):
      start = start.astype(np.float32) if dtype == jnp.bfloat16 else start
      stop = stop.astype(np.float32) if dtype == jnp.bfloat16 else stop
      return np.geomspace(
        start, stop, num, endpoint=endpoint,
        dtype=dtype if dtype != jnp.bfloat16 else np.float32,
        axis=axis).astype(dtype)
    self._CheckAgainstNumpy(np_op, jnp_op, args_maker,
                            check_dtypes=False, tol=tol)
    if dtype in (inexact_dtypes + [None,]):
      self._CompileAndCheck(jnp_op, args_maker,
                            check_dtypes=False, atol=tol, rtol=tol)

  def testDisableNumpyRankPromotionBroadcasting(self):
    try:
      prev_flag = config._read('jax_numpy_rank_promotion')
      FLAGS.jax_numpy_rank_promotion = "allow"
      jnp.ones(2) + jnp.ones((1, 2))  # works just fine
    finally:
      FLAGS.jax_numpy_rank_promotion = prev_flag

    try:
      prev_flag = config._read('jax_numpy_rank_promotion')
      FLAGS.jax_numpy_rank_promotion = "raise"
      self.assertRaises(ValueError, lambda: jnp.ones(2) + jnp.ones((1, 2)))
      jnp.ones(2) + 3  # don't want to raise for scalars
    finally:
      FLAGS.jax_numpy_rank_promotion = prev_flag

    try:
      prev_flag = config._read('jax_numpy_rank_promotion')
      FLAGS.jax_numpy_rank_promotion = "warn"
      self.assertWarnsRegex(UserWarning, "Following NumPy automatic rank promotion for add on "
                            r"shapes \(2,\) \(1, 2\).*", lambda: jnp.ones(2) + jnp.ones((1, 2)))
      jnp.ones(2) + 3  # don't want to warn for scalars
    finally:
      FLAGS.jax_numpy_rank_promotion = prev_flag

  @unittest.skip("Test fails on CI, perhaps due to JIT caching")
  def testDisableNumpyRankPromotionBroadcastingDecorator(self):
    with jax.numpy_rank_promotion("allow"):
      jnp.ones(2) + jnp.ones((1, 2))  # works just fine

    with jax.numpy_rank_promotion("raise"):
      self.assertRaises(ValueError, lambda: jnp.ones(2) + jnp.ones((1, 2)))
      jnp.ones(2) + 3  # don't want to raise for scalars

    with jax.numpy_rank_promotion("warn"):
      self.assertWarnsRegex(UserWarning, "Following NumPy automatic rank promotion for add on "
                            r"shapes \(2,\) \(1, 2\).*", lambda: jnp.ones(2) + jnp.ones((1, 2)))
      jnp.ones(2) + 3  # don't want to warn for scalars

  def testStackArrayArgument(self):
    # tests https://github.com/google/jax/issues/1271
    @jax.jit
    def foo(x):
      return jnp.stack(x)
    foo(np.zeros(2))  # doesn't crash

    @jax.jit
    def foo(x):
      return jnp.concatenate(x)
    foo(np.zeros((2, 2)))  # doesn't crash

  def testReluGradientConstants(self):
    # This is a regression test that verifies that constants associated with the
    # gradient of np.maximum (from lax._balanced_eq) aren't hoisted into the
    # outermost jaxpr. This was producing some large materialized constants for
    # every relu activation in a model.
    def body(i, xy):
      x, y = xy
      y = y + jax.grad(lambda z: jnp.sum(jnp.maximum(z, 0.)))(x)
      return x, y

    f = lambda y: lax.fori_loop(0, 5, body, (y, y))
    jaxpr = jax.make_jaxpr(f)(np.zeros((3, 4), np.float32))
    self.assertFalse(
      any(np.array_equal(x, np.full((3, 4), 2., dtype=np.float32))
          for x in jaxpr.consts))

  @parameterized.named_parameters(
      {"testcase_name": f"_from={from_shape}_to={to_shape}",
       "from_shape": from_shape, "to_shape": to_shape}
      for from_shape, to_shape in [
          [(1, 3), (4, 3)],
          [(3,), (2, 1, 3)],
          [(3,), (3, 3)],
          [(1,), (3,)],
          [(1,), 3],
      ])
  def testBroadcastTo(self, from_shape, to_shape):
    rng = jtu.rand_default(self.rng())
    args_maker = self._GetArgsMaker(rng, [from_shape], [np.float32])
    np_op = lambda x: np.broadcast_to(x, to_shape)
    jnp_op = lambda x: jnp.broadcast_to(x, to_shape)
    self._CheckAgainstNumpy(np_op, jnp_op, args_maker)
    self._CompileAndCheck(jnp_op, args_maker)

  @parameterized.named_parameters(
      {"testcase_name": f"_{shapes}", "shapes": shapes, "broadcasted_shape": broadcasted_shape}
      for shapes, broadcasted_shape in [
        [[], ()],
        [[()], ()],
        [[(1, 3), (4, 3)], (4, 3)],
        [[(3,), (2, 1, 3)], (2, 1, 3)],
        [[(3,), (3, 3)], (3, 3)],
        [[(1,), (3,)], (3,)],
        [[(1,), 3], (3,)],
        [[(6, 7), (5, 6, 1), (7,), (5, 1, 7)], (5, 6, 7)],
        [[[1], [0, 1]], (0, 1)],
        [[(1,), np.array([0, 1])], (0, 1)],
    ])
  def testBroadcastShapes(self, shapes, broadcasted_shape):
    # Test against np.broadcast_shapes once numpy 1.20 is minimum required version
    np.testing.assert_equal(jnp.broadcast_shapes(*shapes), broadcasted_shape)

  def testBroadcastToIssue1522(self):
    self.assertRaisesRegex(
        ValueError, "Incompatible shapes for broadcasting: .*",
        lambda: jnp.broadcast_to(np.ones((2, 3)), (1, 3)))

  def testBroadcastToIntIssue1548(self):
    self.assertAllClose(jnp.broadcast_to(1, (3, 2)), np.ones((3, 2)),
                        check_dtypes=False)

  def testBroadcastToOnScalar(self):
    self.assertIsInstance(jnp.broadcast_to(10.0, ()), jnp.ndarray)
    self.assertIsInstance(np.broadcast_to(10.0, ()), np.ndarray)

  def testPrecision(self):

    ones_1d = np.ones((2,))
    ones_2d = np.ones((2, 2))
    ones_3d = np.ones((2, 2, 2))
    HIGHEST = lax.Precision.HIGHEST

    jtu.assert_dot_precision(None, jnp.dot, ones_1d, ones_1d)
    jtu.assert_dot_precision(
        HIGHEST,
        partial(jnp.dot, precision=HIGHEST),
        ones_1d, ones_1d)
    jtu.assert_dot_precision(
        HIGHEST,
        partial(jnp.dot, precision=HIGHEST),
        ones_3d, ones_3d)
    jtu.assert_dot_precision(
        HIGHEST,
        partial(jnp.matmul, precision=HIGHEST),
        ones_2d, ones_2d)
    jtu.assert_dot_precision(
        HIGHEST,
        partial(jnp.vdot, precision=HIGHEST),
        ones_1d, ones_1d)
    jtu.assert_dot_precision(
        HIGHEST,
        partial(jnp.tensordot, axes=2, precision=HIGHEST),
        ones_2d, ones_2d)
    jtu.assert_dot_precision(
        HIGHEST,
        partial(jnp.tensordot, axes=(0, 0), precision=HIGHEST),
        ones_1d, ones_1d)
    jtu.assert_dot_precision(
        HIGHEST,
        partial(jnp.tensordot, axes=((0,), (0,)), precision=HIGHEST),
        ones_1d, ones_1d)
    jtu.assert_dot_precision(
        HIGHEST,
        partial(jnp.einsum, 'i,i', precision=HIGHEST),
        ones_1d, ones_1d)
    jtu.assert_dot_precision(
        HIGHEST,
        partial(jnp.einsum, 'ij,ij', precision=HIGHEST),
        ones_2d, ones_2d)
    jtu.assert_dot_precision(
        HIGHEST,
        partial(jnp.inner, precision=HIGHEST),
        ones_1d, ones_1d)

  @parameterized.named_parameters(
      jtu.cases_from_list(
        {"testcase_name": "_shape={}_varargs={} axis={}_dtype={}".format(
            shape, varargs, axis, dtype),
         "shape": shape, "varargs": varargs, "axis": axis, "dtype": dtype}
        for shape in [(10,), (10, 15), (10, 15, 20)]
        for _num_axes in range(len(shape))
        for varargs in itertools.combinations(range(1, len(shape) + 1), _num_axes)
        for axis in itertools.combinations(range(len(shape)), _num_axes)
        for dtype in inexact_dtypes))
  def testGradient(self, shape, varargs, axis, dtype):
    rng = jtu.rand_default(self.rng())
    args_maker = self._GetArgsMaker(rng, [shape], [dtype])
    jnp_fun = lambda y: jnp.gradient(y, *varargs, axis=axis)
    np_fun = lambda y: np.gradient(y, *varargs, axis=axis)
    self._CheckAgainstNumpy(
        np_fun, jnp_fun, args_maker, check_dtypes=False)
    self._CompileAndCheck(jnp_fun, args_maker)

  def testZerosShapeErrors(self):
    # see https://github.com/google/jax/issues/1822
    self.assertRaisesRegex(
        TypeError,
        "Shapes must be 1D sequences of concrete values of integer type.*",
        lambda: jnp.zeros(1.))
    self.assertRaisesRegex(
        TypeError,
        r"Shapes must be 1D sequences of concrete values of integer type.*\n"
        "If using `jit`, try using `static_argnums` or applying `jit` to "
        "smaller subfunctions.",
        lambda: jax.jit(jnp.zeros)(2))

  def testTraceMethod(self):
    x = self.rng().randn(3, 4).astype(jnp.float_)
    self.assertAllClose(x.trace(), jnp.array(x).trace())
    self.assertAllClose(x.trace(), jax.jit(lambda y: y.trace())(x))

  def testIntegerPowersArePrecise(self):
    # See https://github.com/google/jax/pull/3036
    # Checks if the squares of float32 integers have no numerical errors.
    # It should be satisfied with all integers less than sqrt(2**24).
    x = jnp.arange(-2**12, 2**12, dtype=jnp.int32)
    np.testing.assert_array_equal(jnp.square(x.astype(jnp.float32)), x * x)
    np.testing.assert_array_equal(x.astype(jnp.float32) ** 2, x * x)

    # Similarly for cubes.
    x = jnp.arange(-2**8, 2**8, dtype=jnp.int32)
    np.testing.assert_array_equal(x.astype(jnp.float32) ** 3, x * x * x)

    x = np.arange(10, dtype=np.float32)
    for i in range(10):
      self.assertAllClose(x.astype(jnp.float32) ** i, x ** i,
                          check_dtypes=False)

  def testToBytes(self):
    v = np.arange(12, dtype=np.int32).reshape(3, 4)
    for order in ['C', 'F']:
      self.assertEqual(jnp.asarray(v).tobytes(order), v.tobytes(order))

  def testToList(self):
    v = np.arange(12, dtype=np.int32).reshape(3, 4)
    self.assertEqual(jnp.asarray(v).tolist(), v.tolist())

  def testReductionWithRepeatedAxisError(self):
    with self.assertRaisesRegex(ValueError, r"duplicate value in 'axis': \(0, 0\)"):
      jnp.sum(jnp.arange(3), (0, 0))

  def testArangeConcretizationError(self):
    msg = r"It arose in jax.numpy.arange argument `{}`".format
    with self.assertRaisesRegex(jax.core.ConcretizationTypeError, msg('stop')):
      jax.jit(jnp.arange)(3)

    with self.assertRaisesRegex(jax.core.ConcretizationTypeError, msg('start')):
      jax.jit(lambda start: jnp.arange(start, 3))(0)

    with self.assertRaisesRegex(jax.core.ConcretizationTypeError, msg('stop')):
      jax.jit(lambda stop: jnp.arange(0, stop))(3)

  @parameterized.named_parameters(jtu.cases_from_list(
    {"testcase_name": str(dtype), "dtype": dtype}
    for dtype in [None] + float_dtypes))
  def testArange64Bit(self, dtype):
    # Test that jnp.arange uses 64-bit arithmetic to define its range, even if the
    # output has another dtype. The issue here is that if python scalar inputs to
    # jnp.arange are cast to float32 before the range is computed, it changes the
    # number of elements output by the range.  It's unclear whether this was deliberate
    # behavior in the initial implementation, but it's behavior that downstream users
    # have come to rely on.
    args = (1.2, 4.8, 0.24)

    # Ensure that this test case leads to differing lengths if cast to float32.
    self.assertLen(np.arange(*args), 15)
    self.assertLen(np.arange(*map(np.float32, args)), 16)

    jnp_fun = lambda: jnp.arange(*args, dtype=dtype)
    np_fun = jtu.with_jax_dtype_defaults(lambda: np.arange(*args, dtype=dtype), dtype is None)
    args_maker = lambda: []
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)
    self._CompileAndCheck(jnp_fun, args_maker)

  def testIssue2347(self):
    # https://github.com/google/jax/issues/2347
    object_list = List[Tuple[jnp.array, float, float, jnp.array, bool]]
    self.assertRaises(TypeError, jnp.array, object_list)

    np_object_list = np.array(object_list)
    self.assertRaises(TypeError, jnp.array, np_object_list)

  @parameterized.named_parameters(jtu.cases_from_list(
    {"testcase_name": jtu.format_test_name_suffix("", shapes, dtypes),
     "shapes": shapes, "dtypes": dtypes}
    for shapes in filter(
      _shapes_are_broadcast_compatible,
      itertools.combinations_with_replacement(all_shapes, 2))
    for dtypes in itertools.product(
      *(_valid_dtypes_for_shape(s, complex_dtypes) for s in shapes))))
  @jax.numpy_rank_promotion('allow')  # This test explicitly exercises implicit rank promotion.
  def testLogaddexpComplex(self, shapes, dtypes):
    @jtu.ignore_warning(category=RuntimeWarning, message="invalid value.*")
    def np_op(x1, x2):
      return np.log(np.exp(x1) + np.exp(x2))

    rng = jtu.rand_some_nan(self.rng())
    args_maker = lambda: tuple(rng(shape, dtype) for shape, dtype in zip(shapes, dtypes))
    if jtu.device_under_test() == 'tpu':
      tol = {np.complex64: 1e-3, np.complex128: 1e-10}
    else:
      tol = {np.complex64: 1e-5, np.complex128: 1e-14}
    self._CheckAgainstNumpy(_promote_like_jnp(np_op), jnp.logaddexp, args_maker, tol=tol)
    self._CompileAndCheck(jnp.logaddexp, args_maker, rtol=tol, atol=tol)

  @parameterized.named_parameters(jtu.cases_from_list(
    {"testcase_name": jtu.format_test_name_suffix("", shapes, dtypes),
     "shapes": shapes, "dtypes": dtypes}
    for shapes in filter(
      _shapes_are_broadcast_compatible,
      itertools.combinations_with_replacement(all_shapes, 2))
    for dtypes in itertools.product(
      *(_valid_dtypes_for_shape(s, complex_dtypes) for s in shapes))))
  @jax.numpy_rank_promotion('allow')  # This test explicitly exercises implicit rank promotion.
  def testLogaddexp2Complex(self, shapes, dtypes):
    @jtu.ignore_warning(category=RuntimeWarning, message="invalid value.*")
    def np_op(x1, x2):
      return np.log2(np.exp2(x1) + np.exp2(x2))

    rng = jtu.rand_some_nan(self.rng())
    args_maker = lambda: tuple(rng(shape, dtype) for shape, dtype in zip(shapes, dtypes))
    if jtu.device_under_test() == 'tpu':
      tol = {np.complex64: 1e-3, np.complex128: 1e-10}
    else:
      tol = {np.complex64: 1e-5, np.complex128: 1e-14}
    self._CheckAgainstNumpy(_promote_like_jnp(np_op), jnp.logaddexp2, args_maker, tol=tol)
    self._CompileAndCheck(jnp.logaddexp2, args_maker, rtol=tol, atol=tol)

  def testDefaultDtypes(self):
    precision = config.jax_default_dtype_bits
    assert precision in ['32', '64']
    self.assertEqual(jnp.bool_, np.bool_)
    self.assertEqual(jnp.int_, np.int32 if precision == '32' else np.int64)
    self.assertEqual(jnp.uint, np.uint32 if precision == '32' else np.uint64)
    self.assertEqual(jnp.float_, np.float32 if precision == '32' else np.float64)
    self.assertEqual(jnp.complex_, np.complex64 if precision == '32' else np.complex128)

  def testFromBuffer(self):
    buf = b'\x01\x02\x03'
    expected = np.frombuffer(buf, dtype='uint8')
    actual = jnp.frombuffer(buf, dtype='uint8')
    self.assertArraysEqual(expected, actual)

  def testFromFunction(self):
    def f(x, y, z):
      return x + 2 * y + 3 * z
    shape = (3, 4, 5)
    expected = np.fromfunction(f, shape=shape)
    actual = jnp.fromfunction(f, shape=shape)
    self.assertArraysEqual(expected, actual)

  def testFromString(self):
    s = "1,2,3"
    expected = np.fromstring(s, sep=',', dtype=int)
    actual = jnp.fromstring(s, sep=',', dtype=int)
    self.assertArraysEqual(expected, actual)


# Most grad tests are at the lax level (see lax_test.py), but we add some here
# as needed for e.g. particular compound ops of interest.

GradTestSpec = collections.namedtuple(
    "GradTestSpec",
    ["op", "nargs", "order", "rng_factory", "dtypes", "name", "tol"])
def grad_test_spec(op, nargs, order, rng_factory, dtypes, name=None, tol=None):
  return GradTestSpec(
      op, nargs, order, rng_factory, dtypes, name or op.__name__, tol)

GRAD_TEST_RECORDS = [
    grad_test_spec(jnp.arcsinh, nargs=1, order=2,
                   rng_factory=jtu.rand_positive,
                   dtypes=[np.float64, np.complex64],
                   tol={np.complex64: 2e-2}),
    grad_test_spec(jnp.arccosh, nargs=1, order=2,
                   rng_factory=jtu.rand_positive,
                   dtypes=[np.float64, np.complex64],
                   tol={np.complex64: 2e-2}),
    grad_test_spec(jnp.arctanh, nargs=1, order=2,
                   rng_factory=partial(jtu.rand_uniform, low=-0.9, high=0.9),
                   dtypes=[np.float64, np.complex64],
                   tol={np.complex64: 2e-2}),
    grad_test_spec(jnp.logaddexp, nargs=2, order=1,
                   rng_factory=partial(jtu.rand_uniform, low=-0.9, high=0.9),
                   dtypes=[np.float64], tol=1e-4),
    grad_test_spec(jnp.logaddexp2, nargs=2, order=2,
                   rng_factory=partial(jtu.rand_uniform, low=-0.9, high=0.9),
                   dtypes=[np.float64], tol=1e-4),
]

GradSpecialValuesTestSpec = collections.namedtuple(
    "GradSpecialValuesTestSpec", ["op", "values", "order"])

GRAD_SPECIAL_VALUE_TEST_RECORDS = [
    GradSpecialValuesTestSpec(jnp.arcsinh, [0., 1000.], 2),
    GradSpecialValuesTestSpec(jnp.arccosh, [1000.], 2),
    GradSpecialValuesTestSpec(jnp.arctanh, [0.], 2),
    GradSpecialValuesTestSpec(jnp.sinc, [0.], 1),
]


class NumpyGradTests(jtu.JaxTestCase):

  @parameterized.named_parameters(itertools.chain.from_iterable(
      jtu.cases_from_list(
        {"testcase_name": jtu.format_test_name_suffix(
            rec.name, shapes, itertools.repeat(dtype)),
         "op": rec.op, "rng_factory": rec.rng_factory, "shapes": shapes, "dtype": dtype,
         "order": rec.order, "tol": rec.tol}
        for shapes in itertools.combinations_with_replacement(nonempty_shapes, rec.nargs)
        for dtype in rec.dtypes)
      for rec in GRAD_TEST_RECORDS))
  @jax.numpy_rank_promotion('allow')  # This test explicitly exercises implicit rank promotion.
  def testOpGrad(self, op, rng_factory, shapes, dtype, order, tol):
    rng = rng_factory(self.rng())
    tol = jtu.join_tolerance(tol, {np.float32: 1e-1, np.float64: 1e-3,
                                   np.complex64: 1e-1, np.complex128: 1e-3})
    args = tuple(rng(shape, dtype) for shape in shapes)
    check_grads(op, args, order, ["fwd", "rev"], tol, tol)

  @parameterized.named_parameters(itertools.chain.from_iterable(
      jtu.cases_from_list(
          {"testcase_name": f"_{rec.op.__name__}_{special_value}",
           "op": rec.op, "special_value": special_value, "order": rec.order}
          for special_value in rec.values)
      for rec in GRAD_SPECIAL_VALUE_TEST_RECORDS))
  def testOpGradSpecialValue(self, op, special_value, order):
    check_grads(op, (special_value,), order, ["fwd", "rev"],
                atol={np.float32: 3e-3})

  def testSincAtZero(self):
    # Some manual tests for sinc at zero, since it doesn't have well-behaved
    # numerical derivatives at zero
    def deriv(f):
      return lambda x: jax.jvp(f, (x,), (1.,))[1]

    def apply_all(fns, x):
      for f in fns:
        x = f(x)
      return x

    d1 = 0.
    for ops in itertools.combinations_with_replacement([deriv, jax.grad], 1):
      self.assertAllClose(apply_all(ops, jnp.sinc)(0.), d1)

    d2 = -np.pi ** 2 / 3
    for ops in itertools.combinations_with_replacement([deriv, jax.grad], 2):
      self.assertAllClose(apply_all(ops, jnp.sinc)(0.), d2)

    d3 = 0.
    for ops in itertools.combinations_with_replacement([deriv, jax.grad], 3):
      self.assertAllClose(apply_all(ops, jnp.sinc)(0.), d3)

    d4 = np.pi ** 4 / 5
    for ops in itertools.combinations_with_replacement([deriv, jax.grad], 4):
      self.assertAllClose(apply_all(ops, jnp.sinc)(0.), d4)

  def testSincGradArrayInput(self):
    # tests for a bug almost introduced in #5077
    jax.grad(lambda x: jnp.sinc(x).sum())(jnp.arange(10.))  # doesn't crash

  def testTakeAlongAxisIssue1521(self):
    # https://github.com/google/jax/issues/1521
    idx = jnp.repeat(jnp.arange(3), 10).reshape((30, 1))

    def f(x):
      y = x * jnp.arange(3.).reshape((1, 3))
      return jnp.take_along_axis(y, idx, -1).sum()

    check_grads(f, (1.,), order=1)

  @parameterized.named_parameters(jtu.cases_from_list(
    {"testcase_name": jtu.format_test_name_suffix("", shapes, itertools.repeat(dtype)),
     "shapes": shapes, "dtype": dtype}
    for shapes in filter(
      _shapes_are_broadcast_compatible,
      itertools.combinations_with_replacement(nonempty_shapes, 2))
    for dtype in (np.complex128, )))
  @jax.numpy_rank_promotion('allow')  # This test explicitly exercises implicit rank promotion.
  def testGradLogaddexpComplex(self, shapes, dtype):
    rng = jtu.rand_default(self.rng())
    args = tuple(rng(shape, dtype) for shape in shapes)
    if jtu.device_under_test() == "tpu":
      tol = 5e-2
    else:
      tol = 3e-2
    check_grads(jnp.logaddexp, args, 1, ["fwd", "rev"], tol, tol)

  @parameterized.named_parameters(jtu.cases_from_list(
    {"testcase_name": jtu.format_test_name_suffix("", shapes, itertools.repeat(dtype)),
     "shapes": shapes, "dtype": dtype}
    for shapes in filter(
      _shapes_are_broadcast_compatible,
      itertools.combinations_with_replacement(nonempty_shapes, 2))
    for dtype in (np.complex128, )))
  @jax.numpy_rank_promotion('allow')  # This test explicitly exercises implicit rank promotion.
  def testGradLogaddexp2Complex(self, shapes, dtype):
    rng = jtu.rand_default(self.rng())
    args = tuple(rng(shape, dtype) for shape in shapes)
    if jtu.device_under_test() == "tpu":
      tol = 5e-2
    else:
      tol = 3e-2
    check_grads(jnp.logaddexp2, args, 1, ["fwd", "rev"], tol, tol)


class NumpySignaturesTest(jtu.JaxTestCase):

  def testWrappedSignaturesMatch(self):
    """Test that jax.numpy function signatures match numpy."""
    jnp_funcs = {name: getattr(jnp, name) for name in dir(jnp)}
    func_pairs = {name: (fun, fun.__np_wrapped__) for name, fun in jnp_funcs.items()
                  if hasattr(fun, '__np_wrapped__')}
    assert len(func_pairs) > 0

    # TODO(jakevdp): fix some of the following signatures. Some are due to wrong argument names.
    unsupported_params = {
      'asarray': ['like'],
      'broadcast_to': ['subok', 'array'],
      'clip': ['kwargs'],
      'copy': ['subok'],
      'corrcoef': ['ddof', 'bias', 'dtype'],
      'cov': ['dtype'],
      'empty_like': ['subok', 'order'],
      'einsum': ['kwargs'],
      'einsum_path': ['einsum_call'],
      'eye': ['order', 'like'],
      'identity': ['like'],
      'full': ['order', 'like'],
      'full_like': ['subok', 'order'],
      'fromfunction': ['like'],
      'histogram': ['normed'],
      'histogram2d': ['normed'],
      'histogramdd': ['normed'],
      'ones': ['order', 'like'],
      'ones_like': ['subok', 'order'],
      'tri': ['like'],
      'zeros_like': ['subok', 'order']
    }

    extra_params = {
      'broadcast_to': ['arr'],
      'einsum': ['precision'],
      'einsum_path': ['subscripts'],
      'take_along_axis': ['mode'],
    }

    mismatches = {}

    for name, (jnp_fun, np_fun) in func_pairs.items():
      # broadcast_shapes is not available in numpy < 1.20
      if numpy_version < (1, 20) and name == "broadcast_shapes":
        continue
      # Some signatures have changed; skip for older numpy versions.
      if numpy_version < (1, 19) and name in ['einsum_path', 'gradient', 'isscalar']:
        continue
      if numpy_version < (1, 22) and name in ['quantile', 'nanquantile',
                                              'percentile', 'nanpercentile']:
        continue
      # Note: can't use inspect.getfullargspec due to numpy issue
      # https://github.com/numpy/numpy/issues/12225
      try:
        np_params = inspect.signature(np_fun).parameters
      except ValueError:
        # Some functions cannot be inspected
        continue
      jnp_params = inspect.signature(jnp_fun).parameters
      extra = set(extra_params.get(name, []))
      unsupported = set(unsupported_params.get(name, []))

      # Checks to prevent tests from becoming out-of-date. If these fail,
      # it means that extra_params or unsupported_params need to be updated.
      assert extra.issubset(jnp_params), f"{name}: extra={extra} is not a subset of jnp_params={set(jnp_params)}."
      assert not unsupported.intersection(jnp_params), f"{name}: unsupported={unsupported} overlaps with jnp_params={set(jnp_params)}."

      # Skip functions that only have *args and **kwargs; we can't introspect these further.
      var_args = (inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD)
      if all(p.kind in var_args for p in jnp_params.values()):
        continue
      if all(p.kind in var_args for p in np_params.values()):
        continue

      # Remove known extra parameters.
      jnp_params = {a: p for a, p in jnp_params.items() if a not in extra}

      # Remove known unsupported parameters.
      np_params = {a: p for a, p in np_params.items() if a not in unsupported}

      # Older versions of numpy may have fewer parameters; to avoid extraneous errors on older numpy
      # versions, we allow for jnp to have more parameters.
      if list(jnp_params)[:len(np_params)] != list(np_params):
        mismatches[name] = {'np_params': list(np_params), 'jnp_params': list(jnp_params)}

    self.assertEqual(mismatches, {})


_available_numpy_dtypes: List[str] = [dtype.__name__ for dtype in jtu.dtypes.all
                                      if dtype != dtypes.bfloat16]


def _all_numpy_ufuncs() -> Iterator[str]:
  """Generate the names of all ufuncs in the top-level numpy namespace."""
  for name in dir(np):
    f = getattr(np, name)
    if isinstance(f, np.ufunc):
      yield name


def _dtypes_for_ufunc(name: str) -> Iterator[Tuple[str, ...]]:
  """Generate valid dtypes of inputs to the given numpy ufunc."""
  func = getattr(np, name)
  for arg_dtypes in itertools.product(_available_numpy_dtypes, repeat=func.nin):
    args = (np.ones(1, dtype=dtype) for dtype in arg_dtypes)
    try:
      with warnings.catch_warnings():
        warnings.filterwarnings("ignore", "divide by zero", RuntimeWarning)
        _ = func(*args)
    except TypeError:
      pass
    else:
      yield arg_dtypes


class NumpyUfuncTests(jtu.JaxTestCase):

  @parameterized.named_parameters(
    {"testcase_name": f"_{name}_{','.join(arg_dtypes)}",
     "name": name, "arg_dtypes": arg_dtypes}
    for name in _all_numpy_ufuncs()
    for arg_dtypes in jtu.cases_from_list(_dtypes_for_ufunc(name)))
  def testUfuncInputTypes(self, name, arg_dtypes):
    # TODO(jakevdp): fix following failures and remove from this exception list.
    if (name in ['divmod', 'floor_divide', 'fmod', 'gcd', 'left_shift', 'mod',
                 'power', 'remainder', 'right_shift', 'rint', 'square']
        and 'bool_' in arg_dtypes):
      self.skipTest(f"jax.numpy does not support {name}{tuple(arg_dtypes)}")
    if name == 'arctanh' and jnp.issubdtype(arg_dtypes[0], jnp.complexfloating):
      self.skipTest("np.arctanh & jnp.arctanh have mismatched NaNs for complex input.")

    jnp_op = getattr(jnp, name)
    np_op = getattr(np, name)
    np_op = jtu.ignore_warning(category=RuntimeWarning,
                               message="divide by zero.*")(np_op)
    args_maker = lambda: tuple(np.ones(1, dtype=dtype) for dtype in arg_dtypes)

    try:
      jnp_op(*args_maker())
    except NotImplementedError:
      self.skipTest(f"jtu.{name} is not yet implemented.")

    # large tol comes from the fact that numpy returns float16 in places
    # that jnp returns float32. e.g. np.cos(np.uint8(0))
    self._CheckAgainstNumpy(np_op, jnp_op, args_maker, check_dtypes=False, tol=1E-2)


class NumpyDocTests(jtu.JaxTestCase):

  def test_lax_numpy_docstrings(self):
    # Test that docstring wrapping & transformation didn't fail.

    # Functions that have their own docstrings & don't wrap numpy.
    known_exceptions = {'broadcast_arrays', 'fromfile', 'fromiter', 'vectorize'}

    for name in dir(jnp):
      if name in known_exceptions or name.startswith('_'):
        continue

      # We only check signatures of functions.
      obj = getattr(jnp, name)
      if isinstance(obj, type) or not callable(obj):
        continue

      # Some jnp functions are imported from numpy or jax.dtypes directly.
      if any(obj is getattr(mod, obj.__name__, None) for mod in [np, dtypes]):
        continue

      wrapped_fun = obj.__np_wrapped__

      # If the wrapped function has a docstring, obj should too
      if wrapped_fun.__doc__ and not obj.__doc__:
        raise Exception(f"jnp.{name} does not contain wrapped docstring.")

      if obj.__doc__ and "*Original docstring below.*" not in obj.__doc__:
        raise Exception(f"jnp.{name} does not have a wrapped docstring.")

  @parameterized.named_parameters(
    {"testcase_name": "_jit" if jit else "", "jit": jit} for jit in [True, False])
  def test_wrapped_function_parameters(self, jit):
    def orig(x):
      """Example Docstring

      Parameters
      ----------
      x : array_like
        Input Data

        .. versionadded:: 1.8.0
      out : array_like, optional
        Output to overwrite
      other_arg : Any
        not used

      Returns
      -------
      x : input
      """
      return x

    def wrapped(x, out=None):
      return x

    if jit:
      wrapped = jax.jit(wrapped)

    wrapped = _wraps(orig, skip_params=['out'])(wrapped)
    doc = wrapped.__doc__

    self.assertStartsWith(doc, "Example Docstring")
    self.assertIn("Original docstring below", doc)
    self.assertIn("Parameters", doc)
    self.assertIn("Returns", doc)
    self.assertNotIn('out', doc)
    self.assertNotIn('other_arg', doc)
    self.assertNotIn('versionadded', doc)


  def test_parse_numpydoc(self):
    # Unit test ensuring that _parse_numpydoc correctly parses docstrings for all
    # functions in NumPy's top-level namespace.
    section_titles = {'Attributes', 'Examples', 'Notes',
                      'Parameters', 'Raises', 'References',
                      'Returns', 'See also', 'See Also', 'Warnings', 'Warns'}
    headings = [title + '\n' + '-'*len(title) for title in section_titles]

    for name in dir(np):
      if name.startswith('_'):
        continue
      obj = getattr(np, name)
      if isinstance(obj, type):
        continue
      if not callable(obj):
        continue
      if 'built-in function' in repr(obj):
        continue
      parsed = _parse_numpydoc(obj.__doc__)

      # Check that no docstring is handled gracefully.
      if not obj.__doc__:
        self.assertEqual(parsed, ParsedDoc(obj.__doc__))
        continue

      # Check that no unexpected section names are found.
      extra_keys = parsed.sections.keys() - section_titles
      if extra_keys:
        raise ValueError(f"Extra section headers found in np.{name}: {extra_keys}")

      # Check that every docstring has a summary.
      if not parsed.summary:
        raise ValueError(f"No summary found for np.{name}")

      # Check that no expected headings are missed.
      for heading in headings:
        assert heading not in parsed.front_matter


if __name__ == "__main__":
  absltest.main(testLoader=jtu.JaxTestLoader())