Allow broadcasting in specialized numba dispatch of AdvancedIncSubtensor

Author: ricardoV94

pytensor/link/numba/dispatch/subtensor.py

@@ -130,15 +130,6 @@ def numba_funcify_AdvancedSubtensor(op, node, **kwargs):
  if isinstance(idx.type, TensorType)
  ]
 
- def broadcasted_to(x_bcast: tuple[bool, ...], to_bcast: tuple[bool, ...]):
- # Check that x is not broadcasted to y based on broadcastable info
- if len(x_bcast) < len(to_bcast):
- return True
- for x_bcast_dim, to_bcast_dim in zip(x_bcast, to_bcast, strict=True):
- if x_bcast_dim and not to_bcast_dim:
- return True
- return False
-
  # Special implementation for consecutive integer vector indices
  if (
  not basic_idxs
@@ -151,17 +142,6 @@ def broadcasted_to(x_bcast: tuple[bool, ...], to_bcast: tuple[bool, ...]):
  )
  # Must be consecutive
  and not op.non_consecutive_adv_indexing(node)
- # y in set/inc_subtensor cannot be broadcasted
- and (
- y is None
- or not broadcasted_to(
- y.type.broadcastable,
- (
- x.type.broadcastable[: adv_idxs[0]["axis"]]
- + x.type.broadcastable[adv_idxs[-1]["axis"] :]
- ),
- )
- )
  ):
  return numba_funcify_multiple_integer_vector_indexing(op, node, **kwargs)
 
@@ -191,14 +171,24 @@ def broadcasted_to(x_bcast: tuple[bool, ...], to_bcast: tuple[bool, ...]):
  return numba_funcify_default_subtensor(op, node, **kwargs)
 
 
+def _broadcasted_to(x_bcast: tuple[bool, ...], to_bcast: tuple[bool, ...]):
+ # Check that x is not broadcasted to y based on broadcastable info
+ if len(x_bcast) < len(to_bcast):
+ return True
+ for x_bcast_dim, to_bcast_dim in zip(x_bcast, to_bcast, strict=True):
+ if x_bcast_dim and not to_bcast_dim:
+ return True
+ return False
+
+
 def numba_funcify_multiple_integer_vector_indexing(
  op: AdvancedSubtensor | AdvancedIncSubtensor, node, **kwargs
 ):
  # Special-case implementation for multiple consecutive vector integer indices (and set/incsubtensor)
  if isinstance(op, AdvancedSubtensor):
- y, idxs = None, node.inputs[1:]
+ x, y, idxs = None, node.inputs
  else:
- y, *idxs = node.inputs[1:]
+ x, y, *idxs = node.inputs
 
  first_axis = next(
  i for i, idx in enumerate(idxs) if isinstance(idx.type, TensorType)
@@ -211,6 +201,10 @@ def numba_funcify_multiple_integer_vector_indexing(
  )
  except StopIteration:
  after_last_axis = len(idxs)
+ last_axis = after_last_axis - 1
+
+ vector_indices = idxs[first_axis:after_last_axis]
+ assert all(v.type.broadcastable == (False,) for v in vector_indices)
 
  if isinstance(op, AdvancedSubtensor):
 
@@ -231,43 +225,58 @@ def advanced_subtensor_multiple_vector(x, *idxs):
 
  return advanced_subtensor_multiple_vector
 
- elif op.set_instead_of_inc:
+ else:
  inplace = op.inplace
 
- @numba_njit
- def advanced_set_subtensor_multiple_vector(x, y, *idxs):
- vec_idxs = idxs[first_axis:after_last_axis]
- x_shape = x.shape
+ # Check if y must be broadcasted
+ # Includes the last integer vector index,
+ indexed_bcast_dims = (
+ *x.type.broadcastable[:first_axis],
+ *x.type.broadcastable[last_axis:],
+ )
+ y_is_broadcasted = _broadcasted_to(y.type.broadcastable, indexed_bcast_dims)
 
- if inplace:
- out = x
- else:
- out = x.copy()
+ if op.set_instead_of_inc:
 
- for outer in np.ndindex(x_shape[:first_axis]):
-  for i, scalar_idxs in enumerate(zip(*vec_idxs)): # noqa: B905
-  out[(*outer, *scalar_idxs)] = y[(*outer, i)]
- return out
+ @numba_njit
+ def advanced_set_subtensor_multiple_vector(x, y, *idxs):
+ vec_idxs = idxs[first_axis:after_last_axis]
+  x_shape = x.shape
 
- return advanced_set_subtensor_multiple_vector
+ if inplace:
+ out = x
+ else:
+ out = x.copy()
 
- else:
- inplace = op.inplace
+  if y_is_broadcasted:
+  y = np.broadcast_to(y, x_shape[:first_axis] + x_shape[last_axis:])
 
- @numba_njit
- def advanced_inc_subtensor_multiple_vector(x, y, *idxs):
- vec_idxs = idxs[first_axis:after_last_axis]
- x_shape = x.shape
+ for outer in np.ndindex(x_shape[:first_axis]):
+ for i, scalar_idxs in enumerate(zip(*vec_idxs)): # noqa: B905
+ out[(*outer, *scalar_idxs)] = y[(*outer, i)]
+ return out
+
+ return advanced_set_subtensor_multiple_vector
+
+ else:
+
+ @numba_njit
+ def advanced_inc_subtensor_multiple_vector(x, y, *idxs):
+ vec_idxs = idxs[first_axis:after_last_axis]
+ x_shape = x.shape
+
+ if inplace:
+ out = x
+ else:
+ out = x.copy()
 
- if inplace:
- out = x
- else:
- out = x.copy()
+ if y_is_broadcasted:
+ y = np.broadcast_to(y, x_shape[:first_axis] + x_shape[last_axis:])
 
- for outer in np.ndindex(x_shape[:first_axis]):
- for i, scalar_idxs in enumerate(zip(*vec_idxs)): # noqa: B905
- out[(*outer, *scalar_idxs)] += y[(*outer, i)]
- return out
+  for outer in np.ndindex(x_shape[:first_axis]):
+  for i, scalar_idxs in enumerate(zip(*vec_idxs)): # noqa: B905
+  out[(*outer, *scalar_idxs)] += y[(*outer, i)]
+  return out
 
  return advanced_inc_subtensor_multiple_vector
 

tests/link/numba/test_subtensor.py

@@ -392,8 +392,8 @@ def test_AdvancedIncSubtensor1(x, y, indices):
  np.array(-99), # Broadcasted value
  ([1, 2], [2, 3]), # 2 vector indices
  False,
- True,
- True,
+ False,
+ False,
  ),
  (
  np.arange(3 * 4 * 5).reshape((3, 4, 5)),