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

Code before:
                              for shape in shapes]
                samples = [xp.reshape(sample, new_shape)
                           for sample, new_shape in zip(samples, new_shapes)]
            axis = -1  # work over the last axis

            NaN = _get_nan(*samples) if samples else xp.nan

            # if axis is not needed, just handle nan_policy and return
            ndims = np.array([sample.ndim for sample in samples])  # NumPy OK for ndims
            if np.all(ndims <= 1):
                # Addresses nan_policy == "raise"
                if nan_policy != 'propagate' or override['nan_propagation']:
                    contains_nan = [_contains_nan(sample, nan_policy)
                                    for sample in samples]
                else:
                    # Behave as though there are no NaNs (even if there are)
                    contains_nan = [False] * len(samples)

                any_contains_nan = not is_jax(xp) and any(contains_nan)
                # Addresses nan_policy == "propagate"
                if any_contains_nan and (nan_policy == 'propagate'
                                         and override['nan_propagation']):
                    res = xp.full(n_out, xp.nan, dtype=NaN.dtype)
                    res = _add_reduced_axes(res, reduced_axes, keepdims)
                    return tuple_to_result(*res)

                # Addresses nan_policy == "omit"
                too_small_msg = too_small_1d_not_omit
                if any_contains_nan and nan_policy == 'omit':
                    # consider passing in contains_nan
                    samples = _remove_nans(samples, paired)
                    too_small_msg = too_small_1d_omit

                if sentinel:
                    samples = _remove_sentinel(samples, paired, sentinel)

                if is_too_small(samples, kwds):
                    warnings.warn(too_small_msg, SmallSampleWarning, stacklevel=2)
                    res = xp.full(n_out, xp.nan, dtype=NaN.dtype)
                    res = _add_reduced_axes(res, reduced_axes, keepdims)
                    return tuple_to_result(*res)

                res = hypotest_fun_out(*samples, **kwds)
                res = result_to_tuple(res, n_out)
                res = _add_reduced_axes(res, reduced_axes, keepdims)
                return tuple_to_result(*res)

            # check for empty input
            empty_output = _check_empty_inputs(samples, axis, xp=xp)
            # only return empty output if zero sized input is too small.
            if (

Review comment (suggestion):
```suggestion
                # To give JAX most benefits of the `_axis_nan_policy` decorator without
                # value-dependent branching, the decorator will always treat JAX arrays
                # as if there are no NaNs. Typically, this means that NaNs will
                # multidimensional naturally. However, there may be cases in which this
                # doesn't produce the desired result, so for now, JAX users should
                # treat all functions as though there is no explicit`nan_policy`.
                # Future work tracked in gh-14651.
                any_contains_nan = not is_jax(xp) and any(contains_nan)
```

Code after:
                              for shape in shapes]
                samples = [xp.reshape(sample, new_shape)
                           for sample, new_shape in zip(samples, new_shapes)]
            axis = -1  # work over the last axis

            NaN = _get_nan(*samples) if samples else xp.nan

            # if axis is not needed, just handle nan_policy and return
            ndims = np.array([sample.ndim for sample in samples])  # NumPy OK for ndims
            if np.all(ndims <= 1):
                # Addresses nan_policy == "raise"
                if nan_policy != 'propagate' or override['nan_propagation']:
                    contains_nan = [_contains_nan(sample, nan_policy)
                                    for sample in samples]
                else:
                    # Behave as though there are no NaNs (even if there are)
                    contains_nan = [False] * len(samples)

                # To give JAX most benefits of the `_axis_nan_policy` decorator without
                # value-dependent branching, the decorator will always treat JAX arrays
                # as if there are no NaNs. Typically, this means that NaNs will
                # multidimensional naturally. However, there may be cases in which this
                # doesn't produce the desired result, so for now, JAX users should
                # treat all functions as though there is no explicit`nan_policy`.
                # Future work tracked in gh-14651.
                any_contains_nan = not is_jax(xp) and any(contains_nan)
                # Addresses nan_policy == "propagate"
                if any_contains_nan and (nan_policy == 'propagate'
                                         and override['nan_propagation']):
                    res = xp.full(n_out, xp.nan, dtype=NaN.dtype)
                    res = _add_reduced_axes(res, reduced_axes, keepdims)
                    return tuple_to_result(*res)

                # Addresses nan_policy == "omit"
                too_small_msg = too_small_1d_not_omit
                if any_contains_nan and nan_policy == 'omit':
                    # consider passing in contains_nan
                    samples = _remove_nans(samples, paired)
                    too_small_msg = too_small_1d_omit

                if sentinel:
                    samples = _remove_sentinel(samples, paired, sentinel)

                if is_too_small(samples, kwds):
                    warnings.warn(too_small_msg, SmallSampleWarning, stacklevel=2)
                    res = xp.full(n_out, xp.nan, dtype=NaN.dtype)
                    res = _add_reduced_axes(res, reduced_axes, keepdims)
                    return tuple_to_result(*res)

                res = hypotest_fun_out(*samples, **kwds)
                res = result_to_tuple(res, n_out)