Extend the implementation of csr_matrix and csc_matrix in numba. Use those…

pytensor/link/numba/dispatch/sparse/basic.py

 import numpy as np
 import scipy as sp
-from numba.core import types
 from numba.extending import overload
 
 from pytensor import config
@@ -96,9 +95,7 @@ def numba_funcify_DenseFromSparse(op, node, **kwargs):
 
 @register_funcify_default_op_cache_key(SparseFromDense)
 def numba_funcify_SparseFromDense(op, node, **kwargs):
-    sparse_format = op.format
-
-    if sparse_format == "csr":
+    if op.format == "csr":
 
         @numba_basic.numba_njit
         def dense_to_csr(matrix):

pytensor/link/numba/dispatch/sparse/variable.py

@@ -185,33 +185,140 @@ def cs_matrix_constant(context, builder, ty, pyval):
 
 
 @overload(sp.sparse.csr_matrix)
-def overload_csr_matrix(arg1, shape, dtype=None):
-    if not isinstance(arg1, types.BaseAnonymousTuple) or len(arg1) != 3:
-        return None
-    if isinstance(shape, types.NoneType):
-        return None
+def overload_csr_matrix(arg1, shape=None, dtype=None):
+    if isinstance(arg1, CSCMatrixType):
 
-    def impl(arg1, shape, dtype=None):
+        def csr_from_csc(arg1, shape=None, dtype=None):
+            return arg1.tocsr()
+
+        return csr_from_csc
+
+    if isinstance(arg1, CSRMatrixType):
+
+        def csr_from_csr(arg1, shape=None, dtype=None):
+            return arg1.copy()
+
+        return csr_from_csr
+
+    if isinstance(arg1, types.Array) and arg1.ndim == 2:
+
+        def csr_from_dense(arg1, shape=None, dtype=None):
+            n_rows = types.int32(arg1.shape[0])
+            n_cols = types.int32(arg1.shape[1])
+
+            # Pass 1: Count non-zeros to pre-allocate
+            nnz = 0
+            for i in range(n_rows):
+                for j in range(n_cols):
+                    if arg1[i, j] != 0:
+                        nnz += 1
+
+            if dtype is not None:
+                data_dtype = dtype
+            else:
+                data_dtype = arg1.dtype
+            data = np.empty(nnz, dtype=data_dtype)
+            indices = np.empty(nnz, dtype=np.int32)
+            indptr = np.zeros(n_rows + 1, dtype=np.int32)
+
+            # Pass 2: Fill the arrays
+            pos = 0
+            for i in range(n_rows):
+                for j in range(n_cols):
+                    value = arg1[i, j]
+                    if value != 0:
+                        data[pos] = value
+                        indices[pos] = j
+                        pos += 1
+                indptr[i + 1] = pos
+
+            return csr_matrix_from_components(data, indices, indptr, (n_rows, n_cols))
+
+        return csr_from_dense
+
+    if (
+        isinstance(arg1, types.BaseAnonymousTuple)
+        and len(arg1) == 3
+        and shape is not None
+    ):
+
+        def csr_from_components(arg1, shape=None, dtype=None):
             data, indices, indptr = arg1
             int32_shape = (types.int32(shape[0]), types.int32(shape[1]))
             return csr_matrix_from_components(data, indices, indptr, int32_shape)
 
-    return impl
+        return csr_from_components
 
 
 @overload(sp.sparse.csc_matrix)
-def overload_csc_matrix(arg1, shape, dtype=None):
-    if not isinstance(arg1, types.BaseAnonymousTuple) or len(arg1) != 3:
-        return None
-    if isinstance(shape, types.NoneType):
-        return None
+def overload_csc_matrix(arg1, shape=None, dtype=None):
+    if isinstance(arg1, CSRMatrixType):
 
-    def impl(arg1, shape, dtype=None):
+        def csc_from_csr(arg1, shape=None, dtype=None):
+            return arg1.tocsc()
+
+        return csc_from_csr
+
+    if isinstance(arg1, CSCMatrixType):
+
+        def csc_from_csc(arg1, shape=None, dtype=None):
+            return arg1.copy()
+
+        return csc_from_csc
+
+    if isinstance(arg1, types.Array) and arg1.ndim == 2:
+
+        def csc_from_dense(arg1, shape=None, dtype=None):
+            if shape is not None:
+                n_rows = types.int32(shape[0])
+                n_cols = types.int32(shape[1])
+            else:
+                n_rows = types.int32(arg1.shape[0])
+                n_cols = types.int32(arg1.shape[1])
+
+            # Pass 1: Count non-zeros to pre-allocate
+            nnz = 0
+            for j in range(n_cols):
+                for i in range(n_rows):
+                    if arg1[i, j] != 0:
+                        nnz += 1
+
+            # Pre-allocate internal containers
+            if dtype is not None:
+                data_dtype = dtype
+            else:
+                data_dtype = arg1.dtype
+            data = np.empty(nnz, dtype=data_dtype)
+            indices = np.empty(nnz, dtype=np.int32)
+            indptr = np.zeros(n_cols + 1, dtype=np.int32)
+
+            # Pass 2: Fill the arrays
+            pos = 0
+            for j in range(n_cols):
+                for i in range(n_rows):
+                    value = arg1[i, j]
+                    if value != 0:
+                        data[pos] = value
+                        indices[pos] = i
+                        pos += 1
+                indptr[j + 1] = pos
+
+            return csc_matrix_from_components(data, indices, indptr, (n_rows, n_cols))
+
+        return csc_from_dense
+
+    if (
+        isinstance(arg1, types.BaseAnonymousTuple)
+        and len(arg1) == 3
+        and shape is not None
+    ):
+
+        def csc_from_components(arg1, shape=None, dtype=None):
             data, indices, indptr = arg1
             int32_shape = (types.int32(shape[0]), types.int32(shape[1]))
             return csc_matrix_from_components(data, indices, indptr, int32_shape)
 
-    return impl
+        return csc_from_components
 
 
 @overload(np.shape)