ENH: sparse.spmatrix: allow fast import

[ricardoV94]

Is your feature request related to a problem? Please describe.

Trying to import spmatrix for dispatch registering triggers a whole load of imports.

Plotting the output of python -X importtime -c "from scipy.sparse._matrix import spmatrix" 2> import.log && python -m tuna import.log

Describe the solution you'd like.

Could we have a way to import spmatrix directly without triggering everything in between?

Describe alternatives you've considered.

No response

Additional context (e.g. screenshots, GIFs)

No response

[perimosocordiae]

The issue here is that while the _matrix.py file is quite trivial, Python will import all parent modules (scipy.sparse and scipy) automatically.

It may be possible to remove some of the cost by making some imports lazy: https://scientific-python.org/specs/spec-0001/
For example, doing this for the csgraph import would save ~18% based on your plot.

[lucascolley]

It may be possible to remove some of the cost by making some imports lazy

If I understand correctly from https://github.com/scipy/scipy/pull/15230/files#r771091534, we currently have lazy loading for the top-level submodules, but not for nested submodules, like sparse.csgraph - is that correct?

[dschult]

If I understand correctly from https://github.com/scipy/scipy/pull/15230/files#r771091534, we currently have lazy loading for the top-level submodules, but not for nested submodules, like sparse.csgraph - is that correct?

That is correct. We have lazy loading for the top level submodules. It turns out sparse.csgraph gets imported with sparse, but sparse.linalg does not. The sparse.linalg subpackage requires direct importing rather than lazy loading or relying on sparse being imported. I don't know much about other non-top-level-subpackages. But it makes sense to me that csgraph and linalg both be lazy-loaded when sparse is loaded. I haven't seen any issues appearing with the currently top-level lazy-loading. And it has been a few years since that went in.

It might be a good time to add lazy-loading for lower level subpackages. At least for sparse.linalg and sparse.csgraph.

@ricardoV94, the biggest chunk (44%) of import time is from importing numpy, which scipy does at the top level to verify that a valid numpy version is available. So we may be able to save you the 0.035s of csgraph, but not the 0.083s of numpy (unless we convince numpy to use lazy-load methods or we implement a way to make numpy load lazily within scipy.)

Why do you need to import spmatrix for dispatch registering? Can you lazy-register? It seems as shame to import all the possible packages you might want to dispatch to. Perhaps there is a way to delay that.

[ricardoV94]

I'm not worried about the numpy import, I already have to import it anyway. Making those other submodules lazy sounds perfect.

We have to singledispatch in case the user calls our conversion function on a sparse matrix, which will only happen for a very small part of users.

The current version of the code is broadly importing scipy.sparse but the point is to allow this line to happen with little extra overhead:

https://github.com/pymc-devs/pytensor/blob/450efff6c7e65475c95d9b7c4d656c102c80414e/pytensor/sparse/sharedvar.py#L16

[ricardoV94]

While we are here, may also look at what from scipy import stats triggers:

[dschult]

Can you import just as the dispatch is occurring? (lazy-dispatch?)
There has got to be a better way than importing all the possible libraries just to see what is possible.

[ricardoV94]

Can you import just as the dispatch is occurring? (lazy-dispatch?)

That may mean I override the dispatch from users instead of them overriding ours. Dispatch is already tricky in that it requires stuff to be imported explicitly, and this would also require it be called once before it's there.

It also doesn't avoid the problem of importing a module that most users won't need unless I replace singledispatch by a sequence of if statements with local imports.

Note this is not a blocker for us just a quality of life improvement. I imagine it would benefit others besides us though as these modules seem to all import each other for what I assume are specific cases? Like stats importing csgraph

[rgommers]

+1 to making import scipy.sparse.csgraph lazy (and sparse.linalg too while we're at it indeed, that's more minor but improves tab-completion). That's an independent public module, so there's no need for sparse users to pay the overhead of that import.

I imagine it would benefit others besides us though as these modules seem to all import each other for what I assume are specific cases? Like stats importing csgraph

New cross-submodule dependencies come in sometimes in PRs, and it's hard to enforce boundaries consistently. We do have a test that there are no import cycles, but that's about it. I'm having a look at enforcing these internal dependencies and preventing new dependencies coming in unnoticed by using tach.

[rgommers]

I'm having a look at enforcing these internal dependencies and preventing new dependencies coming in unnoticed by using tach.

That shows it's a bit too spaghetti-like. See https://github.com/rgommers/scipy/blob/use-tach/tach.toml, which passes with tach check --exact. Some dependencies are conditional though, they're from within a single function - there's no way to distinguish that from top-level imports.

The end result looks like this:

There is too much coupling; some modules like fft, special and ndimage that you'd expect to be more or less standalone still depend on quite a few other modules, as you can verify with:

>>> import sys
>>> import scipy.fft
>>> [s for s in sys.modules.keys() if s.startswith('scipy')]

Or, with a bit of cleaning:

>>> import sys
>>> import scipy.fft  # check what scipy.fft drags in
>>> all_mods = [s for s in sys.modules.keys() if s.startswith('scipy')]
>>> submodules = [
...     'cluster',
...     'constants',
...     'datasets',
...     'differentiate',
...     'fft',
...     'fftpack',
...     'integrate',
...     'interpolate',
...     'io',
...     'linalg',
...     'ndimage',
...     'odr',
...     'optimize',
...     'signal',
...     'sparse',
...     'spatial',
...     'special',
...     'stats'
... ]
>>> depends_on = set()
... for mod in all_mods:
...     for public_mod in submodules:
...         if mod.startswith(f"scipy.{public_mod}"):
...             depends_on.update([public_mod])
...
>>> depends_on.remove('fft')
>>> depends_on
{'special', 'linalg', 'sparse'}

[lucascolley]

(slightly off-topic but while we are here)

• the fft dependency on special looks "fine" - the Fast Hankel Transform implementation uses loggamma and poch
• the dependency on sparse is I believe due to two imports of https://docs.scipy.org/doc/scipy/reference/generated/scipy.sparse.issparse.html in _lib, for input validation in _array_api.array_namespace and _utils._asarray_validated. These imports are local to the functions.
• the linalg dependency is again local to the function and via _lib, but looks like we could just use np.linalg instead?
  

  scipy/scipy/_lib/_finite_differences.py

  Lines 57 to 65 in abcbaec

  	from scipy import linalg
  	ho = Np >> 1
  	x = arange(-ho, ho + 1.0)
  	x = x[:, newaxis]
  	X = x**0.0
  	for k in range(1, Np):
  	X = hstack([X, x**k])
  	w = prod(arange(1, ndiv + 1), axis=0) * linalg.inv(X)[ndiv]

I suppose this shows that the situation may not be as bad as the spaghetti makes it seem (if we can be smart about things in _lib).

[lucascolley]

Another example - the dependency of special on linalg looks to be this single line:

scipy/scipy/special/_orthogonal.py

Line 179 in abcbaec

x = linalg.eigvals_banded(c, overwrite_a_band=True)