Repository: scipy/scipy
File: scipy/special/_support_alternative_backends.py (Python)

Code before:
            func = self.func

        capabilities = self.xp_capabilities or xp_capabilities()
        # In order to retain a naked ufunc when SCIPY_ARRAY_API is
        # disabled, xp_capabilities must apply its changes in place.
        cap_func = capabilities(func)
        assert cap_func is func
        return func

    @functools.lru_cache(1000)
    def _wrapper_for(self, xp):
        if is_numpy(xp):
            return self.func

        # If a native implementation is available, use that
        spx = scipy_namespace_for(xp)
        f = _get_native_func(
            xp, spx, self.name, alt_names_map=self.alt_names_map, in_xp=self.in_xp
        )
        if f is not None:
            return f

        # If generic Array API implementation is available, use that
        if self.generic_impl is not None:
            f = self.generic_impl(xp, spx)
            if f is not None:
                return f

        if is_marray(xp):
            # Unwrap the array, apply the function on the wrapped namespace,
            # and then re-wrap it.
            # IMPORTANT: this only works because all functions in this module
            # are elementwise. Otherwise, we would not be able to define a
            # general rule for mask propagation.

            _f = globals()[self.name]  # Allow nested wrapping
            def f(*args, _f=_f, xp=xp, **kwargs):
                data_args = [getattr(arg, 'data', arg) for arg in args]
                out = _f(*data_args, **kwargs)
                mask = functools.reduce(operator.or_,
                                        (getattr(arg, 'mask', False) for arg in args))
                return xp.asarray(out, mask=mask)

            return f

        if is_dask(xp):
            # Apply the function to each block of the Dask array.
            # IMPORTANT: map_blocks works only because all functions in this module
            # are elementwise. It would be a grave mistake to apply this to gufuncs
            # or any other function with reductions, as they would change their
            # output depending on chunking!

Review comment (refactor):
Can we simplify the signature here so that we only need to fetch `spx` if `not self.in_xp`?

```suggestion
        namespace = xp if self.in_xp else scipy_namespace_for(xp)
        f = _get_native_func(
            xp, namespace, self.name, alt_names_map=self.alt_names_map,
        )
```


Code after:
            # Allow pickling the function. Normally this is done by @wraps,
            # but in this case it doesn't work because self.func is a ufunc.
            wrapped.__module__ = "scipy.special"
            wrapped.__qualname__ = self.name
            func = wrapped
        else:
            func = self.func

        capabilities = self.xp_capabilities or xp_capabilities()
        # In order to retain a naked ufunc when SCIPY_ARRAY_API is
        # disabled, xp_capabilities must apply its changes in place.
        cap_func = capabilities(func)
        assert cap_func is func
        return func

    @functools.lru_cache(1000)
    def _wrapper_for(self, xp):
        if is_numpy(xp):
            return self.func

        # If a native implementation is available, use that
        in_xp = get_native_namespace_name(xp) in self.backends_with_func_in_xp
        namespace = xp if in_xp else _special_namespace_for(xp)
        f = _get_native_func(
            xp, namespace, self.name, alt_names_map=self.alt_names_map
        )
        if f is not None:
            return f

        if in_xp:
            # when namespace is passed to self.generic_impl below, we want to
            # make sure that it is the special namespace for xp and not xp
            # itself, so raise if xp was incorrectly placed in
            # `backends_with_func_in_xp`.
            raise RuntimeError(
                f"func {self.func} is not available as {xp.__name__}.{self.func}"
                f" but {xp.__name__} was passed in ``in_xp``."
            )

        # If generic Array API implementation is available, use that
        if self.generic_impl is not None:
            f = self.generic_impl(xp, namespace)
            if f is not None:
                return f

        if is_marray(xp):
            # Unwrap the array, apply the function on the wrapped namespace,
            # and then re-wrap it.
            # IMPORTANT: this only works because all functions in this module
            # are elementwise. Otherwise, we would not be able to define a
            # general rule for mask propagation.