dtypes_test.py

# Copyright 2019 The JAX Authors.
#
# 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 enum
import functools
import itertools
import operator
import unittest

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

import numpy as np

import jax
from jax._src import dtypes
from jax import numpy as jnp

from jax._src import test_util as jtu
from jax._src.lax import lax as lax_internal

from jax._src.config import config
config.parse_flags_with_absl()

FLAGS = config.FLAGS

# FP8 types not supported on TPU
_fp8_enabled = jtu.device_under_test() != 'tpu'
_int4_enabled = dtypes.int4 is not None

bool_dtypes = [np.dtype('bool')]

np_signed_dtypes = [np.dtype('int8'), np.dtype('int16'), np.dtype('int32'),
                   np.dtype('int64')]
signed_dtypes = list(np_signed_dtypes)

np_unsigned_dtypes = [np.dtype('uint8'), np.dtype('uint16'), np.dtype('uint32'),
                     np.dtype('uint64')]
unsigned_dtypes = list(np_unsigned_dtypes)

if _int4_enabled:
  int4_dtypes = [np.dtype('int4'), np.dtype('uint4')]
  signed_dtypes += [np.dtype('int4')]
  unsigned_dtypes += [np.dtype('uint4')]

np_float_dtypes = [np.dtype('float16'), np.dtype('float32'),
                   np.dtype('float64')]

float_dtypes = [np.dtype(dtypes.bfloat16)] + np_float_dtypes
custom_float_dtypes = [np.dtype(dtypes.bfloat16)]
if _fp8_enabled:
  fp8_dtypes = [np.dtype(dtypes.float8_e4m3b11fnuz), np.dtype(dtypes.float8_e4m3fn),
                np.dtype(dtypes.float8_e4m3fnuz), np.dtype(dtypes.float8_e5m2),
                np.dtype(dtypes.float8_e5m2fnuz)]
  float_dtypes += fp8_dtypes
  custom_float_dtypes += fp8_dtypes

complex_dtypes = [np.dtype('complex64'), np.dtype('complex128')]


all_dtypes = (bool_dtypes + signed_dtypes + unsigned_dtypes + float_dtypes +
              complex_dtypes)

scalar_types = [jnp.bool_, jnp.int8, jnp.int16, jnp.int32, jnp.int64,
                jnp.uint8, jnp.uint16, jnp.uint32, jnp.uint64,
                jnp.bfloat16, jnp.float16, jnp.float32, jnp.float64,
                jnp.complex64, jnp.complex128]

python_scalar_types = [bool, int, float, complex]

_EXPECTED_CANONICALIZE_X64 = {value: value for value in scalar_types}

_EXPECTED_CANONICALIZE_X32 = {value: value for value in scalar_types}
_EXPECTED_CANONICALIZE_X32[np.int64] = np.int32
_EXPECTED_CANONICALIZE_X32[np.uint64] = np.uint32
_EXPECTED_CANONICALIZE_X32[np.float64] = np.float32
_EXPECTED_CANONICALIZE_X32[np.complex128] = np.complex64
_EXPECTED_CANONICALIZE_X32[np.longlong] = np.int32

UINT_DTYPES = {
  8: np.uint8,
  16: np.uint16,
  32: np.uint32,
  64: np.uint64,
}

def identity(x):
  """A named identity function for use in tests"""
  return x


class DtypesTest(jtu.JaxTestCase):

  def test_canonicalize_type(self):
    expected = {
        True: _EXPECTED_CANONICALIZE_X64,
        False: _EXPECTED_CANONICALIZE_X32,
    }
    for in_dtype, expected_dtype in expected[config.x64_enabled].items():
      self.assertEqual(dtypes.canonicalize_dtype(in_dtype), expected_dtype)

  @parameterized.named_parameters(
    {"testcase_name": f"_type={type_.__name__}", "type_": type_}
    for type_ in python_scalar_types)
  def testDefaultTypes(self, type_):
    expected_dtype = dtypes.canonicalize_dtype(dtypes.python_scalar_dtypes[type_])
    for f in [jnp.array, jax.jit(jnp.array), jax.jit(lambda x: x)]:
      y = f(type_(0))
      self.assertTrue(isinstance(y, jax.Array), msg=(f, y))
      self.assertEqual(y.dtype, expected_dtype, msg=(f, y))

  def testUnsupportedType(self):
    with self.assertRaisesRegex(TypeError, "nonsense.* not understood"):
      dtypes.canonicalize_dtype("nonsense")

  @parameterized.named_parameters(
    {"testcase_name": f"_{swap=}_{jit=}",
     "swap": swap, "jit": jit}
    for swap in [False, True] for jit in [False, True])
  @jtu.ignore_warning(category=UserWarning,
                      message="Explicitly requested dtype.*")
  @jax.numpy_dtype_promotion('standard')
  def testBinaryPromotion(self, swap, jit):
    testcases = [
      (jnp.array(1.), 0., jnp.float64),
      (jnp.array(1.), jnp.array(0.), jnp.float64),
      (jnp.array(1.), jnp.array(0., dtype=jnp.float16), jnp.float16),
      (jnp.array(1.), jnp.array(0., dtype=jnp.float32), jnp.float32),
      (jnp.array(1.), jnp.array(0., dtype=jnp.float64), jnp.float64),
      (jnp.array(1., dtype=jnp.float16), 0., jnp.float16),
      (jnp.array(1., dtype=jnp.float32), 0., jnp.float32),
      (jnp.array(1., dtype=jnp.float64), 0., jnp.float64),
      (jnp.array(1., dtype=jnp.float16), jnp.array(0., dtype=jnp.float16), jnp.float16),
      (jnp.array(1., dtype=jnp.float16), jnp.array(0., dtype=jnp.float32), jnp.float32),
      (jnp.array(1., dtype=jnp.float16), jnp.array(0., dtype=jnp.float64), jnp.float64),
      (jnp.array(1., dtype=jnp.float32), jnp.array(0., dtype=jnp.float32), jnp.float32),
      (jnp.array(1., dtype=jnp.float32), jnp.array(0., dtype=jnp.float64), jnp.float64),
      (jnp.array(1., dtype=jnp.float64), jnp.array(0., dtype=jnp.float64), jnp.float64),
      (jnp.array([1.]), 0., jnp.float64),
      (jnp.array([1.]), jnp.array(0.), jnp.float64),
      (jnp.array([1.]), jnp.array(0., dtype=jnp.float16), jnp.float64),
      (jnp.array([1.]), jnp.array(0., dtype=jnp.float32), jnp.float64),
      (jnp.array([1.]), jnp.array(0., dtype=jnp.float64), jnp.float64),
      (jnp.array([1.], dtype=jnp.float32), jnp.array(0., dtype=jnp.float16), jnp.float32),
      (jnp.array([1.], dtype=jnp.float16), jnp.array(0., dtype=jnp.float32), jnp.float32),
      (jnp.array([1.], dtype=jnp.float16), 0., jnp.float16),
    ]
    op = jax.jit(operator.add) if jit else operator.add
    for x, y, dtype in testcases:
      x, y = (y, x) if swap else (x, y)
      z = op(x, y)
      self.assertTrue(isinstance(z, jax.Array), msg=(x, y, z))
      self.assertEqual(z.dtype, dtypes.canonicalize_dtype(dtype), msg=(x, y, z))

  @jax.numpy_dtype_promotion('strict')
  def testPromoteDtypesStrict(self):
    msg = ("Input dtypes .* have no available implicit dtype promotion "
           "path when jax_numpy_dtype_promotion=strict")

    assertTypePromotionError = functools.partial(
      self.assertRaisesRegex, dtypes.TypePromotionError, msg,
      dtypes.promote_types)

    # Check that strong types have diagonal promotion table:
    for t1 in all_dtypes:
      for t2 in all_dtypes:
        if t1 == t2:
          self.assertEqual(t1, dtypes.promote_types(t1, t2))
        else:
          assertTypePromotionError(t1, t2)

    # Promotion between weak types matches numpy promotion
    for t1 in [int, float, complex]:
      for t2 in [int, float, complex]:
        py_result = type(t1(0) + t2(0))
        # np.dtype(int) is int32 on Windows and int64 on Linux/Mac.
        py_result_dtype = (np.dtype(np.int64) if py_result is int
                           else np.dtype(py_result))
        lattice_dtype, lattice_weak_type = dtypes._lattice_result_type(t1, t2)
        self.assertTrue(lattice_weak_type)
        self.assertEqual(lattice_dtype, py_result_dtype)

    # Check that weak promotion only works if strong value is not cast:
    for t1 in bool_dtypes:
      assertTypePromotionError(t1, int)
      assertTypePromotionError(t1, float)
      assertTypePromotionError(t1, complex)
    for t1 in signed_dtypes + unsigned_dtypes:
      self.assertEqual(dtypes.promote_types(t1, int), t1)
      assertTypePromotionError(t1, float)
      assertTypePromotionError(t1, complex)
    for t1 in float_dtypes:
      self.assertEqual(dtypes.promote_types(t1, int), t1)
      self.assertEqual(dtypes.promote_types(t1, float), t1)
      assertTypePromotionError(t1, complex)
    for t1 in complex_dtypes:
      self.assertEqual(dtypes.promote_types(t1, int), t1)
      self.assertEqual(dtypes.promote_types(t1, float), t1)
      self.assertEqual(dtypes.promote_types(t1, complex), t1)

  @jax.numpy_dtype_promotion('standard')
  def testPromoteDtypesStandard(self):
    for t1 in all_dtypes:
      self.assertEqual(t1, dtypes.promote_types(t1, t1))

      self.assertEqual(t1, dtypes.promote_types(t1, np.bool_))
      # TODO(zhangqiaorjc): Consider more dtype promotion rules for fp8.
      if _fp8_enabled and t1 in fp8_dtypes:
        continue
      if _int4_enabled and t1 in int4_dtypes:
        continue
      self.assertEqual(np.dtype(np.complex128),
                       dtypes.promote_types(t1, np.complex128))

      for t2 in all_dtypes:
        # TODO(zhangqiaorjc): Consider more dtype promotion rules for fp8.
        if _fp8_enabled and t2 in fp8_dtypes:
          continue
        if _int4_enabled and t2 in int4_dtypes:
          continue
        # Symmetry
        self.assertEqual(dtypes.promote_types(t1, t2),
                         dtypes.promote_types(t2, t1))

    self.assertEqual(np.dtype(np.float32),
                     dtypes.promote_types(np.float16, dtypes.bfloat16))

    # Promotions of non-inexact types against inexact types always prefer
    # the inexact types.
    for t in float_dtypes + complex_dtypes:
      for i in bool_dtypes + signed_dtypes + unsigned_dtypes:
        # TODO(zhangqiaorjc): Consider more dtype promotion rules for fp8.
        if _fp8_enabled and t in fp8_dtypes:
          continue
        if _int4_enabled and (t in int4_dtypes or i in int4_dtypes):
          continue
        self.assertEqual(t, dtypes.promote_types(t, i))

    # Promotions between exact types, or between inexact types, match NumPy.
    for groups in [bool_dtypes + np_signed_dtypes + np_unsigned_dtypes,
                   np_float_dtypes + complex_dtypes]:
      for t1, t2 in itertools.combinations(groups, 2):
        self.assertEqual(np.promote_types(t1, t2),
                         dtypes.promote_types(t1, t2))

    # Promotion between weak types matches numpy promotion
    for t1 in [int, float, complex]:
      for t2 in [int, float, complex]:
        py_result = type(t1(0) + t2(0))
        # np.dtype(int) is int32 on Windows and int64 on Linux/Mac.
        py_result_dtype = (np.dtype(np.int64) if py_result is int
                           else np.dtype(py_result))
        lattice_dtype, lattice_weak_type = dtypes._lattice_result_type(t1, t2)
        self.assertTrue(lattice_weak_type)
        self.assertEqual(lattice_dtype, py_result_dtype)

  @parameterized.parameters([jnp.bool_, jnp.int32, jnp.bfloat16, jnp.float32, jnp.complex64])
  def testScalarInstantiation(self, scalar_type):
    a = scalar_type(1)
    self.assertEqual(a.dtype, jnp.dtype(scalar_type))
    self.assertIsInstance(a, jax.Array)
    self.assertEqual(0, jnp.ndim(a))
    self.assertIsInstance(np.dtype(scalar_type).type(1), scalar_type)

  def testIsSubdtype(self):
    for t in scalar_types:
      self.assertTrue(dtypes.issubdtype(t, t))
      self.assertTrue(dtypes.issubdtype(np.dtype(t).type, t))
      self.assertTrue(dtypes.issubdtype(t, np.dtype(t).type))
      self.assertTrue(dtypes.issubdtype(t, np.dtype(t)))
      if t != jnp.bfloat16:
        for category in [np.generic, jnp.inexact, jnp.integer, jnp.signedinteger,
                         jnp.unsignedinteger, jnp.floating, jnp.complexfloating]:
          self.assertEqual(dtypes.issubdtype(t, category),
                           np.issubdtype(np.dtype(t).type, category))
          self.assertEqual(dtypes.issubdtype(t, category),
                           np.issubdtype(np.dtype(t).type, category))

  def testIsSubdtypeExtended(self):
    self.assertTrue(dtypes.issubdtype(dtypes.extended, dtypes.extended))
    self.assertTrue(dtypes.issubdtype(dtypes.extended, np.generic))
    self.assertFalse(dtypes.issubdtype(dtypes.extended, np.number))

    self.assertTrue(jnp.issubdtype(dtypes.prng_key, dtypes.prng_key))
    self.assertTrue(jnp.issubdtype(dtypes.prng_key, dtypes.extended))
    self.assertTrue(jnp.issubdtype(dtypes.prng_key, np.generic))
    self.assertFalse(dtypes.issubdtype(dtypes.prng_key, np.number))

  @parameterized.product(dtype=custom_float_dtypes)
  def testIsSubdtypeCustomFloats(self, dtype):
    for dt in [dtype, np.dtype(dtype), str(np.dtype(dtype))]:
      self.assertTrue(dtypes.issubdtype(dt, dt))
      self.assertTrue(dtypes.issubdtype(dt, np.dtype(dtype)))
      self.assertTrue(dtypes.issubdtype(dt, str(np.dtype(dtype))))
      self.assertTrue(dtypes.issubdtype(dt, np.floating))
      self.assertTrue(dtypes.issubdtype(dt, np.inexact))
      self.assertTrue(dtypes.issubdtype(dt, np.number))
      self.assertTrue(dtypes.issubdtype(dt, np.generic))
      self.assertFalse(dtypes.issubdtype(dt, object))
      self.assertFalse(dtypes.issubdtype(dt, np.float64))
      self.assertFalse(dtypes.issubdtype(np.generic, dt))

  @parameterized.product(dtype=int4_dtypes)
  def testIsSubdtypeInt4(self, dtype):
    if dtype == 'int4':
      int_category = np.signedinteger
    elif dtype == 'uint4':
      int_category = np.unsignedinteger
    else:
      raise ValueError("Unexpected dtype: {dtype}")
    for dt in [dtype, np.dtype(dtype), str(np.dtype(dtype))]:
      self.assertTrue(dtypes.issubdtype(dt, dt))
      self.assertTrue(dtypes.issubdtype(dt, np.dtype(dtype)))
      self.assertTrue(dtypes.issubdtype(dt, str(np.dtype(dtype))))
      self.assertTrue(dtypes.issubdtype(dt, int_category))
      self.assertTrue(dtypes.issubdtype(dt, np.integer))
      self.assertTrue(dtypes.issubdtype(dt, np.number))
      self.assertTrue(dtypes.issubdtype(dt, np.generic))
      self.assertFalse(dtypes.issubdtype(dt, object))
      self.assertFalse(dtypes.issubdtype(dt, np.int64))
      self.assertFalse(dtypes.issubdtype(np.generic, dt))

  def testArrayCasts(self):
    for t in [jnp.bool_, jnp.int32, jnp.bfloat16, jnp.float32, jnp.complex64]:
      a = np.array([1, 2.5, -3.7])
      self.assertEqual(a.astype(t).dtype, jnp.dtype(t))
      self.assertEqual(jnp.array(a).astype(t).dtype, jnp.dtype(t))

  def testEnumPromotion(self):
    class AnEnum(enum.IntEnum):
      A = 42
      B = 101
    np.testing.assert_equal(np.array(42), np.array(AnEnum.A))
    np.testing.assert_equal(jnp.array(42), jnp.array(AnEnum.A))
    np.testing.assert_equal(np.int32(101), np.int32(AnEnum.B))
    np.testing.assert_equal(jnp.int32(101), jnp.int32(AnEnum.B))

  def testScalarCastInsideJitWorks(self):
    # jnp.int32(tracer) should work.
    self.assertEqual(jnp.int32(101),
                     jax.jit(lambda x: jnp.int32(x))(jnp.float32(101.4)))

  @parameterized.parameters(python_scalar_types)
  def testDtypeFromScalarType(self, typ):
    self.assertEqual(dtypes.dtype(typ), dtypes.python_scalar_dtypes[typ])

  @parameterized.parameters(python_scalar_types)
  def testDtypeFromScalarValue(self, typ):
    self.assertEqual(dtypes.dtype(typ(0)), dtypes.python_scalar_dtypes[typ])

  @parameterized.parameters(all_dtypes)
  def testDtypeFromValue(self, dtype):
    self.assertEqual(dtypes.dtype(dtype.type(0)), dtype)

  @parameterized.parameters(all_dtypes)
  def testDtypeFromDtype(self, dtype):
    self.assertEqual(dtypes.dtype(dtype), dtype)

  @parameterized.parameters(all_dtypes)
  def testDtypeFromString(self, dtype):
    self.assertEqual(dtypes.dtype(str(dtype)), dtype)

  def testDtypeFromNone(self):
    with self.assertRaisesRegex(ValueError, "Invalid argument to dtype"):
      dtypes.dtype(None)

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


class TestPromotionTables(jtu.JaxTestCase):

  @parameterized.named_parameters(
      {"testcase_name": f"_{jaxtype=}", "jaxtype": jaxtype}
      for jaxtype in dtypes._jax_types + dtypes._weak_types)
  def testJaxTypeFromType(self, jaxtype):
    self.assertIs(dtypes._jax_type(*dtypes._dtype_and_weaktype(jaxtype)), jaxtype)

  @parameterized.named_parameters(
      {"testcase_name": f"_{jaxtype=}", "jaxtype": jaxtype}
      for jaxtype in dtypes._jax_types + dtypes._weak_types)
  def testJaxTypeFromVal(self, jaxtype):
    try:
      val = jaxtype(0)
    except TypeError:
      val = jaxtype.type(0)
    self.assertIs(dtypes._jax_type(*dtypes._dtype_and_weaktype(val)), jaxtype)

  @parameterized.named_parameters(
      {"testcase_name": f"_{dtype=}", "dtype": dtype}
      for dtype in dtypes._jax_types)
  def testJaxTypeWeak(self, dtype):
    jax_type = dtypes._jax_type(dtype, weak_type=True)
    if dtypes.issubdtype(jax_type, np.complexfloating):
      self.assertIs(jax_type, complex)
    elif dtypes.issubdtype(jax_type, np.floating):
      self.assertIs(jax_type, float)
    elif dtypes.issubdtype(jax_type, np.integer):
      self.assertIs(jax_type, int)
    else:
      self.assertIs(jax_type, np.dtype(bool))

  @parameterized.named_parameters(
      {"testcase_name": f"_{typ}", "typ": typ}
       for typ in [bool, int, float, complex])
  def testScalarWeakTypes(self, typ):
    # Regression test for https://github.com/google/jax/issues/11377
    val = typ(0)

    result1 = jnp.array(val)
    result2 = jax.jit(jnp.array)(val)
    self.assertEqual(result1.aval, result2.aval)

    with jax.numpy_dtype_promotion('standard'):
      f = lambda x: x / 2
      result1 = jnp.array(f(val))
      result2 = jax.jit(f)(val)
    self.assertEqual(result1.aval, result2.aval)

  def testResultTypeNone(self):
    # This matches the behavior of np.result_type(None) => np.float64
    self.assertEqual(dtypes.result_type(None), dtypes.canonicalize_dtype(dtypes.float_))

  def testResultTypeWeakFlag(self):
    float_ = dtypes.canonicalize_dtype(dtypes.float_)
    x_weak = jnp.array(1.)
    x_strong = x_weak.astype(float_)
    self.assertEqual(dtypes.result_type(x_weak), float_)
    self.assertEqual(dtypes.result_type(x_weak, return_weak_type_flag=True), (float_, True))
    self.assertEqual(dtypes.result_type(x_strong), float_)
    self.assertEqual(dtypes.result_type(x_strong, return_weak_type_flag=True), (float_, False))

  @jtu.ignore_warning(category=UserWarning,
                      message="Explicitly requested dtype.*")
  @jax.numpy_dtype_promotion('standard')
  def testObservedPromotionTable(self):
    """Test that the weak & strong dtype promotion table does not change over time."""
    # Note: * here refers to weakly-typed values
    typecodes = \
        ['b1','u1','u2','u4','u8','i1','i2','i4','i8','bf','f2','f4','f8','c4','c8','i*','f*','c*']
    if config.x64_enabled:
      expected = [
        ['b1','u1','u2','u4','u8','i1','i2','i4','i8','bf','f2','f4','f8','c4','c8','i*','f*','c*'],
        ['u1','u1','u2','u4','u8','i2','i2','i4','i8','bf','f2','f4','f8','c4','c8','u1','f*','c*'],
        ['u2','u2','u2','u4','u8','i4','i4','i4','i8','bf','f2','f4','f8','c4','c8','u2','f*','c*'],
        ['u4','u4','u4','u4','u8','i8','i8','i8','i8','bf','f2','f4','f8','c4','c8','u4','f*','c*'],
        ['u8','u8','u8','u8','u8','f*','f*','f*','f*','bf','f2','f4','f8','c4','c8','u8','f*','c*'],
        ['i1','i2','i4','i8','f*','i1','i2','i4','i8','bf','f2','f4','f8','c4','c8','i1','f*','c*'],
        ['i2','i2','i4','i8','f*','i2','i2','i4','i8','bf','f2','f4','f8','c4','c8','i2','f*','c*'],
        ['i4','i4','i4','i8','f*','i4','i4','i4','i8','bf','f2','f4','f8','c4','c8','i4','f*','c*'],
        ['i8','i8','i8','i8','f*','i8','i8','i8','i8','bf','f2','f4','f8','c4','c8','i8','f*','c*'],
        ['bf','bf','bf','bf','bf','bf','bf','bf','bf','bf','f4','f4','f8','c4','c8','bf','bf','c4'],
        ['f2','f2','f2','f2','f2','f2','f2','f2','f2','f4','f2','f4','f8','c4','c8','f2','f2','c4'],
        ['f4','f4','f4','f4','f4','f4','f4','f4','f4','f4','f4','f4','f8','c4','c8','f4','f4','c4'],
        ['f8','f8','f8','f8','f8','f8','f8','f8','f8','f8','f8','f8','f8','c8','c8','f8','f8','c8'],
        ['c4','c4','c4','c4','c4','c4','c4','c4','c4','c4','c4','c4','c8','c4','c8','c4','c4','c4'],
        ['c8','c8','c8','c8','c8','c8','c8','c8','c8','c8','c8','c8','c8','c8','c8','c8','c8','c8'],
        ['i*','u1','u2','u4','u8','i1','i2','i4','i8','bf','f2','f4','f8','c4','c8','i*','f*','c*'],
        ['f*','f*','f*','f*','f*','f*','f*','f*','f*','bf','f2','f4','f8','c4','c8','f*','f*','c*'],
        ['c*','c*','c*','c*','c*','c*','c*','c*','c*','c4','c4','c4','c8','c4','c8','c*','c*','c*'],
      ]
    else:
      expected = [
        ['b1','u1','u2','u4','u4','i1','i2','i4','i4','bf','f2','f4','f4','c4','c4','i*','f*','c*'],
        ['u1','u1','u2','u4','u4','i2','i2','i4','i4','bf','f2','f4','f4','c4','c4','u1','f*','c*'],
        ['u2','u2','u2','u4','u4','i4','i4','i4','i4','bf','f2','f4','f4','c4','c4','u2','f*','c*'],
        ['u4','u4','u4','u4','u4','i4','i4','i4','i4','bf','f2','f4','f4','c4','c4','u4','f*','c*'],
        ['u4','u4','u4','u4','u4','i4','i4','i4','i4','bf','f2','f4','f4','c4','c4','u4','f*','c*'],
        ['i1','i2','i4','i4','i4','i1','i2','i4','i4','bf','f2','f4','f4','c4','c4','i1','f*','c*'],
        ['i2','i2','i4','i4','i4','i2','i2','i4','i4','bf','f2','f4','f4','c4','c4','i2','f*','c*'],
        ['i4','i4','i4','i4','i4','i4','i4','i4','i4','bf','f2','f4','f4','c4','c4','i4','f*','c*'],
        ['i4','i4','i4','i4','i4','i4','i4','i4','i4','bf','f2','f4','f4','c4','c4','i4','f*','c*'],
        ['bf','bf','bf','bf','bf','bf','bf','bf','bf','bf','f4','f4','f4','c4','c4','bf','bf','c4'],
        ['f2','f2','f2','f2','f2','f2','f2','f2','f2','f4','f2','f4','f4','c4','c4','f2','f2','c4'],
        ['f4','f4','f4','f4','f4','f4','f4','f4','f4','f4','f4','f4','f4','c4','c4','f4','f4','c4'],
        ['f4','f4','f4','f4','f4','f4','f4','f4','f4','f4','f4','f4','f4','c4','c4','f4','f4','c4'],
        ['c4','c4','c4','c4','c4','c4','c4','c4','c4','c4','c4','c4','c4','c4','c4','c4','c4','c4'],
        ['c4','c4','c4','c4','c4','c4','c4','c4','c4','c4','c4','c4','c4','c4','c4','c4','c4','c4'],
        ['i*','u1','u2','u4','u4','i1','i2','i4','i4','bf','f2','f4','f4','c4','c4','i*','f*','c*'],
        ['f*','f*','f*','f*','f*','f*','f*','f*','f*','bf','f2','f4','f4','c4','c4','f*','f*','c*'],
        ['c*','c*','c*','c*','c*','c*','c*','c*','c*','c4','c4','c4','c4','c4','c4','c*','c*','c*'],
      ]
    typecode_to_dtype = {
      'b1': jnp.bool_,
      'u1': jnp.uint8, 'u2': jnp.uint16, 'u4': jnp.uint32, 'u8': jnp.uint64,
      'i1': jnp.int8, 'i2': jnp.int16, 'i4': jnp.int32, 'i8': jnp.int64,
      'bf': jnp.bfloat16, 'f2': jnp.float16, 'f4': jnp.float32, 'f8': jnp.float64,
      'c4': jnp.complex64, 'c8': jnp.complex128,
      'i*': jnp.int64, 'f*': jnp.float64, 'c*': jnp.complex128,
    }
    dtype_to_typecode = {jnp.dtype(v): k for k, v in typecode_to_dtype.items()
                        if not k.endswith('*')}

    def typecode_to_val(typecode):
      weak_type = typecode.endswith('*')
      dtype = typecode_to_dtype[typecode]
      val = dtype(0)
      if weak_type:
        val = val.item()
      return val

    def val_to_typecode(val):
      dtype = dtypes.result_type(val)
      weak_type = dtypes.is_weakly_typed(val)
      typecode = dtype_to_typecode[dtype]
      if weak_type:
        typecode = typecode[:-1] + '*'
      return typecode

    vals = [typecode_to_val(t) for t in typecodes]
    table = [[val_to_typecode(v1 + v2) for v1 in vals] for v2 in vals]

    def show_differences(expected, actual):
      diffs = ""
      for i, t1 in enumerate(typecodes):
        for j, t2 in enumerate(typecodes):
          if expected[i][j] != actual[i][j]:
            diffs += f"\n{t1}, {t2} -> want {expected[i][j]}, got {actual[i][j]}"
      return diffs

    self.assertEqual(table, expected, show_differences(expected, table))

  @parameterized.named_parameters(
    {"testcase_name": "_xtype={}_ytype={}_xfun={}_yfun={}".format(
      xtype.__name__, ytype.__name__, xfun.__name__, yfun.__name__),
     "xtype": xtype, "ytype": ytype, "xfun": xfun, "yfun": yfun}
    for xtype, ytype in itertools.product(
      [int, float, jnp.int16, jnp.int32, jnp.float16, jnp.float32], repeat=2)
    for xfun, yfun in itertools.product(
      [identity, abs, jnp.array], repeat=2)
    )
  @jax.numpy_dtype_promotion('standard')
  def testBinaryPromotionJitInvariance(self, xtype, ytype, xfun, yfun):
    """Test jit invariance of simple binary promotion rules with and without weak types."""
    f = lambda x, y: xfun(x) + yfun(y)
    args_maker = lambda: [xtype(1), ytype(1)]
    self._CompileAndCheck(f, args_maker, check_dtypes=True)

  @parameterized.named_parameters(
    {"testcase_name": f"_{dtype=}_{weak_type=}",
     "dtype": dtype, "weak_type": weak_type}
    for dtype in all_dtypes
    for weak_type in [True, False]
  )
  def testUnaryPromotion(self, dtype, weak_type):
    # Regression test for https://github.com/google/jax/issues/6051
    if _int4_enabled and dtype in int4_dtypes:
      self.skipTest("XLA support for int4 is incomplete.")
    x = lax_internal._convert_element_type(0, dtype, weak_type=weak_type)
    if weak_type:
      if _fp8_enabled:
        expected = dtypes.canonicalize_dtype(
          dtypes._default_types['f' if x.dtype in ["bfloat16", *fp8_dtypes] else x.dtype.kind])
      else:
        expected = dtypes.canonicalize_dtype(
          dtypes._default_types['f' if x.dtype == 'bfloat16' else x.dtype.kind])
    else:
      expected = x.dtype
    self.assertEqual(dtypes.result_type(x), expected)

  @unittest.skipIf(not _fp8_enabled, "requires fp8 dtypes")
  @jax.numpy_dtype_promotion('standard')
  def testFloat8PromotionError(self):
    for dtype in fp8_dtypes:
      x = jnp.array(1, dtype=dtype)
      y = jnp.array(1, dtype='float32')
      with self.assertRaisesRegex(dtypes.TypePromotionError,
                                  ".*8-bit floats do not support implicit promotion"):
        x + y

  @unittest.skipIf(not _int4_enabled, "requires int4 dtypes")
  @jax.numpy_dtype_promotion('standard')
  def testInt4PromotionError(self):
    for dtype in int4_dtypes:
      x = jnp.array(1, dtype=dtype)
      y = jnp.array(1, dtype='int32')
      with self.assertRaisesRegex(dtypes.TypePromotionError,
                                  ".*4-bit integers do not support implicit promotion"):
        x + y

  @jtu.sample_product(
    dtype=all_dtypes,
    weak_type=[True, False],
    promotion=['standard', 'strict'],
  )
  def testBinaryNonPromotion(self, dtype, weak_type, promotion):
    if _fp8_enabled and dtype in fp8_dtypes:
      self.skipTest("XLA support for float8 is incomplete.")
    if _int4_enabled and dtype in int4_dtypes:
      self.skipTest("XLA support for int4 is incomplete.")
    # Regression test for https://github.com/google/jax/issues/6051
    x = lax_internal._convert_element_type(0, dtype, weak_type=weak_type)
    with jax.numpy_dtype_promotion(promotion):
      y = (x + x)

    if promotion == 'standard' or not weak_type or dtype == dtypes.bool_:
      expected_dtype = dtype
    elif dtypes.issubdtype(dtype, np.complexfloating):
      expected_dtype = np.complex128
    elif dtypes.issubdtype(dtype, np.floating):
      expected_dtype = np.float64
    else:
      expected_dtype = np.int64

    # No boolean weak types.
    expected_weak_type = weak_type and dtype != bool
    expected_dtype = dtypes.canonicalize_dtype(expected_dtype)

    self.assertEqual(y.dtype, expected_dtype)
    self.assertEqual(dtypes.is_weakly_typed(y), expected_weak_type)

  @parameterized.named_parameters(
    {"testcase_name": f"_{dtype=}_{weak_type=}",
     "dtype": dtype, "weak_type": weak_type}
    for dtype in all_dtypes
    for weak_type in [True, False]
  )
  def testArrayRepr(self, dtype, weak_type):
    val = lax_internal._convert_element_type(0, dtype, weak_type=weak_type)
    rep = repr(val)
    self.assertStartsWith(rep, 'Array(')
    if weak_type:
      self.assertEndsWith(rep, f"dtype={val.dtype.name}, weak_type=True)")
    else:
      self.assertEndsWith(rep, f"dtype={val.dtype.name})")

  @jtu.sample_product(
      input_dtype=jtu.dtypes.all + [bool, int, float, complex],
      output_dtype=jtu.dtypes.all + [bool, int, float, complex],
      numpy_dtype_promotion=['strict', 'standard']
  )
  def testSafeToCast(self, input_dtype, output_dtype, numpy_dtype_promotion):
    print(input_dtype, output_dtype)
    with jax.numpy_dtype_promotion(numpy_dtype_promotion):
      # First the special cases which are always safe:
      always_safe = (
        (input_dtype == output_dtype) or
        (dtypes.issubdtype(output_dtype, np.integer) and input_dtype in {int}) or
        (dtypes.issubdtype(output_dtype, np.floating) and input_dtype in {int, float}) or
        (dtypes.issubdtype(output_dtype, np.complexfloating) and input_dtype in {int, float, complex})
      )
      if always_safe:
        self.assertTrue(dtypes.safe_to_cast(input_dtype, output_dtype))

      try:
        result_dtype = dtypes.result_type(input_dtype, dtypes.canonicalize_dtype(output_dtype))
      except dtypes.TypePromotionError:
        result_dtype = None

      if result_dtype is None and input_dtype != output_dtype:
        with self.assertRaises(dtypes.TypePromotionError):
          dtypes.safe_to_cast(input_dtype, output_dtype)
      else:
        self.assertEqual(dtypes.result_type(output_dtype) == result_dtype,
                         dtypes.safe_to_cast(input_dtype, output_dtype))


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