def interpolate(
    input: mx.array,
    size: Optional[Union[int, Tuple[int, ...], List[int]]] = None,
    scale_factor: Optional[Union[float, List[float], Tuple[float, ...]]] = None,
    mode: str = "nearest",
    align_corners: Optional[bool] = None,
) -> mx.array:
    """Interpolate array with correct shape handling.

    Args:
        input (mx.array): Input tensor [N, C, ...] where ... represents spatial dimensions
        size (int or tuple): Output size
        scale_factor (float or tuple): Multiplier for spatial size
        mode (str): 'nearest' or 'linear'
        align_corners (bool): If True, align corners of input and output tensors
    """
    ndim = input.ndim
    if ndim < 3:
        raise ValueError(f"Expected at least 3D input (N, C, D1), got {ndim}D")

    spatial_dims = ndim - 2

    # Handle size and scale_factor
    if size is not None and scale_factor is not None:
        raise ValueError("Only one of size or scale_factor should be defined")
    elif size is None and scale_factor is None:
        raise ValueError("One of size or scale_factor must be defined")

    # Convert single values to tuples
    if size is not None and not isinstance(size, (list, tuple)):
        size = [size] * spatial_dims
    if scale_factor is not None and not isinstance(scale_factor, (list, tuple)):
        scale_factor = [scale_factor] * spatial_dims

    # Calculate output size from scale factor if needed
    if size is None:
        size = []
        for i in range(spatial_dims):
            # Use ceiling instead of floor to match PyTorch behavior
            curr_size = max(1, int(mx.ceil(input.shape[i + 2] * scale_factor[i])))
            size.append(curr_size)

    # Handle 1D case (N, C, W)
    if spatial_dims == 1:
        return interpolate1d(input, size[0], mode, align_corners)
    else:
        raise ValueError(
            f"Only 1D interpolation currently supported, got {spatial_dims}D"
        )
Interpolate array with correct shape handling.

    Args:
        input (mx.array): Input tensor [N, C, ...] where ... represents spatial dimensions
        size (int or tuple): Output size
        scale_factor (float or tuple): Multiplier for spatial size
        mode (str): 'nearest' or 'linear'
        align_corners (bool): If True, align corners of input and output tensors
    
mlx_audio/tts/models/interpolate.py