Skip to content

xwr.rsp

Radar Signal Processing for batched 4D spectrum.

Image Axis Order

Elevation and azimuth axes are in "image order": increasing index is down and to the right, respectively.

Byte order: when do I sample_swap?

If you are using the xwr stack, use the default sample_swap=False, which corresponds to MSB_LSB_IQ, the only option supported by the source-available TI firmware. If processing data collected using other systems (in particular, mmWave studio, which has its own closed-source firmware which supports MSB_LSB_QI), you may need to set sample_swap=True if this option was enabled.

To use the RSP:

  1. Pick your backend. Currently, we support numpy, jax, and pytorch.

    Note

    Each RSP backend is not imported by default; you must explicitly import the backend you want to use, and make sure its dependencies are installed (pytorch, jax, etc). 

    from xwr.rsp import numpy as rsp
# or
from xwr.rsp import jax as rsp

  2. Select the appropriate radar model.

  3. Import the RSP class matching your backend and radar: 

    from xwr.rsp import RSP
from xwr.rsp.torch import AWR1843AOP

rsp: RSP = AWR1843AOP()

    Tip

    Use xwr.rsp.RSP as the type for a generic RSP, and RSP[np.ndarray], RSP[jax.Array], RSP[torch.Tensor], etc for a RSP with a specific backend.

xwr.rsp.RSP

Bases: ABC, Generic[TArray]

Abstract, backend-agnostic Radar Signal Processing base class.

Info

This class documents the public interface for all radar signal processing (RSP) classes, except where otherwise noted.

Type Parameters
  • TArray: Generic backend, e.g., np.ndarray, jax jax.Array, or torch Tensor.

Parameters:

Name Type Description Default
window bool | Mapping[Literal['range', 'doppler', 'azimuth', 'elevation'], bool]

whether to apply a hanning window. If bool, the same option is applied to all axes. If dict, specify per axis with keys "range", "doppler", "azimuth", and "elevation".

 False
size Mapping[Literal['range', 'doppler', 'azimuth', 'elevation'], int]

target size for each axis after zero-padding, specified by axis. If an axis is not spacified, it is not padded.

 {}
sample_swap bool

if True, swap the I and Q components when un-interleaving IIQQ data.

 False
Source code in src/xwr/rsp/generic.py 
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
class RSP(ABC, Generic[TArray]):
    """Abstract, backend-agnostic Radar Signal Processing base class.

    !!! info

        This class documents the public interface for all radar signal
        processing (RSP) classes, except where otherwise noted.

    Type Parameters:
        - `TArray`: Generic backend, e.g., `np.ndarray`, jax `jax.Array`, or
            torch `Tensor`.

    Args:
        window: whether to apply a hanning window. If `bool`, the same option
            is applied to all axes. If `dict`, specify per axis with keys
            "range", "doppler", "azimuth", and "elevation".
        size: target size for each axis after zero-padding, specified by axis.
            If an axis is not spacified, it is not padded.
        sample_swap: if `True`, swap the I and Q components when
            un-interleaving IIQQ data.
    """

    SAMPLE_TYPE: Literal["IQ", "I"] = "IQ"

    def __init__(
        self, window: bool | Mapping[
            Literal["range", "doppler", "azimuth", "elevation"], bool] = False,
        size: Mapping[
            Literal["range", "doppler", "azimuth", "elevation"], int] = {},
        sample_swap: bool = False,
    ) -> None:
        self.window: dict[
            Literal["range", "doppler", "azimuth", "elevation"], bool]
        self._default_window: bool | dict[
            Literal["range", "doppler", "azimuth", "elevation"], bool]

        if isinstance(window, bool):
            self.window = {}
            self._default_window = window
        else:
            self.window = dict(window)
            self._default_window = False

        self.size = size
        self.sample_swap = sample_swap

    @abstractmethod
    def fft(
        self, array: Complex64[TArray, "..."] | Float32[TArray, "..."],
        axes: tuple[int, ...],
        size: tuple[int, ...] | None = None,
        shift: tuple[int, ...] | None = None
    ) -> Complex64[TArray, "..."]:
        """Compute FFT on the specified axes of the array.

        Args:
            array: Input array.
            size: Target size for each axis after FFT (or `None` to use the
                input size).
            axes: Axes along which to compute the FFT.
            shift: Axes to shift after FFT, if any.

        Returns:
            FFT of the input array along the specified axes. If the input
                array is real-valued, the output is the non-negative frequency
                terms of the FFT along the specified axes (with length
                `n // 2 + 1`).
        """
        ...

    @staticmethod
    @abstractmethod
    def hann(
        x: Complex64[TArray, "..."] | Float32[TArray, "..."], axis: int
    ) -> Complex64[TArray, "..."] | Float32[TArray, "..."]:
        """Apply a Hann window to the specified axis of the time signal data.

        Args:
            x: time signal data.
            axis: Axis along which to apply the Hann window.

        Returns:
            Time signal data with the Hann window applied along the specified
                axis.
        """
        ...

    def doppler_range(
        self, x: Complex64[TArray, "#batch doppler tx rx range"]
            | Float32[TArray, "#batch doppler tx rx range"]
    ) -> Complex64[TArray, "#batch doppler2 tx rx range2"]:
        """Calculate range-doppler spectrum from time signal data.

        Args:
            x: IQ (complex64) or in-phase-only (float32) data.

        Returns:
            Computed range-doppler spectrum, with windowing if specified.
        """
        if self.window.get("range", self._default_window):
            x = self.hann(x, 4)
        if self.window.get("doppler", self._default_window):
            x = self.hann(x, 1)

        if self.SAMPLE_TYPE == "I":
            # Double range bins for in-phase-only.
            nrange = x.shape[-1] // 2
            range_bins = self.size.get("range", nrange) * 2
        else:
            range_bins = self.size.get("range", x.shape[4])

        rd = self.fft(
            x, axes=(1, 4), shift=(1,),
            size=(self.size.get("doppler", x.shape[1]), range_bins))

        return rd

    @abstractmethod
    def mimo_virtual_array(
        self, rd: Complex64[TArray, "#batch doppler tx rx range"]
    ) -> Complex64[TArray, "#batch doppler elevation azimuth range"]:
        """Set up MIMO virtual array from range-doppler spectrum.

        Args:
            rd: complex range-doppler spectrum.

        Returns:
            Computed MIMO virtual array, in elevation-azimuth order.
        """
        ...

    def elevation_azimuth(
        self, rd: Complex64[TArray, "#batch doppler tx rx range"]
    ) -> Complex64[TArray, "#batch doppler el az range"]:
        """Calculate elevation-azimuth spectrum from range-doppler spectrum.

        Args:
            rd: range-doppler spectrum.

        Returns:
            Computed elevation-azimuth spectrum, with windowing and padding if
                specified.
        """
        mimo = self.mimo_virtual_array(rd)

        if self.window.get("elevation", self._default_window):
            mimo = self.hann(mimo, 2)
        if self.window.get("azimuth", self._default_window):
            mimo = self.hann(mimo, 3)

        return self.fft(
            mimo, axes=(2, 3), shift=(2, 3),
            size=(
                self.size.get("elevation", mimo.shape[2]),
                self.size.get("azimuth", mimo.shape[3])))

    def __call__(
        self, x: Complex64[TArray, "#batch doppler tx rx _range"]
            | Float32[TArray, "#batch doppler tx rx _range"]
            | Int16[TArray, "#batch doppler tx rx _range"]
    ) -> Complex64[TArray, "#batch doppler2 el az _range"]:
        """Process time signal data to compute elevation-azimuth spectrum.

        Args:
            x: IQ data in complex or interleaved int16 IQ format, or
                in-phase-only data in float32 format.

        Returns:
            Computed doppler-elevation-azimuth-range spectrum.
        """
        if self.SAMPLE_TYPE == "IQ":
            x = iq_from_iiqq(x, sample_swap=self.sample_swap)
        else:
            x = _to_float32(x)

        dr = self.doppler_range(x)
        drae = self.elevation_azimuth(dr)
        return drae

__call__ 

__call__(
    x: Complex64[TArray, "#batch doppler tx rx _range"]
    | Float32[TArray, "#batch doppler tx rx _range"]
    | Int16[TArray, "#batch doppler tx rx _range"],
) -> Complex64[TArray, "#batch doppler2 el az _range"]

Process time signal data to compute elevation-azimuth spectrum.

Parameters:

Name Type Description Default
x Complex64[TArray, '#batch doppler tx rx _range'] | Float32[TArray, '#batch doppler tx rx _range'] | Int16[TArray, '#batch doppler tx rx _range']

IQ data in complex or interleaved int16 IQ format, or in-phase-only data in float32 format.

 required

Returns:

Type Description
Complex64[TArray, '#batch doppler2 el az _range']

Computed doppler-elevation-azimuth-range spectrum.
Source code in src/xwr/rsp/generic.py 
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
def __call__(
    self, x: Complex64[TArray, "#batch doppler tx rx _range"]
        | Float32[TArray, "#batch doppler tx rx _range"]
        | Int16[TArray, "#batch doppler tx rx _range"]
) -> Complex64[TArray, "#batch doppler2 el az _range"]:
    """Process time signal data to compute elevation-azimuth spectrum.

    Args:
        x: IQ data in complex or interleaved int16 IQ format, or
            in-phase-only data in float32 format.

    Returns:
        Computed doppler-elevation-azimuth-range spectrum.
    """
    if self.SAMPLE_TYPE == "IQ":
        x = iq_from_iiqq(x, sample_swap=self.sample_swap)
    else:
        x = _to_float32(x)

    dr = self.doppler_range(x)
    drae = self.elevation_azimuth(dr)
    return drae

doppler_range 

doppler_range(
    x: Complex64[TArray, "#batch doppler tx rx range"]
    | Float32[TArray, "#batch doppler tx rx range"],
) -> Complex64[TArray, "#batch doppler2 tx rx range2"]

Calculate range-doppler spectrum from time signal data.

Parameters:

Name Type Description Default
x Complex64[TArray, '#batch doppler tx rx range'] | Float32[TArray, '#batch doppler tx rx range']

IQ (complex64) or in-phase-only (float32) data.

 required

Returns:

Type Description
Complex64[TArray, '#batch doppler2 tx rx range2']

Computed range-doppler spectrum, with windowing if specified.
Source code in src/xwr/rsp/generic.py 
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
def doppler_range(
    self, x: Complex64[TArray, "#batch doppler tx rx range"]
        | Float32[TArray, "#batch doppler tx rx range"]
) -> Complex64[TArray, "#batch doppler2 tx rx range2"]:
    """Calculate range-doppler spectrum from time signal data.

    Args:
        x: IQ (complex64) or in-phase-only (float32) data.

    Returns:
        Computed range-doppler spectrum, with windowing if specified.
    """
    if self.window.get("range", self._default_window):
        x = self.hann(x, 4)
    if self.window.get("doppler", self._default_window):
        x = self.hann(x, 1)

    if self.SAMPLE_TYPE == "I":
        # Double range bins for in-phase-only.
        nrange = x.shape[-1] // 2
        range_bins = self.size.get("range", nrange) * 2
    else:
        range_bins = self.size.get("range", x.shape[4])

    rd = self.fft(
        x, axes=(1, 4), shift=(1,),
        size=(self.size.get("doppler", x.shape[1]), range_bins))

    return rd

elevation_azimuth 

elevation_azimuth(
    rd: Complex64[TArray, "#batch doppler tx rx range"],
) -> Complex64[TArray, "#batch doppler el az range"]

Calculate elevation-azimuth spectrum from range-doppler spectrum.

Parameters:

Name Type Description Default
rd Complex64[TArray, '#batch doppler tx rx range']

range-doppler spectrum.

 required

Returns:

Type Description
Complex64[TArray, '#batch doppler el az range']

Computed elevation-azimuth spectrum, with windowing and padding if specified.
Source code in src/xwr/rsp/generic.py 
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
def elevation_azimuth(
    self, rd: Complex64[TArray, "#batch doppler tx rx range"]
) -> Complex64[TArray, "#batch doppler el az range"]:
    """Calculate elevation-azimuth spectrum from range-doppler spectrum.

    Args:
        rd: range-doppler spectrum.

    Returns:
        Computed elevation-azimuth spectrum, with windowing and padding if
            specified.
    """
    mimo = self.mimo_virtual_array(rd)

    if self.window.get("elevation", self._default_window):
        mimo = self.hann(mimo, 2)
    if self.window.get("azimuth", self._default_window):
        mimo = self.hann(mimo, 3)

    return self.fft(
        mimo, axes=(2, 3), shift=(2, 3),
        size=(
            self.size.get("elevation", mimo.shape[2]),
            self.size.get("azimuth", mimo.shape[3])))

fft abstractmethod 

fft(
    array: Complex64[TArray, ...] | Float32[TArray, ...],
    axes: tuple[int, ...],
    size: tuple[int, ...] | None = None,
    shift: tuple[int, ...] | None = None,
) -> Complex64[TArray, ...]

Compute FFT on the specified axes of the array.

Parameters:

Name Type Description Default
array Complex64[TArray, ...] | Float32[TArray, ...]

Input array.

 required
size tuple[int, ...] | None

Target size for each axis after FFT (or None to use the input size).

 None
axes tuple[int, ...]

Axes along which to compute the FFT.

 required
shift tuple[int, ...] | None

Axes to shift after FFT, if any.

 None

Returns:

Type Description
Complex64[TArray, ...]

FFT of the input array along the specified axes. If the input array is real-valued, the output is the non-negative frequency terms of the FFT along the specified axes (with length n // 2 + 1).
Source code in src/xwr/rsp/generic.py 
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
@abstractmethod
def fft(
    self, array: Complex64[TArray, "..."] | Float32[TArray, "..."],
    axes: tuple[int, ...],
    size: tuple[int, ...] | None = None,
    shift: tuple[int, ...] | None = None
) -> Complex64[TArray, "..."]:
    """Compute FFT on the specified axes of the array.

    Args:
        array: Input array.
        size: Target size for each axis after FFT (or `None` to use the
            input size).
        axes: Axes along which to compute the FFT.
        shift: Axes to shift after FFT, if any.

    Returns:
        FFT of the input array along the specified axes. If the input
            array is real-valued, the output is the non-negative frequency
            terms of the FFT along the specified axes (with length
            `n // 2 + 1`).
    """
    ...

hann abstractmethod staticmethod 

hann(
    x: Complex64[TArray, ...] | Float32[TArray, ...], axis: int
) -> Complex64[TArray, ...] | Float32[TArray, ...]

Apply a Hann window to the specified axis of the time signal data.

Parameters:

Name Type Description Default
x Complex64[TArray, ...] | Float32[TArray, ...]

time signal data.

 required
axis int

Axis along which to apply the Hann window.

 required

Returns:

Type Description
Complex64[TArray, ...] | Float32[TArray, ...]

Time signal data with the Hann window applied along the specified axis.
Source code in src/xwr/rsp/generic.py 
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
@staticmethod
@abstractmethod
def hann(
    x: Complex64[TArray, "..."] | Float32[TArray, "..."], axis: int
) -> Complex64[TArray, "..."] | Float32[TArray, "..."]:
    """Apply a Hann window to the specified axis of the time signal data.

    Args:
        x: time signal data.
        axis: Axis along which to apply the Hann window.

    Returns:
        Time signal data with the Hann window applied along the specified
            axis.
    """
    ...

mimo_virtual_array abstractmethod 

mimo_virtual_array(
    rd: Complex64[TArray, "#batch doppler tx rx range"],
) -> Complex64[TArray, "#batch doppler elevation azimuth range"]

Set up MIMO virtual array from range-doppler spectrum.

Parameters:

Name Type Description Default
rd Complex64[TArray, '#batch doppler tx rx range']

complex range-doppler spectrum.

 required

Returns:

Type Description
Complex64[TArray, '#batch doppler elevation azimuth range']

Computed MIMO virtual array, in elevation-azimuth order.
Source code in src/xwr/rsp/generic.py 
330
331
332
333
334
335
336
337
338
339
340
341
342
@abstractmethod
def mimo_virtual_array(
    self, rd: Complex64[TArray, "#batch doppler tx rx range"]
) -> Complex64[TArray, "#batch doppler elevation azimuth range"]:
    """Set up MIMO virtual array from range-doppler spectrum.

    Args:
        rd: complex range-doppler spectrum.

    Returns:
        Computed MIMO virtual array, in elevation-azimuth order.
    """
    ...

xwr.rsp.iq_from_iiqq 

iq_from_iiqq(
    iiqq: Int16[TArray, "... n"] | Complex64[TArray, "... _n"],
    sample_swap: bool = False,
) -> Complex64[TArray, "... n2"]

Un-interleave IIQQ data.

Info

The default sample_swap = False corresponds to the MSB_LSB_IQ byte order used by xwr.

In this case, MSB_LSB_IQ means that I is the MSB and Q is the LSB. However, the data stream is little-endian, which means Q actually comes before I, leading to the actual physical layout being QQII and so on.

Type Parameters
  • TArray: This function is multi-backend, and supports numpy np.ndarray, jax jax.Array, and torch Tensor.

Parameters:

Name Type Description Default
iiqq Int16[TArray, '... n'] | Complex64[TArray, '... _n']

interleaved IIQQ data; see RadarFrame. If already complex, leave it as is.

 required
sample_swap bool

if True, swap the I and Q components so that the output is Q + j*I instead of I + j*Q.

 False

Returns:

Type Description
Complex64[TArray, '... n2']

Complex IQ data.
Source code in src/xwr/rsp/generic.py 
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
def iq_from_iiqq(
    iiqq: Int16[TArray, "... n"] | Complex64[TArray, "... _n"],
    sample_swap: bool = False,
) -> Complex64[TArray, "... n2"]:
    """Un-interleave IIQQ data.

    !!! info

        The default `sample_swap = False` corresponds to the
        [`MSB_LSB_IQ`][xwr.radar.defines.SampleSwap] byte order used by `xwr`.

        In this case, `MSB_LSB_IQ` means that I is the MSB and Q is the LSB.
        However, the data stream is little-endian, which means Q actually comes
        before I, leading to the actual physical layout being QQII and so on.

    Type Parameters:
        - `TArray`: This function is multi-backend, and supports numpy
            `np.ndarray`, jax `jax.Array`, and torch `Tensor`.

    Args:
        iiqq: interleaved IIQQ data; see [`RadarFrame`][xwr.capture.types.].
            If already complex, leave it as is.
        sample_swap: if `True`, swap the I and Q components so that the
            output is `Q + j*I` instead of `I + j*Q`.

    Returns:
        Complex IQ data.
    """
    shape = (*iiqq.shape[:-1], iiqq.shape[-1] // 2)

    backend = _check_backend(iiqq)
    if backend == "numpy":
        assert isinstance(iiqq, np.ndarray)

        if iiqq.dtype == np.complex64:
            return iiqq
        iq = np.zeros(shape, dtype=np.complex64)
        if sample_swap:
            iq[..., 0::2] = iiqq[..., 0::4] + 1j * iiqq[..., 2::4]
            iq[..., 1::2] = iiqq[..., 1::4] + 1j * iiqq[..., 3::4]
        else:
            iq[..., 0::2] = 1j * iiqq[..., 0::4] + iiqq[..., 2::4]
            iq[..., 1::2] = 1j * iiqq[..., 1::4] + iiqq[..., 3::4]
        return cast(Complex64[TArray, "... n/2"], iq)

    elif backend == "jax":
        from jax import numpy as jnp
        assert isinstance(iiqq, jnp.ndarray)

        if iiqq.dtype == jnp.complex64:
            return iiqq
        if sample_swap:
            iq = jnp.zeros(
                shape, dtype=jnp.complex64
            ).at[..., 0::2].set(iiqq[..., 0::4] + 1j * iiqq[..., 2::4]
            ).at[..., 1::2].set(iiqq[..., 1::4] + 1j * iiqq[..., 3::4])
        else:
            iq = jnp.zeros(
                shape, dtype=jnp.complex64
            ).at[..., 0::2].set(1j * iiqq[..., 0::4] + iiqq[..., 2::4]
            ).at[..., 1::2].set(1j * iiqq[..., 1::4] + iiqq[..., 3::4])
        return cast(Complex64[TArray, "... n/2"], iq)

    else: # backend == "torch"
        import torch
        assert isinstance(iiqq, torch.Tensor)

        if iiqq.dtype == torch.complex64:
            return iiqq
        iq = torch.zeros(shape, dtype=torch.complex64, device=iiqq.device)
        if sample_swap:
            iq[..., 0::2] = iiqq[..., 0::4] + 1j * iiqq[..., 2::4]
            iq[..., 1::2] = iiqq[..., 1::4] + 1j * iiqq[..., 3::4]
        else:
            iq[..., 0::2] = 1j * iiqq[..., 0::4] + iiqq[..., 2::4]
            iq[..., 1::2] = 1j * iiqq[..., 1::4] + iiqq[..., 3::4]
        return cast(Complex64[TArray, "... n/2"], iq)

xwr.rsp.iqiq_from_iiqq 

iqiq_from_iiqq(
    iiqq: Int16[TArray, "... n"], sample_swap: bool = False
) -> Int16[TArray, "... n/2 2"]

Un-interleave IIQQ data.

Type Parameters
  • TArray: This function is multi-backend, and supports numpy np.ndarray, jax jax.Array, and torch Tensor.

Parameters:

Name Type Description Default
iiqq Int16[TArray, '... n']

interleaved IIQQ data; see RadarFrame.

 required
sample_swap bool

if True, swap the I and Q components in the output.

 False

Returns:

Type Description
Int16[TArray, '... n/2 2']

IQ data in an uninterleaved format with a trailing I/Q axis.
Source code in src/xwr/rsp/generic.py 
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
def iqiq_from_iiqq(
    iiqq: Int16[TArray, "... n"],
    sample_swap: bool = False,
) -> Int16[TArray, "... n/2 2"]:
    """Un-interleave IIQQ data.

    Type Parameters:
        - `TArray`: This function is multi-backend, and supports numpy
            `np.ndarray`, jax `jax.Array`, and torch `Tensor`.

    Args:
        iiqq: interleaved IIQQ data; see [`RadarFrame`][xwr.capture.types.].
        sample_swap: if `True`, swap the I and Q components in the output.

    Returns:
        IQ data in an uninterleaved format with a trailing I/Q axis.
    """
    shape = (*iiqq.shape[:-1], iiqq.shape[-1] // 2)
    i_idx, q_idx = (1, 0) if sample_swap else (0, 1)

    backend = _check_backend(iiqq)
    if backend == "numpy":
        assert isinstance(iiqq, np.ndarray)

        iq = np.zeros((*shape, 2), dtype=np.int16)
        iq[..., 0::2, q_idx] = iiqq[..., 0::4]
        iq[..., 1::2, q_idx] = iiqq[..., 1::4]
        iq[..., 0::2, i_idx] = iiqq[..., 2::4]
        iq[..., 1::2, i_idx] = iiqq[..., 3::4]
        return cast(Int16[TArray, "... n/2 2"], iq)

    elif backend == "jax":
        from jax import numpy as jnp
        assert isinstance(iiqq, jnp.ndarray)

        iq = jnp.zeros(
            (*shape, 2), dtype=jnp.int16
        ).at[..., 0::2, q_idx].set(iiqq[..., 0::4]
        ).at[..., 1::2, q_idx].set(iiqq[..., 1::4]
        ).at[..., 0::2, i_idx].set(iiqq[..., 2::4]
        ).at[..., 1::2, i_idx].set(iiqq[..., 3::4])
        return cast(Int16[TArray, "... n/2 2"], iq)

    else:  # backend == "torch"
        import torch
        assert isinstance(iiqq, torch.Tensor)

        iq = torch.zeros((*shape, 2), dtype=torch.int16, device=iiqq.device)
        iq[..., 0::2, q_idx] = iiqq[..., 0::4]
        iq[..., 1::2, q_idx] = iiqq[..., 1::4]
        iq[..., 0::2, i_idx] = iiqq[..., 2::4]
        iq[..., 1::2, i_idx] = iiqq[..., 3::4]
        return cast(Int16[TArray, "... n/2 2"], iq)