Finish Copilot code

Author: jaanerik

pytensor/link/jax/dispatch/subtensor.py

@@ -31,11 +31,18 @@
 """
 
 
+@jax_funcify.register(AdvancedSubtensor1)
+def jax_funcify_AdvancedSubtensor1(op, node, **kwargs):
+    def advanced_subtensor1(x, ilist):
+        return x[ilist]
+
+    return advanced_subtensor1
+
+
 @jax_funcify.register(Subtensor)
 @jax_funcify.register(AdvancedSubtensor)
-@jax_funcify.register(AdvancedSubtensor1)
 def jax_funcify_Subtensor(op, node, **kwargs):
-    idx_list = getattr(op, "idx_list", None)
+    idx_list = op.idx_list
 
     def subtensor(x, *ilists):
         indices = indices_from_subtensor(ilists, idx_list)
@@ -47,10 +54,24 @@ def subtensor(x, *ilists):
     return subtensor
 
 
-@jax_funcify.register(IncSubtensor)
 @jax_funcify.register(AdvancedIncSubtensor1)
+def jax_funcify_AdvancedIncSubtensor1(op, node, **kwargs):
+    if getattr(op, "set_instead_of_inc", False):
+
+        def jax_fn(x, y, ilist):
+            return x.at[ilist].set(y)
+
+    else:
+
+        def jax_fn(x, y, ilist):
+            return x.at[ilist].add(y)
+
+    return jax_fn
+
+
+@jax_funcify.register(IncSubtensor)
 def jax_funcify_IncSubtensor(op, node, **kwargs):
-    idx_list = getattr(op, "idx_list", None)
+    idx_list = op.idx_list
 
     if getattr(op, "set_instead_of_inc", False):
 
@@ -77,8 +98,8 @@ def incsubtensor(x, y, *ilist, jax_fn=jax_fn, idx_list=idx_list):
 
 @jax_funcify.register(AdvancedIncSubtensor)
 def jax_funcify_AdvancedIncSubtensor(op, node, **kwargs):
-    idx_list = getattr(op, "idx_list", None)
-    
+    idx_list = op.idx_list
+
     if getattr(op, "set_instead_of_inc", False):
 
         def jax_fn(x, indices, y):

pytensor/link/numba/dispatch/subtensor.py

@@ -20,7 +20,6 @@
 )
 from pytensor.link.numba.dispatch.compile_ops import numba_deepcopy
 from pytensor.tensor import TensorType
-from pytensor.tensor.rewriting.subtensor import is_full_slice
 from pytensor.tensor.subtensor import (
     AdvancedIncSubtensor,
     AdvancedIncSubtensor1,
@@ -29,7 +28,7 @@
     IncSubtensor,
     Subtensor,
 )
-from pytensor.tensor.type_other import MakeSlice, NoneTypeT, SliceType
+from pytensor.tensor.type_other import MakeSlice
 
 
 def slice_new(self, start, stop, step):
@@ -239,15 +238,15 @@ def {function_name}({", ".join(input_names)}):
 @register_funcify_and_cache_key(AdvancedIncSubtensor)
 def numba_funcify_AdvancedSubtensor(op, node, **kwargs):
     if isinstance(op, AdvancedSubtensor):
-        x, y, tensor_inputs = node.inputs[0], None, node.inputs[1:]
+        tensor_inputs = node.inputs[1:]
     else:
-        x, y, *tensor_inputs = node.inputs
+        tensor_inputs = node.inputs[2:]
 
     # Reconstruct indexing information from idx_list and tensor inputs
     basic_idxs = []
     adv_idxs = []
     input_idx = 0
-    
+
     for i, entry in enumerate(op.idx_list):
         if isinstance(entry, slice):
             # Basic slice index
@@ -256,12 +255,14 @@ def numba_funcify_AdvancedSubtensor(op, node, **kwargs):
             # Advanced tensor index
             if input_idx < len(tensor_inputs):
                 idx_input = tensor_inputs[input_idx]
-                adv_idxs.append({
-                    "axis": i,
-                    "dtype": idx_input.type.dtype,
-                    "bcast": idx_input.type.broadcastable,
-                    "ndim": idx_input.type.ndim,
-                })
+                adv_idxs.append(
+                    {
+                        "axis": i,
+                        "dtype": idx_input.type.dtype,
+                        "bcast": idx_input.type.broadcastable,
+                        "ndim": idx_input.type.ndim,
+                    }
+                )
                 input_idx += 1
 
     # Special implementation for consecutive integer vector indices

pytensor/link/pytorch/dispatch/subtensor.py

@@ -9,7 +9,7 @@
     Subtensor,
     indices_from_subtensor,
 )
-from pytensor.tensor.type_other import MakeSlice, SliceType
+from pytensor.tensor.type_other import MakeSlice
 
 
 def check_negative_steps(indices):
@@ -63,8 +63,8 @@ def makeslice(start, stop, step):
 @pytorch_funcify.register(AdvancedSubtensor1)
 @pytorch_funcify.register(AdvancedSubtensor)
 def pytorch_funcify_AdvSubtensor(op, node, **kwargs):
-    idx_list = getattr(op, "idx_list", None)
-    
+    idx_list = op.idx_list
+
     def advsubtensor(x, *flattened_indices):
         indices = indices_from_subtensor(flattened_indices, idx_list)
         check_negative_steps(indices)
@@ -105,7 +105,7 @@ def inc_subtensor(x, y, *flattened_indices):
 @pytorch_funcify.register(AdvancedIncSubtensor)
 @pytorch_funcify.register(AdvancedIncSubtensor1)
 def pytorch_funcify_AdvancedIncSubtensor(op, node, **kwargs):
-    idx_list = getattr(op, "idx_list", None)
+    idx_list = op.idx_list
     inplace = op.inplace
     ignore_duplicates = getattr(op, "ignore_duplicates", False)
 
@@ -139,7 +139,9 @@ def adv_inc_subtensor_no_duplicates(x, y, *flattened_indices):
 
     else:
         # Check if we have slice indexing in idx_list
-        has_slice_indexing = any(isinstance(entry, slice) for entry in idx_list) if idx_list else False
+        has_slice_indexing = (
+            any(isinstance(entry, slice) for entry in idx_list) if idx_list else False
+        )
         if has_slice_indexing:
             raise NotImplementedError(
                 "IncSubtensor with potential duplicates indexes and slice indexing not implemented in PyTorch"

pytensor/tensor/basic.py

@@ -1818,6 +1818,33 @@ def do_constant_folding(self, fgraph, node):
         return True
 
 
+@_vectorize_node.register(Alloc)
+def vectorize_alloc(op: Alloc, node: Apply, batch_val, *batch_shapes):
+    # batch_shapes are usually not batched (they are scalars for the shape)
+    # batch_val is the value being allocated.
+
+    # If shapes are batched, we fall back (complex case)
+    if any(
+        b_shp.type.ndim > shp.type.ndim
+        for b_shp, shp in zip(batch_shapes, node.inputs[1:], strict=True)
+    ):
+        return vectorize_node_fallback(op, node, batch_val, *batch_shapes)
+
+    # If value is batched, we need to prepend batch dims to the output shape
+    val = node.inputs[0]
+    batch_ndim = batch_val.type.ndim - val.type.ndim
+
+    if batch_ndim == 0:
+        return op.make_node(batch_val, *batch_shapes)
+
+    # We need the size of the batch dimensions
+    # batch_val has shape (B1, B2, ..., val_dims...)
+    batch_dims = [batch_val.shape[i] for i in range(batch_ndim)]
+
+    new_shapes = batch_dims + list(batch_shapes)
+    return op.make_node(batch_val, *new_shapes)
+
+
 alloc = Alloc()
 pprint.assign(alloc, printing.FunctionPrinter(["alloc"]))
 

pytensor/tensor/rewriting/subtensor.py

@@ -14,6 +14,7 @@
     in2out,
     node_rewriter,
 )
+from pytensor.graph.type import Type
 from pytensor.raise_op import Assert
 from pytensor.scalar import Add, ScalarConstant, ScalarType
 from pytensor.scalar import constant as scalar_constant
@@ -212,6 +213,20 @@ def get_advsubtensor_axis(indices):
         return axis
 
 
+def reconstruct_indices(idx_list, tensor_inputs):
+    """Reconstruct indices from idx_list and tensor inputs."""
+    indices = []
+    input_idx = 0
+    for entry in idx_list:
+        if isinstance(entry, slice):
+            indices.append(entry)
+        elif isinstance(entry, Type):
+            if input_idx < len(tensor_inputs):
+                indices.append(tensor_inputs[input_idx])
+                input_idx += 1
+    return indices
+
+
 @register_specialize
 @node_rewriter([AdvancedSubtensor])
 def local_replace_AdvancedSubtensor(fgraph, node):
@@ -229,17 +244,9 @@ def local_replace_AdvancedSubtensor(fgraph, node):
 
     indexed_var = node.inputs[0]
     tensor_inputs = node.inputs[1:]
-    
+
     # Reconstruct indices from idx_list and tensor inputs
-    indices = []
-    input_idx = 0
-    for entry in node.op.idx_list:
-        if isinstance(entry, slice):
-            indices.append(entry)
-        elif isinstance(entry, Type):
-            if input_idx < len(tensor_inputs):
-                indices.append(tensor_inputs[input_idx])
-                input_idx += 1
+    indices = reconstruct_indices(node.op.idx_list, tensor_inputs)
 
     axis = get_advsubtensor_axis(indices)
 
@@ -267,17 +274,9 @@ def local_AdvancedIncSubtensor_to_AdvancedIncSubtensor1(fgraph, node):
     res = node.inputs[0]
     val = node.inputs[1]
     tensor_inputs = node.inputs[2:]
-    
+
     # Reconstruct indices from idx_list and tensor inputs
-    indices = []
-    input_idx = 0
-    for entry in node.op.idx_list:
-        if isinstance(entry, slice):
-            indices.append(entry)
-        elif isinstance(entry, Type):
-            if input_idx < len(tensor_inputs):
-                indices.append(tensor_inputs[input_idx])
-                input_idx += 1
+    indices = reconstruct_indices(node.op.idx_list, tensor_inputs)
 
     axis = get_advsubtensor_axis(indices)
 
@@ -1112,6 +1111,7 @@ def local_inplace_AdvancedIncSubtensor1(fgraph, node):
 def local_inplace_AdvancedIncSubtensor(fgraph, node):
     if isinstance(node.op, AdvancedIncSubtensor) and not node.op.inplace:
         new_op = type(node.op)(
+            node.op.idx_list,
             inplace=True,
             set_instead_of_inc=node.op.set_instead_of_inc,
             ignore_duplicates=node.op.ignore_duplicates,
@@ -1376,6 +1376,7 @@ def local_useless_inc_subtensor_alloc(fgraph, node):
                     z_broad[k]
                     and not same_shape(xi, y, dim_x=k, dim_y=k)
                     and shape_of[y][k] != 1
+                    and shape_of[xi][k] == 1
                 )
             ]
 
@@ -1778,17 +1779,9 @@ def ravel_multidimensional_bool_idx(fgraph, node):
     else:
         x, y = node.inputs[0], node.inputs[1]
         tensor_inputs = node.inputs[2:]
-    
+
     # Reconstruct indices from idx_list and tensor inputs
-    idxs = []
-    input_idx = 0
-    for entry in node.op.idx_list:
-        if isinstance(entry, slice):
-            idxs.append(entry)
-        elif isinstance(entry, Type):
-            if input_idx < len(tensor_inputs):
-                idxs.append(tensor_inputs[input_idx])
-                input_idx += 1
+    idxs = reconstruct_indices(node.op.idx_list, tensor_inputs)
 
     if any(
         (
@@ -1829,36 +1822,36 @@ def ravel_multidimensional_bool_idx(fgraph, node):
         # Create new AdvancedSubtensor with updated idx_list
         new_idx_list = list(node.op.idx_list)
         new_tensor_inputs = list(tensor_inputs)
-        
+
         # Update the idx_list and tensor_inputs for the raveled boolean index
         input_idx = 0
         for i, entry in enumerate(node.op.idx_list):
             if isinstance(entry, Type):
                 if input_idx == bool_idx_pos:
                     new_tensor_inputs[input_idx] = raveled_bool_idx
                 input_idx += 1
-        
+
         new_out = AdvancedSubtensor(new_idx_list)(raveled_x, *new_tensor_inputs)
     else:
         # Create new AdvancedIncSubtensor with updated idx_list
         new_idx_list = list(node.op.idx_list)
         new_tensor_inputs = list(tensor_inputs)
-        
+
         # Update the tensor_inputs for the raveled boolean index
         input_idx = 0
         for i, entry in enumerate(node.op.idx_list):
             if isinstance(entry, Type):
                 if input_idx == bool_idx_pos:
                     new_tensor_inputs[input_idx] = raveled_bool_idx
                 input_idx += 1
-        
+
         # The dimensions of y that correspond to the boolean indices
         # must already be raveled in the original graph, so we don't need to do anything to it
         new_out = AdvancedIncSubtensor(
             new_idx_list,
             inplace=node.op.inplace,
             set_instead_of_inc=node.op.set_instead_of_inc,
-            ignore_duplicates=node.op.ignore_duplicates
+            ignore_duplicates=node.op.ignore_duplicates,
         )(raveled_x, y, *new_tensor_inputs)
         # But we must reshape the output to match the original shape
         new_out = new_out.reshape(x_shape)