Make BlockEncoder Sparse

[JerryChen97]

Context:
As a following up of #6889, we make BlockEncode sparse

Description of the Change:
Overloaded many qml.math methods such that BlockEncode can automatically dispatch to sparse methods.

Benefits:
BlockEncode is able to accept and output sparse matrices now

Possible Drawbacks:
Current implementation is a bit fall-back, due to lack of scipy sparse support one many methods, e.g. sum over certain axis, sqrt as a matrix square root, etc

Related GitHub Issues:
[sc-82069]

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 99.59%. Comparing base (4444bac) to head (b046fdf).

Additional details and impacted files

@@           Coverage Diff           @@
##           master    #6963   +/-   ##
=======================================
  Coverage   99.59%   99.59%           
=======================================
  Files         480      480           
  Lines       45695    45726   +31     
=======================================
+ Hits        45510    45541   +31     
  Misses        185      185           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[JerryChen97]

Suggested change

# ar.register_function("scipy", "linalg.eigh", sp.sparse.linalg.eigh)

No method eigh in scipy.sparse.eigh

[mlxd]

Isn't https://docs.scipy.org/doc/scipy/reference/generated/scipy.sparse.linalg.eigsh.html the sparse equivalent? Can we bind to this?

[JerryChen97]

@mlxd np.linalg.eigh will precisely calculate all the eigs of the given square matrix. However, scipy.sparse.linalg.eigsh can only take limited number of suspaces; they denoted that "k must be smaller than N. It is not possible to compute all eigenvectors of a matrix". This is also a quite confusing fact to me and I would like to know why the package was implemented in such a way...

[mlxd]

This is how the algorithm works at the FORTRAN layer (the Python layer binds to https://github.com/opencollab/arpack-ng). It relies on the Lanzcos alg under the hood, which nominally gives you a set of "useful" eigenvals/vecs, rather than all. They accept an argument to choose looking for the largest, smallest, etc using the which argument. It's fine to expect this is how it works in practice, though if you want to use the values, controlling what is returned may be worth modifying the function signature to accommodate the control offered at the scipy layer.

[mlxd]

Thanks @JerryChen97
Just a few suggestions from my side. The main thing I see is that there's quite a few dense conversions happening, that we should be able to avoid.

[mlxd]

Do we want to do this? If the DM is sparse won't this densify it?
Can't we just use https://docs.scipy.org/doc/scipy/reference/generated/scipy.sparse.csr_matrix.sqrt.html on the given sub-terms, or other equivalents in the scipy.sparse module?

Also, not sure if sparse mats support broadcasting/batching

[JerryChen97]

@mlxd scipy.sparse.csr_matrix.sqrt is an elementwise operation. Here the sqrt in quantum.py we would like to calculate the matrix sqrt.
The reason why this method is touched is that in BlockEncode we need to implement something like $\text{sqrtmatrix}(Id - A^\dagger A)$; and the missing of a full-spectrum eigh in scipy.sparse made things difficult, so I forced them to back to numpy as temporary solution. A better solution using exclusively only scipy.sparse methods is wanted.

[JerryChen97]

To get the matrix sqrt, as far as my knowledge, another practical method is to make use of Schur decomp. scipy.linalg has it, but again scipy.sparse.linalg do not have.

[mlxd]

Ah, my bad, I assumed that was the mat sqrt. It seems the sparse linalg routines for power also only accept integers. I think this is something then we'd need to implement for support, otherwise we are limiting ourselves in the overall mat size we can represent

[JerryChen97]

On the $\text{sqrtmatrix}(Id - A^\dagger A)$, I believe we can temporarily employ the Denman–Beavers iterations to ensure only sparse operations were applied and meanwhile get the correct results for most of random matrices. But its performance is notably poor for scipy.sparse (it requires inv in its iterations), and essentially it does not naturally preserve the sparsity of matrices at all.

[mlxd]

Isn't https://docs.scipy.org/doc/scipy/reference/generated/scipy.sparse.linalg.eigsh.html the sparse equivalent? Can we bind to this?

[mlxd]

So the https://docs.scipy.org/doc/scipy/reference/generated/scipy.sparse.csr_matrix.sum.html dispatch is auto-converting to dense? Is there an example I can test this out with, as it shouldn't need to.

Also, any reason with COO instead of the CSR rep used elsewhere?

[JerryChen97]

Indeed. Its falling back to nparray.
import scipy as sp

# Create a 2x2 sparse matrix in CSR format
m = sp.sparse.csr_matrix([[1, 0], [0, 1]])
print(type(m.sum(axis=0)))  # <class 'numpy.matrix'>

# Create a 2D matrix in COO format
m2 = sp.sparse.coo_matrix([[1, 2], [3, 4]])
print(type(m2.sum(axis=0)))  # <class 'numpy.matrix'>

[JerryChen97]

On the COO: somehow I used to have an impression that COO supports multi-dim, so picked it since this sum in our codebase is used for mutli-dim tensor to sum over arbitrary dims; I just double checked and it turns out I was wrong! I think no difference here then. That also explains why I saw sparse.sum has no support for multi axes.

[mlxd]

Ok, so I guess a reduction over one of the axes will (in most cases) reduce the sparsity, so the assumption is sum(SpMAT[m * n], axes=m) -> Dense[n], which should (in most cases) be reasonable (assuming we still keep in mind the size here can be quite large).

Though, this is problematic if somebody does it on something like the following:

mat = sp.csr_matrix(([1,1,1,1], ([0,2,1,3], [0,1,2,3])), shape=(2**30,2**30))
mat.sum(axis=0) # Will create a `size(dtype)*2**30` array

The above should be fast, at the expense of memory. On the other size, using

sum(mat)

will preserve the sparse matrix, at the expense of a lot of overheads in evaluating it. Not sure if we want to add some guard-rails here to avoid the memory blow-up, or if it is fine in practice and to just assume it to be used with reasonable sized systems.

[mlxd]

We shouldn't need to densify the full matrix here. We can compare the data fields directly with np allclose, and use the index and indptr fields to identify the positions.

There's some other solutions available in the wild. For example, you can calculate the difference of the two sparse matrices, and check if the resulting data values are above the required threshold for closeness

[JerryChen97]

Yeah if both sparse its actually simple. Here I eventually took this way because very often we need to compare nparray and sparse😂so I was thinking then lets surrender all back to np

[mlxd]

If we do this though we limit ourselves to the largest matrix sizes we can use.
I see this as a few stages:

• If both are sparse, compare the data as I mentioned earlier
• If one is dense and one sparse, check their intended sizes first; if different return false; if true, continue
• If the size of both are small, then convert sparse to dense and proceed as mentioned; if not, consider the following
• Extract the non-zero entries from the dense matrix (they should be in the same order as for the sparse) and directly compare them, akin to something like:

d # dense numpy matrix
s # sparse numpy matrix
d_nonzero = d[np.nonzero(d)]
np.allclose(d_nonzero, s.data)

I think doing some benchmarking of both approaches to understand the thresholds of interest for the pathways is likely something we want to do here too, as there's likely no one-size-fits-all strategy.

[mlxd]

I'd suggest modifying this (and related) tests to avoid keeping the sparse matrix fully dense --- we want to make sure the lack of entries in the sparse rep doesn't cause problems. I'd recommend looking into the (data, index, indptr) constructor for the sparse rep and using this to build some additional examples.

[JerryChen97]

True. I'll make a better test for sparse case here later.