def broadcastable_mean_std(
    x: np.ndarray | "torch.Tensor" | "tf.Tensor", mean: np.ndarray, std: np.ndarray
) -> tuple[np.ndarray, np.ndarray]:
    """
    Ensure that the mean and standard deviation are broadcastable with respect to input `x`.

    :param x: Input samples to standardise.
    :param mean: Mean.
    :param std: Standard Deviation.
    """
    if mean.shape != std.shape:  # pragma: no cover
        raise ValueError("The shape of mean and the standard deviation must be identical.")

    # catch non-broadcastable input, when mean and std are vectors
    if mean.ndim == 1 and mean.shape[0] > 1 and mean.shape[0] != x.shape[-1]:
        # allow input shapes NC* (batch) and C* (non-batch)
        channel_idx = 1 if x.shape[1] == mean.shape[0] else 0
        broadcastable_shape = [1] * x.ndim
        broadcastable_shape[channel_idx] = mean.shape[0]

        # expand mean and std to new shape
        mean = mean.reshape(broadcastable_shape)
        std = std.reshape(broadcastable_shape)
    return mean, std

    Ensure that the mean and standard deviation are broadcastable with respect to input `x`.

    :param x: Input samples to standardise.
    :param mean: Mean.
    :param std: Standard Deviation.
    
art/preprocessing/standardisation_mean_std/utils.py