Add numba overload for Nonzero

[Abhinav-Khot]

Description

Added a numba overload in Nonzero by adding the @nb.njit decorator as np.nonzero is supported by numba already.

Related Issue

• Closes Add numba overload for Nonzero #1279
• 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):

---

📚 Documentation preview 📚: https://pytensor--1289.org.readthedocs.build/en/1289/

[jessegrabowski]

This is not how a numba dispatch is implemented. Please review the relevant documentation if you'd like to tackle this issue: https://pytensor.readthedocs.io/en/latest/extending/creating_a_numba_jax_op.html

[Abhinav-Khot]

This is not how a numba dispatch is implemented. Please review the relevant documentation if you'd like to tackle this issue: https://pytensor.readthedocs.io/en/latest/extending/creating_a_numba_jax_op.html

Hello,

I have tried implementing the backend for numba using the documentation and have written some tests too. Would be grateful if you could review them and give feedback

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 82.00%. Comparing base (2a7f3e1) to head (b336df9).
Report is 36 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1289      +/-   ##
==========================================
+ Coverage   81.99%   82.00%   +0.01%     
==========================================
  Files         188      188              
  Lines       48553    48488      -65     
  Branches     8673     8666       -7     
==========================================
- Hits        39812    39765      -47     
+ Misses       6579     6575       -4     
+ Partials     2162     2148      -14     

Files with missing lines	Coverage Δ
pytensor/link/numba/dispatch/basic.py	78.50% <100.00%> (-0.59%)	⬇️

... and 27 files with indirect coverage changes

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[Abhinav-Khot]

@jessegrabowski Could you please review my changes?

[jessegrabowski]

ndim == 1 is not a special case in the C-backend. In that case, you get a 1-tuple. All backends should return the same thing.

[Abhinav-Khot]

The modification returns also 1-tuple because this is already implemented in the nonzero function itself

class Nonzero(Op):
    """
    Return the indices of the elements that are non-zero.

    Parameters
    ----------
    a: array_like
        Input array.

    Returns
    -------
    indices: list
        A list containing the indices of the non-zero elements of `a`.

    See Also
    --------
    nonzero_values : Return the non-zero elements of the input array
    flatnonzero : Return the indices of the non-zero elements of the
        flattened input array.

    """

    __props__ = ()

    def make_node(self, a):
        a = as_tensor_variable(a)
        if a.ndim == 0:
            raise ValueError("Nonzero only supports non-scalar arrays.")
        output = [TensorType(dtype="int64", shape=(None,))() for i in range(a.ndim)]
        return Apply(self, [a], output)

    def perform(self, node, inp, out_):
        a = inp[0]

        result_tuple = np.nonzero(a)
        for i, res in enumerate(result_tuple):
            out_[i][0] = res.astype("int64")

    def grad(self, inp, grads):
        return [grad_undefined(self, 0, inp[0])]


_nonzero = Nonzero()


def nonzero(a, return_matrix=False):
    """
    Returns one of the following:

        If return_matrix is False (default, same as NumPy):
            A tuple of vector arrays such that the ith element of the jth array
            is the index of the ith non-zero element of the input array in the
            jth dimension.

        If return_matrix is True (same as PyTensor Op):
            Returns a matrix of shape (ndim, number of nonzero elements) such
            that element (i,j) is the index in the ith dimension of the jth
            non-zero element.

    Parameters
    ----------
    a : array_like
        Input array.
    return_matrix : bool
        If True, returns a symbolic matrix. If False, returns a tuple of
        arrays. Defaults to False.

    Returns
    -------
    tuple of vectors or matrix

    See Also
    --------
    nonzero_values : Return the non-zero elements of the input array
    flatnonzero : Return the indices of the non-zero elements of the
        flattened input array.

    """
    res = _nonzero(a)
    if isinstance(res, list):
        res = tuple(res)
    else:
        res = (res,)

    if return_matrix:
        if len(res) > 1:
            return stack(res, 0)
        elif len(res) == 1:
            return shape_padleft(res[0])
    else:
        return res

If we do not handle the ndim = 1 case seperately, for the array [1,2,0] the result would be (([2],),) which is a tuple of tuple of lists. Instead it should be a tuple of lists which is what we get with the modification. I have replaced this modification with :

if(a.ndim == 1):
       return result_tuple[0]

for efficiency.

[jessegrabowski]

This input validation is done by the nonzero Op itself, there's no need to repeat it here.

[jessegrabowski]

Looks great! Thanks for this