Repository: jax-ml/jax (python)

Issue:
scipy.signal.istft error with numpy window
### Description

`jax.signal.istft` raises `ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()` when the `window` argument is set to a numpy array.

```python
import jax.numpy as jnp
import numpy as np 

Zxx = jnp.ones((257, 41))
nperseg = 512
window = np.ones(nperseg) 
jax.scipy.signal.istft(Zxx, fs=1024, nperseg=nperseg, window=window)
```

Changing [this line](https://github.com/jax-ml/jax/blob/2ebb756056396e825c32b493ea242cd6c9db0765/jax/_src/scipy/signal.py#L1150) to `if isinstance(window, str) and window == 'hann':` instead of attempting to compare the array directly to a string should fix it.


### System info (python version, jaxlib version, accelerator, etc.)

```
jax:    0.6.0
jaxlib: 0.6.0
numpy:  2.2.5
python: 3.10.14 (main, Aug 14 2024, 05:11:29) [Clang 18.1.8 ]
device info: NVIDIA GeForce GTX 1660 Ti-1, 1 local devices"
process_count: 1
platform: uname_result(system='Linux', node='Legion-5', release='6.11.0-21-generic', version='#21~24.04.1-Ubuntu SMP PREEMPT_DYNAMIC Mon Feb 24 16:52:15 UTC 2', machine='x86_64')


$ nvidia-smi
Wed May  7 11:44:14 2025       
+---------------------------------------------------------------------------------------+
| NVIDIA-SMI 535.247.01             Driver Version: 535.247.01   CUDA Version: 12.2     |
|-----------------------------------------+----------------------+----------------------+
| GPU  Name                 Persistence-M | Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp   Perf          Pwr:Usage/Cap |         Memory-Usage | GPU-Util  Compute M. |
|                                         |                      |               MIG M. |
|=========================================+======================+======================|
|   0  NVIDIA GeForce GTX 1660 Ti     On  | 00000000:01:00.0  On |                  N[/](https://file+.vscode-resource.vscode-cdn.net/)A |
| N/A   44C    P8               6W [/](https://file+.vscode-resource.vscode-cdn.net/)  80W |   4599MiB /  6144MiB |     17%      Default |
|                                         |                      |                  N/A |
+-----------------------------------------+----------------------+----------------------+
                                                                                         
+---------------------------------------------------------------------------------------+
| Processes:                                                                            |
|  GPU   GI   CI        PID   Type   Process name                            GPU Memory |
|        ID   ID                                                             Usage      |
|=======================================================================================|
|    0   N/A  N/A      2572      G   [/usr/lib/xorg/Xorg](https://file+.vscode-resource.vscode-cdn.net/usr/lib/xorg/Xorg)                           45MiB |
|    0   N/A  N/A     27839      C   .../multivarie_python/.venv/bin/python     4550MiB |
+---------------------------------------------------------------------------------------+
```

Pull request description:
jax.scipy.signal.istft: support array input for window
Fixes #28571

Patch:
diff --git a/jax/_src/dtypes.py b/jax/_src/dtypes.py
index d1e5b7bf4..ae3516ea6 100644
--- a/jax/_src/dtypes.py
+++ b/jax/_src/dtypes.py
@@ -279,6 +279,12 @@ def to_inexact_dtype(dtype: DTypeLike) -> DType:
   return _dtype_to_inexact.get(dtype_, dtype_)
 
 
+def to_floating_dtype(dtype: DTypeLike) -> DType:
+  """Promotes a dtype to a non-complex floating dtype."""
+  dtype_ = np.dtype(dtype)
+  return finfo(_dtype_to_inexact.get(dtype_, dtype_)).dtype
+
+
 def to_complex_dtype(dtype: DTypeLike) -> DType:
   ftype = to_inexact_dtype(dtype)
   if ftype in [np.dtype('float64'), np.dtype('complex128')]:
diff --git a/jax/_src/scipy/signal.py b/jax/_src/scipy/signal.py
index d950cd2ea..565909e8a 100644
--- a/jax/_src/scipy/signal.py
+++ b/jax/_src/scipy/signal.py
@@ -1071,7 +1071,7 @@ def istft(Zxx: Array, fs: ArrayLike = 1.0, window: str = 'hann',
     noverlap: Number of points to overlap between segments (default: ``nperseg // 2``).
     nfft: Number of FFT points used in the STFT. If ``None`` (default), the
       value is determined from the size of ``Zxx``.
-    input_onesided: If Tru` (default), interpret the input as a one-sided STFT
+    input_onesided: If True (default), interpret the input as a one-sided STFT
       (positive frequencies only). If False, interpret the input as a two-sided STFT.
     boundary: If True (default), it is assumed that the input signal was extended at
       its boundaries by ``stft``. If `False`, the input signal is assumed to have been truncated at the boundaries by `stft`.
@@ -1108,7 +1108,7 @@ def istft(Zxx: Array, fs: ArrayLike = 1.0, window: str = 'hann',
     raise ValueError('Must specify differing time and frequency axes!')
 
   Zxx = jnp.asarray(Zxx, dtype=jax.dtypes.canonicalize_dtype(
-      np.result_type(Zxx, np.complex64)))
+    dtypes.to_complex_dtype(Zxx.dtype)))
 
   n_default = (2 * (Zxx.shape[freq_axis] - 1) if input_onesided
                else Zxx.shape[freq_axis])
@@ -1147,7 +1147,7 @@ def istft(Zxx: Array, fs: ArrayLike = 1.0, window: str = 'hann',
   xsubs = ifunc(Zxx, axis=-2, n=nfft)[..., :nperseg_int, :]
 
   # Get window as array
-  if window == 'hann':
+  if isinstance(window, str) and window == 'hann':
     # Implement the default case without scipy
     win = jnp.array([1.0]) if nperseg_int == 1 else jnp.sin(jnp.linspace(0, jnp.pi, nperseg_int, endpoint=False)) ** 2
     win = win.astype(xsubs.dtype)
diff --git a/jaxlib/py_device_list.cc b/jaxlib/py_device_list.cc
index c5004dc57..c80602e14 100644
--- a/jaxlib/py_device_list.cc
+++ b/jaxlib/py_device_list.cc
@@ -366,14 +366,21 @@ void PyDeviceList::PopulateMemoryKindInfo() {
     throw nb::value_error("Unrecognized DeviceList type");
   }
   MemoryKindInfo info;
-  if (std::get<0>(device_list_)->size() == 0) {
+  xla::ifrt::Device* addressable_device = nullptr;
+  const int process_index = py_client_ ? py_client_->process_index() : 0;
+  for (xla::ifrt::Device* device : std::get<0>(device_list_)->devices()) {
+    if (device->ProcessIndex() == process_index) {
+      addressable_device = device;
+      break;
+    }
+  }
+  if (addressable_device == nullptr) {
     info.default_memory_kind = nb::none();
     memory_kind_info_ = std::move(info);
     return;
   }
-  xla::ifrt::Device* device = std::get<0>(device_list_)->devices()[0];
 
-  auto default_memory = device->DefaultMemory();
+  auto default_memory = addressable_device->DefaultMemory();
   if (!default_memory.ok()) {
     // Cache the error.
     memory_kind_info_ = default_memory.status();
@@ -381,9 +388,9 @@ void PyDeviceList::PopulateMemoryKindInfo() {
   }
   info.default_memory_kind = nb::cast(*(*default_memory)->Kind().memory_kind());
   nb::tuple memory_kinds =
-      nb::steal<nb::tuple>(PyTuple_New(device->Memories().size()));
-  for (size_t i = 0; i < device->Memories().size(); ++i) {
-    auto* memory = device->Memories()[i];
+      nb::steal<nb::tuple>(PyTuple_New(addressable_device->Memories().size()));
+  for (size_t i = 0; i < addressable_device->Memories().size(); ++i) {
+    auto* memory = addressable_device->Memories()[i];
     nb::str s = nb::str(memory->Kind().memory_kind()->data(),
                         memory->Kind().memory_kind()->size());
     PyTuple_SET_ITEM(memory_kinds.ptr(), i, s.release().ptr());
@@ -395,17 +402,24 @@ void PyDeviceList::PopulateMemoryKindInfo() {
 void PyDeviceList::PopulateMemoryKindInfoForDuckTypedDevices() {
   MemoryKindInfo info;
   try {
-    if (std::get<1>(device_list_).size() == 0) {
+    nb::handle addressable_device;
+    for (nb::handle device : std::get<1>(device_list_)) {
+      if (nb::cast<int>(device.attr("process_index")) ==
+          nb::cast<int>(device.attr("client").attr("process_index")())) {
+        addressable_device = device;
+        break;
+      }
+    }
+    if (!addressable_device) {
       info.default_memory_kind = nb::none();
       // info.memory_kinds is default-initialized to an empty tuple.
       memory_kind_info_ = std::move(info);
       return;
     }
-    nb::handle device = std::get<1>(device_list_)[0];
-    auto default_memory = device.attr("default_memory")();
+    auto default_memory = addressable_device.attr("default_memory")();
     info.default_memory_kind = default_memory.attr("kind");
     info.memory_kinds = nb::tuple(
-        nb::object(device.attr("addressable_memories")()));
+        nb::object(addressable_device.attr("addressable_memories")()));
     memory_kind_info_ = std::move(info);
   } catch (nb::python_error& e) {
     // Cache the error.


Test patch:
diff --git a/tests/scipy_signal_test.py b/tests/scipy_signal_test.py
index 11923257a..7ff3c8743 100644
--- a/tests/scipy_signal_test.py
+++ b/tests/scipy_signal_test.py
@@ -388,7 +388,7 @@ class LaxBackedScipySignalTests(jtu.JaxTestCase):
     ],
     dtype=default_dtypes,
     fs=[1.0, 16000.0],
-    window=['boxcar', 'triang', 'blackman', 'hamming', 'hann'],
+    window=['boxcar', 'triang', 'blackman', 'hamming', 'hann', 'USE_ARRAY'],
     onesided=[False, True],
     boundary=[False, True],
   )
@@ -399,6 +399,11 @@ class LaxBackedScipySignalTests(jtu.JaxTestCase):
       new_freq_len = (shape[freqaxis] - 1) * 2
       shape = shape[:freqaxis] + (new_freq_len ,) + shape[freqaxis + 1:]
 
+    if window == 'USE_ARRAY':
+      # ensure dtype matches the expected dtype of `xsubs` within the implementation.
+      window = np.ones(nperseg, dtype=(
+        dtypes.to_floating_dtype(dtype) if onesided else dtypes.to_complex_dtype(dtype)))
+
     kwds = dict(fs=fs, window=window, nperseg=nperseg, noverlap=noverlap,
                 nfft=nfft, input_onesided=onesided, boundary=boundary,
                 time_axis=timeaxis, freq_axis=freqaxis)