Support duplicate indices in AdvancedIncSubtensor1's Numba implementation

aesara/link/numba/dispatch/basic.py

@@ -499,7 +499,6 @@ def numba_funcify_Subtensor(op, node, **kwargs):
 
 @numba_funcify.register(IncSubtensor)
 @numba_funcify.register(AdvancedIncSubtensor)
-@numba_funcify.register(AdvancedIncSubtensor1)
 def numba_funcify_IncSubtensor(op, node, **kwargs):
 
     incsubtensor_def_src = create_index_func(
@@ -515,6 +514,39 @@ def numba_funcify_IncSubtensor(op, node, **kwargs):
     return numba_njit(incsubtensor_fn)
 
 
+@numba_funcify.register(AdvancedIncSubtensor1)
+def numba_funcify_AdvancedIncSubtensor1(op, node, **kwargs):
+    inplace = op.inplace
+    set_instead_of_inc = op.set_instead_of_inc
+
+    if set_instead_of_inc:
+
+        @numba_njit
+        def advancedincsubtensor1_inplace(x, vals, idxs):
+            for idx, val in zip(idxs, vals):
+                x[idx] = val
+            return x
+
+    else:
+
+        @numba_njit
+        def advancedincsubtensor1_inplace(x, vals, idxs):
+            for idx, val in zip(idxs, vals):
+                x[idx] += val
+            return x
+
+    if inplace:
+        return advancedincsubtensor1_inplace
+    else:
+
+        @numba_njit
+        def advancedincsubtensor1(x, vals, idxs):
+            x = x.copy()
+            return advancedincsubtensor1_inplace(x, vals, idxs)
+
+        return advancedincsubtensor1
+
+
 @numba_funcify.register(DeepCopyOp)
 def numba_funcify_DeepCopyOp(op, node, **kwargs):
 

tests/link/test_numba.py

@@ -550,6 +550,11 @@ def test_IncSubtensor(x, y, indices):
             at.as_tensor(rng.poisson(size=(2, 4, 5))),
             ([1, 2],),
         ),
+        (
+            at.as_tensor(np.arange(3 * 4 * 5).reshape((3, 4, 5))),
+            at.as_tensor(rng.poisson(size=(2, 4, 5))),
+            ([1, 1],),
+        ),
     ],
 )
 def test_AdvancedIncSubtensor1(x, y, indices):
@@ -583,6 +588,14 @@ def test_AdvancedIncSubtensor1(x, y, indices):
             at.as_tensor(rng.poisson(size=(2, 4))),
             ([1, 2], slice(None), [3, 4]),
         ),
+        pytest.param(
+            at.as_tensor(np.arange(3 * 4 * 5).reshape((3, 4, 5))),
+            at.as_tensor(rng.poisson(size=(2, 5))),
+            ([1, 1], [2, 2]),
+            marks=pytest.mark.xfail(
+                reason="Duplicate index handling hasn't been implemented, yet."
+            ),
+        ),
     ],
 )
 def test_AdvancedIncSubtensor(x, y, indices):