[Suggestion] : Remove strict requirement on backends.

[ASKabalan]

Currently, we have to install pytorch to use the JAX backend of s2fft.

I think it would be a nice to be able to conditionally activate backends depending on the availability of the package

For example I user can use s2FFT using only numpy if he does not have JAX or pytorch
Otherwise he gets a runtime error instead of an import time error

Same for pyssht.
Don't know the best practice to do this, but a try catch around the import + global boolean should do the trick

[jasonmcewen]

Thanks for the suggestion @ASKabalan! At present we only have PyTorch support for the precompute approach but we do plan to add PyTorch support for the on-the-fly transforms as well. So it would be good to have a nice setup like you're suggesting.

@matt-graham do you have any thoughts on this or know of best practices?

[matt-graham]

Hi @jasonmcewen. This is related to what I suggested in #224 - we can easily make all of PyTorch, pyssht and healpy optional dependencies and guard the imports from these packages within try: ... except ImportError: ... logic, either at a module level or within the relevant functions themselves, and raise an informative error if a user tries to use a function for which the relevant optional dependencies are not available.

I would say dropping hard requirements on all external packages which are not 'core' to the package would be good practice for a library like s2fft as it avoids forcing users who only want to rely on a subset of the functionality to install unnecessary dependencies, and as I mentioned in #224 extends the range of systems that users can install the package on. The overhead for those users who wish to use the optional features is minimal as we can add relevant optional dependencies groups so that for example they could just do pip install s2fft[all] to install all optional dependencies, or pip install s2fft[torch] to just install extra dependencies required for PyTorch support and so on.

[matt-graham]

Just noticed suggestion to also apply this to JAX - while we could do this for JAX too, as a lot of the modules (perhaps the majority?) make multiple imports from jax package namespace, this would potentially get a bit unwieldy to guard all these imports. What might help with reducing our explicit imports from JAX and PyTorch APIs (and duplication of functions) is to use their support for the array API, by doing something like

import array_api_compat

...

def spectral_periodic_extension(fm, L: int):
    xp = array_api_compat.array_namespace(fm)
    nphi = fm.shape[0]
    return xnp.concatenate(
        (
            fm[-xnp.arange(L - nphi // 2, 0, -1) % nphi],
            fm,
            fm[xnp.arange(L - (nphi + 1) // 2) % nphi],
        )
    )

This would make the function compatible with all of NumPy, JAX and PyTorch arrays without requiring an explicit import from any of them and also avoid having multiple _jax, _torch variants of functions.

[jasonmcewen]

Hi @jasonmcewen. This is related to what I suggested in #224 - we can easily make all of PyTorch, pyssht and healpy optional dependencies and guard the imports from these packages within try: ... except ImportError: ... logic, either at a module level or within the relevant functions themselves, and raise an informative error if a user tries to use a function for which the relevant optional dependencies are not available.

I would say dropping hard requirements on all external packages which are not 'core' to the package would be good practice for a library like s2fft as it avoids forcing users who only want to rely on a subset of the functionality to install unnecessary dependencies, and as I mentioned in #224 extends the range of systems that users can install the package on. The overhead for those users who wish to use the optional features is minimal as we can add relevant optional dependencies groups so that for example they could just do pip install s2fft[all] to install all optional dependencies, or pip install s2fft[torch] to just install extra dependencies required for PyTorch support and so on.

Ok, fair enough @matt-graham ! Does this mean users would need to install from source rather than PyPi tho? Perhaps that is not a big issue anyway.

[jasonmcewen]

Just noticed suggestion to also apply this to JAX - while we could do this for JAX too, as a lot of the modules (perhaps the majority?) make multiple imports from jax package namespace, this would potentially get a bit unwieldy to guard all these imports. What might help with reducing our explicit imports from JAX and PyTorch APIs (and duplication of functions) is to use their support for the array API, by doing something like

import array_api_compat

...

def spectral_periodic_extension(fm, L: int):
    xp = array_api_compat.array_namespace(fm)
    nphi = fm.shape[0]
    return xnp.concatenate(
        (
            fm[-xnp.arange(L - nphi // 2, 0, -1) % nphi],
            fm,
            fm[xnp.arange(L - (nphi + 1) // 2) % nphi],
        )
    )

This would make the function compatible with all of NumPy, JAX and PyTorch arrays without requiring an explicit import from any of them and also avoid having multiple _jax, _torch variants of functions.

Maybe. This seems like quite a big revision tho. Let's discuss further...

[matt-graham]

Ok, fair enough @matt-graham ! Does this mean users would need to install from source rather than PyPi tho? Perhaps that is not a big issue anyway.

No, the extras / optional dependencies syntax works fine for both packages installed from a local directory and a packaging index like PyPI. It's the same syntax that JAX uses to install from PyPI with CUDA support:

pip install jax[cuda12]

[jasonmcewen]

Ok, sounds like it makes a lot of sense then.

[opeltre]

Hello!

I'm curious about your current intentions on this topic, as installing s2fft indeed pulls torch which itself brings triton==3.2.0, on top of all the nvidia binaries shared or not between backends... The trend seems to be that numerical backends will keep shipping larger and larger binaries, not the opposite, as required by the JIT machineries etc.

If you're willing to make the torch dependency optional, it should be possible with very few code changes as long as torch calls are nested within function definitions. After searching through the codebase I think you could:

• have a conditional import torch statement on a single file, catching ModuleNotFoundError,
• ensure that you otherwise return an object compatible with all type annotations i.e. torch.tensor (sic) and torch.Tensor

That way you shouldn't have to refactor code anywhere else than on top level torch imports.

I'd be happy to open a PR if that's something you still look forward to. I've dealt with similar backend-agnostic issues in a couple other projects. Right now I was planning on depending upon s2fft for a small meteorology toolbox, but the size of dependencies is somewhat deterring.

Here's a recipe proposal:

1. Somewhere, make a reference to torch available otherwise all type annotations will break:

# _torch_mock.py 

class TorchMock:
     """Mock the torch module to support existing type annotations.

     Torch calls within function bodies are not compiled nor evaluated, 
     so Python will not complain with existing code. 

     Raise `ModuleNotFoundError` instead of `AttributeError` in case 
     Torch primitive is called.
     """
     tensor = TypeAlias
     Tensor = TypeAlias

     def __getattr__(self, attr): 
         raise ModuleNotFoundError("torch installation not found.")

torch = TorchMock()

2. Put the try-catch import torch to a single place

# _backends.py

try: 
    import torch

except ModuleNotFoundError:
    logging.info("torch installation not found.")
    from _torch_mock import torch

3. In code using torch, just replace import torch statements with from _backends import torch.

N.B. Regarding cpu/cuda extras, I know the situation was really terrible until a few months ago to handle e.g. different jaxlib platforms within dependencies, so it must have been hardly possible back in 2024. The situation has gotten much smoother though. So out of curiosity, are you able to install jax[cpu] within the CI, without pulling all the nvidia binaries?

[matt-graham]

@opeltre thanks for your comments, detailed suggestion and offer to contribute! I definitely agree that making torch an optional dependency would be sensible given, as you mention, it pulls in such large binary dependencies.

Your suggested approach of mocking the minimal amount we need from the torch interface and guarding imports in a try ... except ImportError block does seem like a reasonable approach to take. We've got an open PR #277 which as part of a large refactor of how we're providing torch functions (changing to wrapping JAX implementations rather than reimplementing in PyTorch) will also make torch an optional dependency in a similar manner, with this being a bit simpler in the refactored approach as we only import torch in a single module providing the utility functions for wrapping the JAX interface. Hopefully that should resolve this issue.