Repository: google/jax

Issue:
Numpy window functions produce large constants
Affected functions are:

- `jnp.blackman`
- `jnp.bartlett`
- `jnp.hamming`
- `jnp.hanning`
- `jnp.kaiser`

For example:
```python
import jax
import jax.numpy as jnp
from functools import partial
print(jax.jit(partial(jnp.hamming, 100)).lower().as_text())
```
```
module @jit__unnamed_wrapped_function_ {
  func.func public @main() -> tensor<100xf32> {
    %0 = mhlo.constant dense<[8.000000e-02, 0.0809261277, 0.0837007835, 0.0883127972, 0.0947435945, 0.102967285, 0.11295075, 0.124653794, 0.138029292, 0.153023377, 0.169575676, 0.187619552, 0.207082346, 0.227885664, 0.249945775, 0.273173809, 0.297476292, 0.322755307, 0.34890908, 0.375832349, 0.403416634, 0.43155089, 0.46012184, 0.489014417, 0.518112302, 0.547298372, 5.764550e-01, 0.605464816, 0.634211063, 0.662577927, 0.690451264, 0.71771878, 0.744270622, 0.76999998, 0.794803202, 0.818580448, 0.841235935, 0.862678468, 0.882821619, 0.901584446, 0.918891251, 0.934672355, 0.948864281, 0.961409866, 0.972258627, 0.981366753, 0.988697707, 0.994221866, 0.997917056, 0.999768435, 0.999768435, 0.997917056, 0.994221866, 0.988697707, 0.981366753, 0.972258627, 0.961409866, 0.948864281, 0.934672355, 0.918891251, 0.901584446, 0.882821619, 0.862678468, 0.841235935, 0.818580448, 0.794803202, 0.76999998, 0.744270622, 0.71771878, 0.690451264, 0.662577927, 0.634211063, 0.605464816, 5.764550e-01, 0.547298372, 0.518112302, 0.489014417, 0.46012184, 0.43155089, 0.403416634, 0.375832349, 0.34890908, 0.322755307, 0.297476292, 0.273173809, 0.249945775, 0.227885664, 0.207082346, 0.187619552, 0.169575676, 0.153023377, 0.138029292, 0.124653794, 0.11295075, 0.102967285, 0.0947435945, 0.0883127972, 0.0837007835, 0.0809261277, 8.000000e-02]> : tensor<100xf32>
    return %0 : tensor<100xf32>
  }
}
```


Discussion:

Resolving patch:
diff --git a/jax/_src/numpy/lax_numpy.py b/jax/_src/numpy/lax_numpy.py
--- a/jax/_src/numpy/lax_numpy.py
+++ b/jax/_src/numpy/lax_numpy.py
@@ -2259,20 +2259,6 @@ def arange(start: core.DimSize, stop: Optional[core.DimSize] = None,
     return array(np.arange(start, stop=stop, step=step, dtype=dtype))
 
 
-def _wrap_numpy_nullary_function(f):
-  """Adapts `f` to return a DeviceArray instead of an np.ndarray.
-
-  `f` cannot have any non-static array arguments.
-  """
-  @_wraps(f, update_doc=False)
-  def wrapper(*args, **kwargs):
-    args = [core.concrete_or_error(None, arg, f"the error occurred in argument {i} jnp.{f.__name__}()")
-            for i, arg in enumerate(args)]
-    kwargs = {key: core.concrete_or_error(None, val, f"the error occurred in argument '{key}' jnp.{f.__name__}()")
-              for key, val in kwargs.items()}
-    return asarray(f(*args, **kwargs))
-  return wrapper
-
 @overload
 def linspace(start: ArrayLike, stop: ArrayLike, num: int = 50,
              endpoint: bool = True, retstep: Literal[False] = False,
@@ -4281,12 +4267,56 @@ def _static_idx(idx: slice, size: core.DimSize):
     return stop + k + 1, start + 1, -step, True
 
 
-blackman = _wrap_numpy_nullary_function(np.blackman)
-bartlett = _wrap_numpy_nullary_function(np.bartlett)
-hamming = _wrap_numpy_nullary_function(np.hamming)
-hanning = _wrap_numpy_nullary_function(np.hanning)
-# TODO: lower `kaiser` via lax to allow non-constant beta values.
-kaiser = _wrap_numpy_nullary_function(np.kaiser)
+@_wraps(np.blackman)
+def blackman(M: int) -> Array:
+  M = core.concrete_or_error(int, M, "M argument of jnp.blackman")
+  dtype = dtypes.canonicalize_dtype(float_)
+  if M <= 1:
+    return ones(M, dtype)
+  n = lax.iota(dtype, M)
+  return 0.42 - 0.5 * cos(2 * pi * n / (M - 1)) + 0.08 * cos(4 * pi * n / (M - 1))
+
+
+@_wraps(np.bartlett)
+def bartlett(M: int) -> Array:
+  M = core.concrete_or_error(int, M, "M argument of jnp.bartlett")
+  dtype = dtypes.canonicalize_dtype(float_)
+  if M <= 1:
+    return ones(M, dtype)
+  n = lax.iota(dtype, M)
+  return 1 - abs(2 * n + 1 - M) / (M - 1)
+
+
+@_wraps(np.hamming)
+def hamming(M: int) -> Array:
+  M = core.concrete_or_error(int, M, "M argument of jnp.hamming")
+  dtype = dtypes.canonicalize_dtype(float_)
+  if M <= 1:
+    return ones(M, dtype)
+  n = lax.iota(dtype, M)
+  return 0.54 - 0.46 * cos(2 * pi * n / (M - 1))
+
+
+@_wraps(np.hanning)
+def hanning(M: int) -> Array:
+  M = core.concrete_or_error(int, M, "M argument of jnp.hanning")
+  dtype = dtypes.canonicalize_dtype(float_)
+  if M <= 1:
+    return ones(M, dtype)
+  n = lax.iota(dtype, M)
+  return 0.5 * (1 - cos(2 * pi * n / (M - 1)))
+
+
+@_wraps(np.kaiser)
+def kaiser(M: int, beta: ArrayLike) -> Array:
+  M = core.concrete_or_error(int, M, "M argument of jnp.kaiser")
+  dtype = dtypes.canonicalize_dtype(float_)
+  if M <= 1:
+    return ones(M, dtype)
+  n = lax.iota(dtype, M)
+  alpha = 0.5 * (M - 1)
+  return i0(beta * sqrt(1 - ((n - alpha) / alpha) ** 2)) / i0(beta)
+
 
 def _gcd_cond_fn(xs):
   x1, x2 = xs


Test patch:
diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py
--- a/tests/lax_numpy_test.py
+++ b/tests/lax_numpy_test.py
@@ -2440,6 +2440,20 @@ def testHVDStack(self, shape, op, dtypes, array_input, out_dtype):
       self._CheckAgainstNumpy(jnp_fun, np_fun, args_maker)
       self._CompileAndCheck(jnp_fun, args_maker)
 
+  @jtu.sample_product(
+    [dict(name=name, **kwds)
+      for name in ['blackman', 'bartlett', 'hamming', 'hanning', 'kaiser']
+      for kwds in ([dict(beta=1), dict(beta=0.5)] if name == 'kaiser' else [{}])
+    ],
+    size = [0, 1, 5, 10],
+  )
+  def testWindowFunction(self, name, size, **kwds):
+    jnp_fun = partial(getattr(jnp, name), size, **kwds)
+    np_fun = partial(getattr(np, name), size, **kwds)
+    args_maker = lambda: []
+    self._CheckAgainstNumpy(jnp_fun, np_fun, args_maker)
+    self._CompileAndCheck(jnp_fun, args_maker)
+
   @jtu.sample_product(
     [dict(shape=shape, fill_value_shape=fill_value_shape)
       for shape in array_shapes + [3, np.array(7, dtype=np.int32)]