Repository: google/jax
Pull request #22571 (2024-07-22 17:43:54)

jax.scipy.fft.dct: implement & test norm='backward'
Related to #22466


Patch:
diff --git a/jax/_src/scipy/fft.py b/jax/_src/scipy/fft.py
index 11de07dbb..a826d4746 100644
--- a/jax/_src/scipy/fft.py
+++ b/jax/_src/scipy/fft.py
@@ -55,7 +55,8 @@ def dct(x: Array, type: int = 2, n: int | None = None,
       If larger than ``x.shape[axis]``, the input will be zero-padded, if
       smaller, the input will be truncated.
     axis: integer, default=-1. The axis along which the dct will be performed.
-    norm: string. The normalization mode. Currently only ``"ortho"`` is supported.
+    norm: string. The normalization mode: one of ``[None, "backward", "ortho"]``.
+      The default is ``None``, which is equivalent to ``"backward"``.
 
   Returns:
     array containing the discrete cosine transform of x
@@ -98,7 +99,7 @@ def dct(x: Array, type: int = 2, n: int | None = None,
   """
   if type != 2:
     raise NotImplementedError('Only DCT type 2 is implemented.')
-  if norm is not None and norm not in ['ortho']:
+  if norm is not None and norm not in ['backward', 'ortho']:
     raise ValueError(f"jax.scipy.fft.dct: {norm=!r} is not implemented")
 
   axis = canonicalize_axis(axis, x.ndim)
@@ -149,7 +150,8 @@ def dctn(x: Array, type: int = 2,
       specified, it will default to the shape of ``x`` along the specified ``axes``.
     axes: integer or sequence of integers. Specifies the axes along which the
       transform will be computed.
-    norm: string. The normalization mode. Currently only ``"ortho"`` is supported.
+    norm: string. The normalization mode: one of ``[None, "backward", "ortho"]``.
+      The default is ``None``, which is equivalent to ``"backward"``.
 
   Returns:
     array containing the discrete cosine transform of x
@@ -198,7 +200,7 @@ def dctn(x: Array, type: int = 2,
   """
   if type != 2:
     raise NotImplementedError('Only DCT type 2 is implemented.')
-  if norm is not None and norm not in ['ortho']:
+  if norm is not None and norm not in ['backward', 'ortho']:
     raise ValueError(f"jax.scipy.fft.dctn: {norm=!r} is not implemented")
 
   if axes is None:
@@ -234,7 +236,8 @@ def idct(x: Array, type: int = 2, n: int | None = None,
       If larger than ``x.shape[axis]``, the input will be zero-padded, if
       smaller, the input will be truncated.
     axis: integer, default=-1. The axis along which the dct will be performed.
-    norm: string. The normalization mode. Currently only ``"ortho"`` is supported.
+    norm: string. The normalization mode: one of ``[None, "backward", "ortho"]``.
+      The default is ``None``, which is equivalent to ``"backward"``.
 
   Returns:
     array containing the inverse discrete cosine transform of x
@@ -286,7 +289,7 @@ def idct(x: Array, type: int = 2, n: int | None = None,
   """
   if type != 2:
     raise NotImplementedError('Only DCT type 2 is implemented.')
-  if norm is not None and norm not in ['ortho']:
+  if norm is not None and norm not in ['backward', 'ortho']:
     raise ValueError(f"jax.scipy.fft.idct: {norm=!r} is not implemented")
 
   axis = canonicalize_axis(axis, x.ndim)
@@ -296,7 +299,7 @@ def idct(x: Array, type: int = 2, n: int | None = None,
                  for a in range(x.ndim)])
   N = x.shape[axis]
   x = x.astype(jnp.float32)
-  if norm is None:
+  if norm is None or norm == 'backward':
     x = _dct_ortho_norm(x, axis)
   x = _dct_ortho_norm(x, axis)
 
@@ -328,7 +331,8 @@ def idctn(x: Array, type: int = 2,
       specified, it will default to the shape of ``x`` along the specified ``axes``.
     axes: integer or sequence of integers. Specifies the axes along which the
       transform will be computed.
-    norm: string. The normalization mode. Currently only ``"ortho"`` is supported.
+    norm: string. The normalization mode: one of ``[None, "backward", "ortho"]``.
+      The default is ``None``, which is equivalent to ``"backward"``.
 
   Returns:
     array containing the inverse discrete cosine transform of x
@@ -384,7 +388,7 @@ def idctn(x: Array, type: int = 2,
   """
   if type != 2:
     raise NotImplementedError('Only DCT type 2 is implemented.')
-  if norm is not None and norm not in ['ortho']:
+  if norm is not None and norm not in ['backward', 'ortho']:
     raise ValueError(f"jax.scipy.fft.idctn: {norm=!r} is not implemented")
 
   if axes is None:


Test patch:
diff --git a/tests/scipy_fft_test.py b/tests/scipy_fft_test.py
index 17c1e9c2d..6c549f5ed 100644
--- a/tests/scipy_fft_test.py
+++ b/tests/scipy_fft_test.py
@@ -50,7 +50,7 @@ class LaxBackedScipyFftTests(jtu.JaxTestCase):
     shape=[(10,), (2, 5)],
     n=[None, 1, 7, 13, 20],
     axis=[-1, 0],
-    norm=[None, 'ortho'],
+    norm=[None, 'ortho', 'backward'],
   )
   def testDct(self, shape, dtype, n, axis, norm):
     rng = jtu.rand_default(self.rng())
@@ -67,7 +67,7 @@ class LaxBackedScipyFftTests(jtu.JaxTestCase):
      for axes in _get_dctn_test_axes(shape)
      for s in _get_dctn_test_s(shape, axes)],
     dtype=real_dtypes,
-    norm=[None, 'ortho'],
+    norm=[None, 'ortho', 'backward'],
   )
   def testDctn(self, shape, dtype, s, axes, norm):
     rng = jtu.rand_default(self.rng())
@@ -83,7 +83,7 @@ class LaxBackedScipyFftTests(jtu.JaxTestCase):
     shape=[(10,), (2, 5)],
     n=[None, 1, 7, 13, 20],
     axis=[-1, 0],
-    norm=[None, 'ortho'],
+    norm=[None, 'ortho', 'backward'],
   )
   # TODO(phawkins): these tests are failing on T4 GPUs in CI with a
   # CUDA_ERROR_ILLEGAL_ADDRESS.
@@ -103,7 +103,7 @@ class LaxBackedScipyFftTests(jtu.JaxTestCase):
      for axes in _get_dctn_test_axes(shape)
      for s in _get_dctn_test_s(shape, axes)],
     dtype=real_dtypes,
-    norm=[None, 'ortho'],
+    norm=[None, 'ortho', 'backward'],
   )
   # TODO(phawkins): these tests are failing on T4 GPUs in CI with a
   # CUDA_ERROR_ILLEGAL_ADDRESS.