Repository: scipy/scipy
File: scipy/stats/_axis_nan_policy.py (Python)

Code before:
            sample = sample.data.copy() if np.any(mask) else sample.data
            sample = np.asarray(sample)  # `sample.data` could be a memoryview?
            sample[mask] = sentinel
        out_samples.append(sample)

    return out_samples, sentinel


def _check_empty_inputs(samples, axis):
    """
    Check for empty sample; return appropriate output for a vectorized hypotest
    """
    # if none of the samples are empty, we need to perform the test
    if not any(sample.size == 0 for sample in samples):
        return None
    # otherwise, the statistic and p-value will be either empty arrays or
    # arrays with NaNs. Produce the appropriate array and return it.
    output_shape = _broadcast_array_shapes_remove_axis(samples, axis)
    output = np.ones(output_shape) * _get_nan(*samples)
    return output


def _add_reduced_axes(res, reduced_axes, keepdims, xp=np):
    """
    Add reduced axes back to all the arrays in the result object
    if keepdims = True.
    """
    return ([xpx.expand_dims(output, axis=reduced_axes)
             if not isinstance(output, int) else output for output in res]
            if keepdims else res)


# Standard docstring / signature entries for `axis`, `nan_policy`, `keepdims`
_name = 'axis'
_desc = (
    """If an int, the axis of the input along which to compute the statistic.
The statistic of each axis-slice (e.g. row) of the input will appear in a
corresponding element of the output.
If ``None``, the input will be raveled before computing the statistic."""
    .split('\n'))


def _get_axis_params(default_axis=0, _name=_name, _desc=_desc):  # bind NOW
    _type = f"int or None, default: {default_axis}"
    _axis_parameter_doc = Parameter(_name, _type, _desc)
    _axis_parameter = inspect.Parameter(_name,
                                        inspect.Parameter.KEYWORD_ONLY,
                                        default=default_axis)
    return _axis_parameter_doc, _axis_parameter



Review comment (suggestion):
The changes in this file are pretty straightforward xp-translations that allow non-numpy arrays to go through the first part of the decorator, which makes a few things much easier for wrapped functions.
  - The dimensions specified by `axis=None` and `axis` tuples are raveled (so functions only have to consider scalar, in `axis`)
  - The remaining `axis` is moved to `-1`, which tends to make indexing easier.
  - `keepdims` is supported automatically by adding back any axes that get reduced away.

It also ensures that when the input array is too small (e.g. empty), the right thing happens (e.g. emit `SmallSampleWarning`, return NaN).

Code after:
            mask = np.broadcast_to(mask, sample.shape)
            sample = sample.data.copy() if np.any(mask) else sample.data
            sample = np.asarray(sample)  # `sample.data` could be a memoryview?
            sample[mask] = sentinel
        out_samples.append(sample)

    return out_samples, sentinel


def _check_empty_inputs(samples, axis, xp=None):
    """
    Check for empty sample; return appropriate output for a vectorized hypotest
    """
    xp = array_namespace(*samples) if xp is None else xp
    # if none of the samples are empty, we need to perform the test
    if not any(xp_size(sample) == 0 for sample in samples):
        return None
    # otherwise, the statistic and p-value will be either empty arrays or
    # arrays with NaNs. Produce the appropriate array and return it.
    output_shape = _broadcast_array_shapes_remove_axis(samples, axis)
    NaN = _get_nan(*samples)
    output = xp.full(output_shape, xp.nan, dtype=NaN.dtype)
    return output


def _add_reduced_axes(res, reduced_axes, keepdims, xp=np):
    """
    Add reduced axes back to all the arrays in the result object
    if keepdims = True.
    """
    return ([xpx.expand_dims(output, axis=reduced_axes)
             if not isinstance(output, int) else output for output in res]
            if keepdims else res)


# Standard docstring / signature entries for `axis`, `nan_policy`, `keepdims`
_name = 'axis'
_desc = (
    """If an int, the axis of the input along which to compute the statistic.
The statistic of each axis-slice (e.g. row) of the input will appear in a
corresponding element of the output.
If ``None``, the input will be raveled before computing the statistic."""
    .split('\n'))


def _get_axis_params(default_axis=0, _name=_name, _desc=_desc):  # bind NOW
    _type = f"int or None, default: {default_axis}"
    _axis_parameter_doc = Parameter(_name, _type, _desc)
    _axis_parameter = inspect.Parameter(_name,
                                        inspect.Parameter.KEYWORD_ONLY,
                                        default=default_axis)