Add `Op` to solve batched linear matrix equation

[purna135]

Addressing #791

This PR adds Op to aesara.tensor.nlinalg that wraps numpy.linalg.solve for solving linear matrix equations, or system of linear scalar equations, and computing the reverse-mode gradients.

Now, since there is already a Solve in aesara.tensor.slinalg, why do we need this Numpy Solve again?

Scipy doesn't appear to have anything that can be used to solve a linear equation with a batch of triangular matrices, but we can deal with the batched data by using the np.vectorized version of the Scipy Solve. Numpy, on the other hand, has solve, which can handle batches of matrices but is expected to be slower.

So, to decide which one to look at, I create a vectorized version of aesara.tensor.slinalg.cholesky and compare its performance to another cholesky method that deals with batched data using numpy.linalg.solve.

The detailed investigation can be found at the link below.

1. Cholesky - Numpy vs Scipy for Large Array
2. Cholesky - Numpy vs Scipy for Small Array

In short, I obtained the following result:

1. Runtime of Numpy vs Scipy for Large Array

2. Runtime of Numpy vs Scipy for Small Array

As we can see, cholesky which uses numpy solve is more performant for reasonably sized arrays and has simpler code, so we decided to use numpy.linalg.solve instead of Scipy solve.

[codecov[bot]]

Codecov Report

Merging #1060 (e95fa12) into main (8794f48) will increase coverage by 0.01%. The diff coverage is 100.00%.

:exclamation: Current head e95fa12 differs from pull request most recent head 79e45be. Consider uploading reports for the commit 79e45be to get more accurate results

@@            Coverage Diff             @@
##             main    #1060      +/-   ##
==========================================
+ Coverage   79.25%   79.27%   +0.01%     
==========================================
  Files         152      152              
  Lines       47875    47921      +46     
  Branches    10906    10911       +5     
==========================================
+ Hits        37942    37988      +46     
  Misses       7436     7436              
  Partials     2497     2497              

Impacted Files	Coverage Δ
aesara/tensor/nlinalg.py	98.74% <100.00%> (+0.16%)	:arrow_up:

[purna135]

Hello, @brandonwillard. Please accept my apologies for forgetting to link my analysis on this; I have now updated the description of this PR with my findings on both solve methods. please take a look at them.

[brandonwillard]

Hello, @brandonwillard. Please accept my apologies for forgetting to link my analysis on this; I have now updated the description of this PR with my findings on both solve methods. please take a look at them.

The questions I'm referring to have to do with how NumPy's solve is added. There's enough inherent overlap between the existing Solve and this one to prioritize code reuse, generalization, etc.

[purna135]

Can we completely replace the existence solve with this Numpy solve ?

[brandonwillard]

Can we completely replace the existence solve with this Numpy solve ?

I don't think so; I recall each as having distinct options/approaches.

[ricardoV94]

Hello, @brandonwillard. Please accept my apologies for forgetting to link my analysis on this; I have now updated the description of this PR with my findings on both solve methods. please take a look at them.

@purna135 did you benchmark cholesky or solve? I remember you only analysing cholesky and that's what your linked notebook does.

It seems though that scipy cholesky also offers one argument that numpy does not (lower), not sure how important is to keep that one around. Anyway that's not the focus of this PR. Edit: the default of scipy and numpy are not even the same :/

[ricardoV94]

For solve, since scipy has many more options than numpy, I would think that the vectorize approach is more straightforward.

[purna135]

@purna135 did you benchmark cholesky or solve?

Yes, I only tested cholesky, but since cholesky is dependent on solve, I thought it would be better to work on solve in a separate PR.

[ricardoV94]

@purna135 did you benchmark cholesky or solve?

Yes, I only tested cholesky, but since cholesky is dependent on solve, I thought it would be better to work on solve in a separate PR.

I agree, it's just that the linked notebook doesn't matter for this PR because it doesn't say anything about Solve

[Sayam753]

@purna135 can you do similar analysis benchmarking the scipy and numpy solve Ops?

[purna135]

Yes, let me test solve op.

[ricardoV94]

I think the benchmark doesn't matter as much here, because scipy has more options that have no numpy equivalent.

[Sayam753]

@purna135 let's revisit the scipy solve Op and explore the option of add broadcasting there.

[purna135]

@purna135 let's revisit the scipy solve Op and explore the option of add broadcasting there. Ok

[purna135]

Although the benchmark is no longer relevant, I tested both solve methods out of curiosity and found that the Numpy implementation of solve is 2-3 times faster than the vectorized version of Scipy solve.

The detailed analysis can be found here: https://github.com/purna135/GSoC-Learnings/blob/main/Performance%20Comparison%20for%20Solve%20method%20-%20Numpy%20vs%20Scipy.ipynb

[brandonwillard]

Although the benchmark is no longer relevant, I tested both solve methods out of curiosity and found that the Numpy implementation of solve is 2-3 times faster than the vectorized version of Scipy solve.

A numpy.vectorized version of a ufunc or C-level vectorized implementation will almost necessarily be slower (e.g. because the former is just a Python for-loop with additional bookkeeping) , so that's to be expected.

Aside from that, it seems like those comparisons are only demonstrating the difference between the implementations and/or defaults of numpy.linalg.solve and scipy.linalg.solve for sparse/broadcasted identity matrices. Does that sound correct?

[purna135]

Aside from that, it seems like those comparisons are only demonstrating the difference between the implementations and/or defaults of numpy.linalg.solve and scipy.linalg.solve for sparse/broadcasted identity matrices. Does that sound correct?

Yes, exactly