bool_idx_to_nonzero rewrite...

Author: ricardoV94

pytensor/link/numba/dispatch/subtensor.py

@@ -30,7 +30,7 @@
     IncSubtensor,
     Subtensor,
 )
-from pytensor.tensor.type_other import MakeSlice
+from pytensor.tensor.type_other import MakeSlice, NoneTypeT
 
 
 def slice_new(self, start, stop, step):
@@ -270,6 +270,8 @@ def numba_funcify_AdvancedSubtensor(op, node, **kwargs):
         not must_ignore_duplicates
         and len(adv_idxs) >= 1
         and all(adv_idx["dtype"] != "bool" for adv_idx in adv_idxs)
+        # Implementation does not support newaxis
+        and not any(isinstance(idx.type, NoneTypeT) for idx in idxs)
     ):
         return vector_integer_advanced_indexing(op, node, **kwargs)
 
@@ -310,6 +312,8 @@ def vector_integer_advanced_indexing(
 ):
     """Implement all forms of advanced indexing (and assignment) that combine basic and vector integer indices.
 
+        It does not support `newaxis` in basic indices
+
         It handles += like `np.add.at` would, accumulating add for duplicate indices.
 
     Examples

pytensor/tensor/rewriting/subtensor.py

@@ -1782,6 +1782,7 @@ def bool_idx_to_nonzero(fgraph, node):
     bool_idx_to_nonzero.__name__,
     bool_idx_to_nonzero,
     "numba",
+    "shape_unsafe",  # It can mask invalid mask sizes
     use_db_name_as_tag=False,  # Not included if only "specialize" is requested
 )