Repository: google/jax
Pull request #21853 (2024-06-13 17:18:41)

jax.dtypes: avoid erroring on non-hashable dtype
Part of addressing [#21824](https://github.com/google/jax/issues/21824)


Patch:
diff --git a/jax/_src/dtypes.py b/jax/_src/dtypes.py
index b4c9cfdf5..0691d910b 100644
--- a/jax/_src/dtypes.py
+++ b/jax/_src/dtypes.py
@@ -752,14 +752,14 @@ def check_user_dtype_supported(dtype, fun_name=None):
     msg = f"JAX only supports number and bool dtypes, got dtype {dtype}"
     msg += f" in {fun_name}" if fun_name else ""
     raise TypeError(msg)
-  if dtype is not None and np_dtype != canonicalize_dtype(dtype):
+  if dtype is not None and np_dtype != canonicalize_dtype(np_dtype):
     msg = ("Explicitly requested dtype {} {} is not available, "
            "and will be truncated to dtype {}. To enable more dtypes, set the "
            "jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell "
            "environment variable. "
            "See https://github.com/google/jax#current-gotchas for more.")
     fun_name = f"requested in {fun_name}" if fun_name else ""
-    truncated_dtype = canonicalize_dtype(dtype).name
+    truncated_dtype = canonicalize_dtype(np_dtype).name
     warnings.warn(msg.format(dtype, fun_name, truncated_dtype), stacklevel=3)
 
 def safe_to_cast(input_dtype_or_value: Any,
diff --git a/jax/_src/numpy/lax_numpy.py b/jax/_src/numpy/lax_numpy.py
index ca8c26051..980e31b0b 100644
--- a/jax/_src/numpy/lax_numpy.py
+++ b/jax/_src/numpy/lax_numpy.py
@@ -397,161 +397,23 @@ def _conv(x: Array, y: Array, mode: str, op: str, precision: PrecisionLike,
   return result[0, 0, out_order]
 
 
+@util.implements(np.convolve, lax_description=_PRECISION_DOC,
+                 extra_params=_CONV_PREFERRED_ELEMENT_TYPE_DESCRIPTION)
 @partial(jit, static_argnames=('mode', 'precision', 'preferred_element_type'))
 def convolve(a: ArrayLike, v: ArrayLike, mode: str = 'full', *,
              precision: PrecisionLike = None,
              preferred_element_type: DTypeLike | None = None) -> Array:
-  r"""Convolution of two one dimensional arrays.
-
-  JAX implementation of :func:`numpy.convolve`.
-
-  Convolution of one dimensional arrays is defined as:
-
-  .. math::
-
-     c_k = \sum_j a_{k - j} v_j
-
-  Args:
-    a: left-hand input to the convolution. Must have ``a.ndim == 1``.
-    v: right-hand input to the convolution. Must have ``v.ndim == 1``.
-    mode: controls the size of the output. Available operations are:
-
-      * ``"full"``: (default) output the full convolution of the inputs.
-      * ``"same"``: return a centered portion of the ``"full"`` output which
-        is the same size as ``a``.
-      * ``"valid"``: return the portion of the ``"full"`` output which do not
-        depend on padding at the array edges.
-
-    precision: Specify the precision of the computation. Refer to
-      :class:`jax.lax.Precision` for a description of available values.
-
-    preferred_element_type: A datatype, indicating to accumulate results to and
-      return a result with that datatype. Default is ``None``, which means the
-      default accumulation type for the input types.
-
-  Returns:
-    Array containing the convolved result.
-
-  See Also:
-    - :func:`jax.scipy.signal.convolve`: ND convolution
-    - :func:`jax.numpy.correlate`: 1D correlation
-
-  Examples:
-    A few 1D convolution examples:
-
-    >>> x = jnp.array([1, 2, 3, 2, 1])
-    >>> y = jnp.array([4, 1, 2])
-
-    ``jax.numpy.convolve``, by default, returns full convolution using implicit
-    zero-padding at the edges:
-
-    >>> jnp.convolve(x, y)
-    Array([ 4.,  9., 16., 15., 12.,  5.,  2.], dtype=float32)
-
-    Specifying ``mode = 'same'`` returns a centered convolution the same size
-    as the first input:
-
-    >>> jnp.convolve(x, y, mode='same')
-    Array([ 9., 16., 15., 12.,  5.], dtype=float32)
-
-    Specifying ``mode = 'valid'`` returns only the portion where the two arrays
-    fully overlap:
-
-    >>> jnp.convolve(x, y, mode='valid')
-    Array([16., 15., 12.], dtype=float32)
-
-    For complex-valued inputs:
-
-    >>> x1 = jnp.array([3+1j, 2, 4-3j])
-    >>> y1 = jnp.array([1, 2-3j, 4+5j])
-    >>> jnp.convolve(x1, y1)
-    Array([ 3. +1.j, 11. -7.j, 15.+10.j,  7. -8.j, 31. +8.j], dtype=complex64)
-  """
   util.check_arraylike("convolve", a, v)
   return _conv(asarray(a), asarray(v), mode=mode, op='convolve',
                precision=precision, preferred_element_type=preferred_element_type)
 
 
+@util.implements(np.correlate, lax_description=_PRECISION_DOC,
+                 extra_params=_CONV_PREFERRED_ELEMENT_TYPE_DESCRIPTION)
 @partial(jit, static_argnames=('mode', 'precision', 'preferred_element_type'))
 def correlate(a: ArrayLike, v: ArrayLike, mode: str = 'valid', *,
               precision: PrecisionLike = None,
               preferred_element_type: DTypeLike | None = None) -> Array:
-  r"""Correlation of two one dimensional arrays.
-
-  JAX implementation of :func:`numpy.correlate`.
-
-  Correlation of one dimensional arrays is defined as:
-
-  .. math::
-
-     c_k = \sum_j a_{k + j} \overline{v_j}
-
-  where :math:`\overline{v_j}` is the complex conjugate of :math:`v_j`.
-
-  Args:
-    a: left-hand input to the correlation. Must have ``a.ndim == 1``.
-    v: right-hand input to the correlation. Must have ``v.ndim == 1``.
-    mode: controls the size of the output. Available operations are:
-
-      * ``"full"``: output the full correlation of the inputs.
-      * ``"same"``: return a centered portion of the ``"full"`` output which
-        is the same size as ``a``.
-      * ``"valid"``: (default) return the portion of the ``"full"`` output which do not
-        depend on padding at the array edges.
-
-    precision: Specify the precision of the computation. Refer to
-      :class:`jax.lax.Precision` for a description of available values.
-
-    preferred_element_type: A datatype, indicating to accumulate results to and
-      return a result with that datatype. Default is ``None``, which means the
-      default accumulation type for the input types.
-
-  Returns:
-    Array containing the cross-correlation result.
-
-  See Also:
-    - :func:`jax.scipy.signal.correlate`: ND correlation
-    - :func:`jax.numpy.convolve`: 1D convolution
-
-  Examples:
-    >>> x = jnp.array([1, 2, 3, 2, 1])
-    >>> y = jnp.array([4, 5, 6])
-
-    Since default ``mode = 'valid'``, ``jax.numpy.correlate`` returns only the
-    portion of correlation where the two arrays fully overlap:
-
-    >>> jnp.correlate(x, y)
-    Array([32., 35., 28.], dtype=float32)
-
-    Specifying ``mode = 'full'`` returns full correlation using implicit
-    zero-padding at the edges.
-
-    >>> jnp.correlate(x, y, mode='full')
-    Array([ 6., 17., 32., 35., 28., 13.,  4.], dtype=float32)
-
-    Specifying ``mode = 'same'`` returns a centered correlation the same size
-    as the first input:
-
-    >>> jnp.correlate(x, y, mode='same')
-    Array([17., 32., 35., 28., 13.], dtype=float32)
-
-    If both the inputs arrays are real-valued and symmetric then the result will
-    also be symmetric and will be equal to the result of ``jax.numpy.convolve``.
-
-    >>> x1 = jnp.array([1, 2, 3, 2, 1])
-    >>> y1 = jnp.array([4, 5, 4])
-    >>> jnp.correlate(x1, y1, mode='full')
-    Array([ 4., 13., 26., 31., 26., 13.,  4.], dtype=float32)
-    >>> jnp.convolve(x1, y1, mode='full')
-    Array([ 4., 13., 26., 31., 26., 13.,  4.], dtype=float32)
-
-    For complex-valued inputs:
-
-    >>> x2 = jnp.array([3+1j, 2, 2-3j])
-    >>> y2 = jnp.array([4, 2-5j, 1])
-    >>> jnp.correlate(x2, y2, mode='full')
-    Array([ 3. +1.j,  3.+17.j, 18.+11.j, 27. +4.j,  8.-12.j], dtype=complex64)
-  """
   util.check_arraylike("correlate", a, v)
   return _conv(asarray(a), asarray(v), mode=mode, op='correlate',
                precision=precision, preferred_element_type=preferred_element_type)
diff --git a/jax/experimental/mosaic/gpu/examples/flash_attention.py b/jax/experimental/mosaic/gpu/examples/flash_attention.py
index 99590168c..4246fe3cd 100644
--- a/jax/experimental/mosaic/gpu/examples/flash_attention.py
+++ b/jax/experimental/mosaic/gpu/examples/flash_attention.py
@@ -47,7 +47,7 @@ _utils_c = c
 
 
 # TODO(apaszke): Implement a Q-scaled, base2 exp implementation.
-class ExpImplementation(enum.Enum):
+class ExpImplementation(enum.StrEnum):
   EXACT = enum.auto()
   APPROX = enum.auto()
 
@@ -371,6 +371,6 @@ if __name__ == "__main__":
     optimal_time = matmul_flops / peak_flops * 1e6  # us
     achieved_tc_util = optimal_time / runtime_us * 100
     print(
-        f"{kv_seq_len=:<6} {q_seq_len=:<6} {num_q_heads=:<4} {head_dim=:<6} exp_impl={exp_impl.name:<6}:"
+        f"{kv_seq_len=:<6} {q_seq_len=:<6} {num_q_heads=:<4} {head_dim=:<6} exp_impl={str(exp_impl):<6}:"
         f"  {runtime_us:<7.1f}us = {achieved_tc_util:4.1f}% TC utilization"
     )
diff --git a/setup.py b/setup.py
index d5cc11059..7f0823b0d 100644
--- a/setup.py
+++ b/setup.py
@@ -95,13 +95,6 @@ setup(
           f"jax-cuda12-plugin[with_cuda]>={_current_jaxlib_version},<={_jax_version}",
         ],
 
-        # Deprecated alias for cuda12, kept to avoid breaking users who wrote
-        # cuda12_pip in their CI.
-        'cuda12_pip': [
-          f"jaxlib=={_current_jaxlib_version}",
-          f"jax-cuda12-plugin[with_cuda]>={_current_jaxlib_version},<={_jax_version}",
-        ],
-
         # Target that does not depend on the CUDA pip wheels, for those who want
         # to use a preinstalled CUDA.
         'cuda12_local': [


Test patch:
diff --git a/tests/dtypes_test.py b/tests/dtypes_test.py
index 668ea4c0b..0a29e6168 100644
--- a/tests/dtypes_test.py
+++ b/tests/dtypes_test.py
@@ -558,6 +558,13 @@ class DtypesTest(jtu.JaxTestCase):
     _, new_scale = jax.jit(jax.grad(outer, (0, 1)))(jnp.float32(3.14), scale)
     self.assertAllClose(new_scale, jnp.float32(1.0))
 
+  def test_check_dtype_non_hashable(self):
+    # regression test for issue with checking non-hashable custom dtype
+    class MyDtype:
+      __hash__ = None
+      dtype = np.dtype('float32')
+    dtypes.check_user_dtype_supported(MyDtype())
+
 class EArrayTest(jtu.JaxTestCase):
 
   @parameterized.parameters([True, False])