Numba: Add several remaining sparse ops

[tomicapretto]

Description

This PR implements the following sparse ops in numba:

• ColScaleCSC
• RowScaleCSC
• GetItemList and GetItemListGrad
• GetItem2Lists and GetItem2ListsGrad
• GetItem2d
• GetItemScalar
• Neg
• Diag
• SquareDiag

The following ops are not implemented, but I'm not sure if they should be:

• CSMGrad, is it used?
• EnsureSortedIndices, I see it used in the clean function together with remove0, but I don't see clean being used.
• Remove0, see previous line
• ConstructSparseFromList AFAIK it's only used in the gradient of AdvancedSubtensor1 when sparse_grad is True. How relevant is this?

I've been assisted by Codex to perform the implementations and I've reviewed everything and included modifications when needed. To my best understanding, implementations are OK, but of course this is open to feedback.

Related Issue

• Closes #
• Related to #

Checklist

• Checked that the pre-commit linting/style checks pass
• Included tests that prove the fix is effective or that the new feature works
• Added necessary documentation (docstrings and/or example notebooks)
• If you are a pro: each commit corresponds to a relevant logical change

Type of change

• New feature / enhancement
• Bug fix
• Documentation
• Maintenance
• Other (please specify):

[tomicapretto]

I'm starting to thing that it would be a good idea to have many of the implementations be at pytensor/link/numba/dispatch/sparse/variable.py instead of at pytensor/link/numba/dispatch/sparse/basic.py. In basic.py we could just call the underlying method. This would make it more straightforward to migrate to numba-sparse in the future.

[ricardoV94]

Why is Neg its own Op, instead of how other Elemwise are handled?

[tomicapretto]

Why is Neg its own Op, instead of how other Elemwise are handled?

I have no idea. Do you mean that we could have something like

@structured_elemwise(ptm.neg)
def neg(x):
    """
    Compute neg(x) for all non-zero elements of x.
    """

?

Update: I guess the author tried to follow a consistent pattern at the time of making the linked commit, and we inherited it until now you raise the question.

42080b8

[ricardoV94]

Yeah let's change it to be structured_elemwise. I don't love how the elemwise look now (there's an open issue) but no reason why this should be special cased. If you implement like the other Elemwise it will work out of the box

[tomicapretto]

Yeah let's change it to be structured_elemwise. I don't love how the elemwise look now (there's an open issue) but no reason why this should be special cased. If you implement like the other Elemwise it will work out of the box

What about doing it in a separate PR? Removing this Neg op requires the modification of several files, including mlx stuff. nevermind, I mixed sparse with scalar

[tomicapretto]

@ricardoV94, in efc3700 I removed Neg and added a neg function decorated as structured_elemwise.
I also added a test following the pattern I found in tests/sparse/test_math.py for elemwise operations, but I'm afraid it's not running, like all other elemwise tests using the elemwise_checker function.

[ricardoV94]

This could be just a helper that builds the equivalent symbolic operations. I don't see the advantage of being its own Op?

def square_diagonal(diag):
  n = diag.shape[0]
  indices = pt.arange(n, dtype=np.int32)
  indptr = pt.arange(n + 1, dtype=np.int32)
  ... whatever op builds a csc_matrix symbolically)

[ricardoV94]

Or do we even need it? Does it exist in scipy / or is it used anywhere else internally?

[tomicapretto]

I think we can overload eye and then use it within this Op. Would that make sense?