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

Code before:

            _f = globals()[self.name]  # Allow nested wrapping
            def f(*args, _f=_f, xp=xp, **kwargs):
                # Hide dtype kwarg from map_blocks
                return xp.map_blocks(functools.partial(_f, **kwargs), *args)

            return f

        # As a final resort, use the NumPy/SciPy implementation
        _f = self.func

        if is_jax(xp) and self.is_ufunc:
            # if this is a ufunc, we can use the resolve_dtypes method to figure
            # out what the output dtype should be and use lazy_apply to make this
            # work with the JAX JIT. TODO: explore making this work for other
            # lazy backends.
            def f(*args, _f=_f, xp=xp, **kwargs):
                # JAX uses NumPy dtypes, so it's easy to get the dtypes for
                # use in resolve_dtypes. Needs the is_jax_array arg to properly
                # handle Python scalars. The (None, ) appended on the end is for the
                # dtype of out, which resolve_dtypes needs. Since JAX arrays are
                # immutable, there can be no out, so this will always be None.
                dtypes = (
                    tuple(arg.dtype if is_jax_array(arg) else type(arg) for arg in args)
                    + (None, )
                )
                out_dtype = getattr(xp, str(_f.resolve_dtypes(dtypes)[-1]))
                return xpx.lazy_apply(
                    _f, *args, xp=xp, as_numpy=True, dtype=out_dtype, **kwargs
                )
        else:
            def f(*args, _f=_f, xp=xp, **kwargs):
                args = [np.asarray(arg) for arg in args]
                out = _f(*args, **kwargs)
                return xp.asarray(out)

        return f


def _get_native_func(xp, namespace, f_name, *, alt_names_map=None):
    if alt_names_map is None:
        alt_names_map = {}
    f_name = alt_names_map.get(get_native_namespace_name(xp), f_name)
    f = getattr(namespace, f_name, None)
    if f is None and hasattr(xp, 'special'):
        # Currently dead branch, in anticipation of 'special' Array API extension
        # https://github.com/data-apis/array-api/issues/725
        f = getattr(xp.special, f_name, None)
    return f



Review comment (suggestion):
If you agree that this is more readable:
```suggestion
                dtypes = (arg.dtype if is_jax_array(arg) else type(arg) for arg in args)
                dtypes = (*dtypes, None)
```
I would also move the comment about appending `None` between the two lines. 

Code after:
            # 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!

            _f = globals()[self.name]  # Allow nested wrapping
            def f(*args, _f=_f, xp=xp, **kwargs):
                # Hide dtype kwarg from map_blocks
                return xp.map_blocks(functools.partial(_f, **kwargs), *args)

            return f

        # As a final resort, use the NumPy/SciPy implementation
        _f = self.func

        if is_jax(xp) and self.is_ufunc:
            # if this is a ufunc, we can use the resolve_dtypes method to figure
            # out what the output dtype should be and use lazy_apply to make this
            # work with the JAX JIT. TODO: explore making this work for other
            # lazy backends.
            def f(*args, _f=_f, xp=xp, **kwargs):
                # JAX uses NumPy dtypes, so it's easy to get the dtypes for
                # use in resolve_dtypes, which only works with NumPy dtypes.
                # This needs to call `is_jax_array(arg)` in the
                # ternary below to properly handle Python scalars.
                dtypes = (arg.dtype if is_jax_array(arg) else type(arg) for arg in args)
                # result_dtypes needs an arg for the dtype of the optional out param.
                # Use `None` since `out` is incompatible with JAX's immutability.
                dtypes = (*dtypes, None)
                out_dtype = getattr(xp, str(_f.resolve_dtypes(dtypes)[-1]))
                return xpx.lazy_apply(
                    _f, *args, xp=xp, as_numpy=True, dtype=out_dtype, **kwargs
                )
        else:
            def f(*args, _f=_f, xp=xp, **kwargs):
                args = [np.asarray(arg) for arg in args]
                out = _f(*args, **kwargs)
                return xp.asarray(out)

        return f


def _get_native_func(xp, namespace, f_name, *, alt_names_map=None):
    if alt_names_map is None:
        alt_names_map = {}
    f_name = alt_names_map.get(get_native_namespace_name(xp), f_name)
    f = getattr(namespace, f_name, None)
    if f is None and hasattr(xp, 'special'):
        # Currently dead branch, in anticipation of 'special' Array API extension
        # https://github.com/data-apis/array-api/issues/725
        f = getattr(xp.special, f_name, None)
    return f