Make Sparse tensor Types extend TensorType

[aerubanov]

This pull request is a continuation of the work in PR #303 on refactors the sparse tensor type structure so that sparse types extend the existing tensor types.

Closes #142.

[aerubanov]

I did a rebase of the branch to the master, which resulted in some of the tests being broken. Now I'm fixing the bugs that occurred, but I decided to open this PR to be able to track progress and discuss changes.

[codecov]

Codecov Report

Merging #766 (d7179fc) into main (31ab8fd) will increase coverage by 0.02%.
The diff coverage is 84.10%.

@@            Coverage Diff             @@
##             main     #766      +/-   ##
==========================================
+ Coverage   78.48%   78.50%   +0.02%     
==========================================
  Files         154      154              
  Lines       47773    47857      +84     
  Branches    10853    10879      +26     
==========================================
+ Hits        37494    37572      +78     
  Misses       7743     7743              
- Partials     2536     2542       +6     

Impacted Files	Coverage Δ
aesara/link/c/type.py	75.66% <ø> (ø)
aesara/sparse/sandbox/sp2.py	82.79% <ø> (ø)
aesara/tensor/blas.py	79.61% <26.31%> (-1.06%)	⬇️
aesara/sparse/type.py	72.11% <81.08%> (+0.85%)	⬆️
aesara/sparse/basic.py	82.47% <94.44%> (+0.39%)	⬆️
aesara/misc/may_share_memory.py	71.42% <100.00%> (ø)
aesara/sparse/opt.py	75.60% <100.00%> (+0.14%)	⬆️
aesara/sparse/sharedvar.py	94.11% <100.00%> (ø)
aesara/tensor/basic.py	87.63% <100.00%> (+0.02%)	⬆️
aesara/tensor/basic_opt.py	85.03% <100.00%> (ø)
... and 8 more

[aerubanov]

@brandonwillard , Looks like I fixed tests that were broken after rebase. In your initial PR #303 you noticed that we should "Convert sparse to dense when used with some Ops (e.g. Elemwise, CAReduce)". Could you provide more details about this?

[brandonwillard]

@brandonwillard , Looks like I fixed tests that were broken after rebase. In your initial PR #303 you noticed that we should "Convert sparse to dense when used with some Ops (e.g. Elemwise, CAReduce)". Could you provide more details about this?

In order to make the old TensorVariable interface, _tensor_py_operators, applicable to sparse matrices, we need to start adding support for currently unsupported sparse operations, like __get__, which constructs all sorts of Subtensor Ops. Right now, sparse matrices implement their own indexing Ops, which is a huge interface problem. We can't reasonably add sparse matrix-specific logic to Subtensor, though.

We can get around this by dispatching to the appropriate sparse Ops at the _sparse_py_operators interface level, but then there are still some related operations that aren't covered, like advanced indexing on sparse matrices. If we simply convert sparse matrices to dense matrices at that interface, we can—at the very least—make things work (i.e. not return graphs that will err when evaluated).

Really, that's all we need to do in order to merge this: i.e. make sure the new interface doesn't generate graphs that necessarily produces errors, and converting to a dense matrix is a quick and dirty work-around in most cases.

[aerubanov]

@brandonwillard , I added tests to check that no errors occurred when using _tensor_py_operators with sparse tensors and added conversion sparse to dense to make them pass. Also, when the conversion happens the warning is generated to inform users about that. Please look at the code and let me know what additional improvements are needed.

[brandonwillard]

It looks like you have all the necessary logic in there. I left a few comments about where certain steps take place (e.g. in the sparse module vs. others). The basic idea is that we want to keep all sparse object related handling within the sparse module. In other words, we don't want the concerns of sparse types to leak into code that doesn't explicitly need to concern itself with sparse types.

My guess is that those lines of code can be moved into the sparse Variable interface before the external code is called.

Aside from that, we'll need to rebase some of those commits (e.g. combine test-adding commits with their respective code changes, merge/remove commits that undo things performed by earlier commits in this branch, etc.) and update their formatting to match these guidelines (just the summary lines for now). Since this is such a large set of changes, these kinds of things are even more important—both for the reviewing overhead and basic organizational reasons.

[aerubanov]

Hi @brandonwillard ! I made the changes according to your comments. Now I override in _sparse_py_operators unimplemented operations for sparse variables, converting them to dense. So this logic stays inside the sparse module. And I cleaned up the commit history.

[brandonwillard]

Hi @brandonwillard ! I made the changes according to your comments. Now I override in _sparse_py_operators unimplemented operations for sparse variables, converting them to dense. So this logic stays inside the sparse module. And I cleaned up the commit history.

Great, thanks! I'm still a bit backlogged, but I'll review soon.

[aerubanov]

Hi @brandonwillard! I just wanted to remind you of this PR just in case. Of course, if you have time for that.

[brandonwillard]

It looks like there were some problems rebasing, because there are a few reversions of upstream changes and unmanaged conflict diffs.

[brandonwillard]

Suggested change

	fgraph.replace(r, new_r, reason=reason, verbose=verbose, **kwargs)
	fgraph.replace(r, new_r, reason=reason, verbose=False, **kwargs)

verbose is set to False here because this method handles that functionality itself.

[brandonwillard]

I'll fix the rebase issues and push in a minute.

[brandonwillard]

There's still more to be done to get the sparse interface into shape, but these changes look like a good step toward that, if only because they explicitly state that we want a more correct type representation.

The biggest issue right now is that TensorTypes aren't an abstract class, and direct instances of TensorTypes represent dense tensors, alongside instances of the new DenseTensor type. It might not take much to make such a change and remove all the meta classes, though.

Anyway, I'm fine moving forward with these changes and opening another set of issues to follow up, but let's get a few more eyes on this first.

[kc611]

LGTM.

The biggest issue right now is that TensorTypes aren't an abstract class, and direct instances of TensorTypes represent dense tensors, alongside instances of the new DenseTensor type. It might not take much to make such a change and remove all the meta classes, though.

We should make an issue page for this. Since, this can be tackled separately from the additional logic that we add towards supporting SparseTypes.

[aerubanov]

Hi @brandonwillard! I rebased on top of the master and resolved the conflicts.

[brandonwillard]

The location and name of the new tests.test_tensor_interface needs to be reconsidered, because it's explicitly testing interfaces that already have corresponding test sub-packages/modules (e.g. tests.sparse and tests.tensor.test_var). If the is_sparse == False case isn't needed (see here), then this module could be relocated/renamed to tests.sparse.test_var.

[brandonwillard]

I've updated the new tests and added a commit that renames SparseType to SparseTensorType (to match the newly introduced DenseTensorType).

[aerubanov]

@brandonwillard, Thanks for the review and the corrections!

[brandonwillard]

@aerubanov, thanks for the PR! This was an important one that required a lot of persistence. It's this kind of help that really moves things forward, and we truly appreciate it.