draft commit to push on CI

Author: ferrine

pytensor/sparse/basic.py

@@ -954,7 +954,7 @@ class DenseFromSparse(Op):
 
     """
 
-    __props__ = ()
+    __props__ = ("sparse_grad",)
 
     def __init__(self, structured=True):
         self.sparse_grad = structured

pytensor/sparse/rewriting.py

@@ -1099,13 +1099,13 @@ def c_code_cache_version(self):
 csm_grad_c = CSMGradC()
 
 
-@node_rewriter([csm_grad(None)])
+@node_rewriter([csm_grad()])
 def local_csm_grad_c(fgraph, node):
     """
     csm_grad(None) -> csm_grad_c
 
     """
-    if node.op == csm_grad(None):
+    if node.op == csm_grad():
         return [csm_grad_c(*node.inputs)]
     return False
 

pytensor/sparse/type.py

@@ -74,7 +74,10 @@ def __init__(
     ):
         if shape is None and broadcastable is None:
             shape = (None, None)
-
+        if broadcastable is None:
+            broadcastable = (False, False)
+        if broadcastable != (False, False):
+            raise ValueError("Broadcasting sparse types is not yet implemented")
         if format not in self.format_cls:
             raise ValueError(
                 f'unsupported format "{format}" not in list',
@@ -96,7 +99,9 @@ def clone(
             dtype = self.dtype
         if shape is None:
             shape = self.shape
-        return type(self)(format, dtype, shape=shape, **kwargs)
+        return type(self)(
+            format, dtype, shape=shape, broadcastable=broadcastable, **kwargs
+        )
 
     def filter(self, value, strict=False, allow_downcast=None):
         if isinstance(value, Variable):

pytensor/tensor/shape.py

@@ -1,7 +1,7 @@
 import warnings
 from numbers import Number
 from textwrap import dedent
-from typing import Dict, List, Tuple, Union
+from typing import Dict, List, Sequence, Tuple, Union
 
 import numpy as np
 
@@ -392,10 +392,15 @@ class SpecifyShape(COp):
     Maybe in the future we will never do the assert!
     """
 
+    __props__ = ("output_broadcastable",)
     view_map = {0: [0]}
     __props__ = ()
     _f16_ok = True
 
+    def __init__(self, output_broadcastable: Sequence) -> None:
+        super().__init__()
+        self.output_broadcastable = tuple(output_broadcastable)
+
     def make_node(self, x, *shape):
         from pytensor.tensor.basic import get_scalar_constant_value
 
@@ -432,7 +437,9 @@ def make_node(self, x, *shape):
                 except NotScalarConstantError:
                     pass
 
-        out_var = x.type.clone(shape=type_shape)()
+        out_var = x.type.clone(
+            shape=type_shape, broadcastable=self.output_broadcastable
+        )()
 
         return Apply(self, [x, *shape], [out_var])
 
@@ -539,12 +546,10 @@ def c_code_cache_version(self):
         return (2,)
 
 
-_specify_shape = SpecifyShape()
-
-
 def specify_shape(
     x: Union[np.ndarray, Number, Variable],
     shape: Union[ShapeValueType, List[ShapeValueType], Tuple[ShapeValueType, ...]],
+    broadcastable=None,
 ):
     """Specify a fixed shape for a `Variable`.
 
@@ -574,8 +579,9 @@ def specify_shape(
     # If shape does not match x.ndim, we rely on the `Op` to raise a ValueError
     if not new_shape_info and len(shape) == x.type.ndim:
         return x
-
-    return _specify_shape(x, *shape)
+    if broadcastable is None:
+        broadcastable = x.broadcastable
+    return SpecifyShape(broadcastable)(x, *shape)
 
 
 @_get_vector_length.register(SpecifyShape)
@@ -934,7 +940,11 @@ def specify_broadcastable(x, *axes):
         raise ValueError("Trying to specify broadcastable of non-existent dimension")
 
     shape_info = [1 if i in axes else s for i, s in enumerate(x.type.shape)]
-    return specify_shape(x, shape_info)
+    broadcastable = [
+        True if i in axes else b for i, b in enumerate(x.type.broadcastable)
+    ]
+
+    return specify_shape(x, shape_info, broadcastable=broadcastable)
 
 
 class Unbroadcast(COp):

pytensor/tensor/var.py

@@ -1018,7 +1018,7 @@ def __init__(self, type: _TensorTypeType, data, name=None):
             )
 
         # We want all the shape information from `data`
-        new_type = type.clone(shape=data_shape)
+        new_type = type.clone(shape=data_shape, broadcastable=type.broadcastable)
 
         assert not any(s is None for s in new_type.shape)
 

tests/sparse/test_basic.py

@@ -1093,21 +1093,23 @@ def check_format_ndim(format, ndim):
         s = SparseFromDense(format)(x)
         s_m = -s
         d = dense_from_sparse(s_m)
+        pytensor.grad(None, x, known_grads={d: d.type()})
         c = d.sum()
         g = pytensor.grad(c, x)
         f = pytensor.function([x], [s, g])
         f(np.array(0, dtype=config.floatX, ndmin=ndim))
         f(np.array(7, dtype=config.floatX, ndmin=ndim))
 
-    def test_format_ndim(self):
-        for format in "csc", "csr":
-            for ndim in 0, 1, 2:
-                self.check_format_ndim(format, ndim)
+    @pytest.mark.parametrize("format", ["csc", "csr"])
+    @pytest.mark.parametrize("ndim", [0, 1, 2])
+    def test_format_ndim(self, format, ndim):
+        self.check_format_ndim(format, ndim)
 
-            with pytest.raises(TypeError):
-                self.check_format_ndim(format, 3)
-            with pytest.raises(TypeError):
-                self.check_format_ndim(format, 4)
+    @pytest.mark.parametrize("format", ["csc", "csr"])
+    @pytest.mark.parametrize("ndim", [3, 4])
+    def test_format_ndim_raises(self, format, ndim):
+        with pytest.raises(TypeError):
+            self.check_format_ndim(format, ndim)
 
 
 class TestCsmProperties:

tests/tensor/test_shape.py

@@ -19,7 +19,6 @@
     Shape_i,
     SpecifyShape,
     Unbroadcast,
-    _specify_shape,
     reshape,
     shape,
     shape_i,
@@ -350,13 +349,13 @@ def test_check_inputs(self):
             specify_shape([[1, 2, 3], [4, 5, 6]], (2.2, 3))
 
         with pytest.raises(TypeError, match="must be integer types"):
-            _specify_shape([[1, 2, 3], [4, 5, 6]], *(2.2, 3))
+            SpecifyShape([False, False])([[1, 2, 3], [4, 5, 6]], *(2.2, 3))
 
         with pytest.raises(ValueError, match="will never match"):
             specify_shape(matrix(), [4])
 
         with pytest.raises(ValueError, match="will never match"):
-            _specify_shape(matrix(), *[4])
+            SpecifyShape([False, False])(matrix(), *[4])
 
         with pytest.raises(ValueError, match="must have fixed dimensions"):
             specify_shape(matrix(), vector(dtype="int32"))