def median_filter(x: np.ndarray, filter_width: int):
    """Apply a median filter of width `filter_width` along the last dimension of `x`"""
    pad_width = filter_width // 2
    if x.shape[-1] <= pad_width:
        # F.pad requires the padding width to be smaller than the input dimension
        return x

    if (ndim := x.ndim) <= 2:
        # `F.pad` does not support 1D or 2D inputs for reflect padding but supports 3D and 4D
        x = x[None, None, :]

    assert (
        filter_width > 0 and filter_width % 2 == 1
    ), "`filter_width` should be an odd number"

    x = np.pad(x, ((0, 0), (0, 0), (pad_width, pad_width)), mode="reflect")

    # todo: more efficient version in mlx
    result = signal.medfilt(x.astype(np.float32), kernel_size=(1, 1, filter_width))[
        ..., pad_width:-pad_width
    ]

    if ndim <= 2:
        result = result[0, 0]

    return result
Apply a median filter of width `filter_width` along the last dimension of `x`
mlx_audio/stt/models/whisper/timing.py