Fix numba AdvancedIncSubtensor1 with broadcasted values

[ricardoV94]

Description

Related Issue

• Closes No broadcasting support in numba AdvancedIncSubtensor1 #93

Checklist

• Checked that the pre-commit linting/style checks pass
• Included tests that prove the fix is effective or that the new feature works
• Added necessary documentation (docstrings and/or example notebooks)
• If you are a pro: each commit corresponds to a relevant logical change

Type of change

• New feature / enhancement
• Bug fix
• Documentation
• Maintenance
• Other (please specify):

[codecov]

Codecov Report

Attention: Patch coverage is 89.74359% with 4 lines in your changes missing coverage. Please review.

Project coverage is 80.85%. Comparing base (1e96b89) to head (85b0c5d).
Report is 209 commits behind head on main.

Files	Patch %	Lines
pytensor/link/numba/dispatch/basic.py	89.74%	2 Missing and 2 partials ⚠️

Additional details and impacted files

@@           Coverage Diff           @@
##             main     #757   +/-   ##
=======================================
  Coverage   80.85%   80.85%           
=======================================
  Files         162      162           
  Lines       47019    47045   +26     
  Branches    11502    11514   +12     
=======================================
+ Hits        38017    38039   +22     
- Misses       6748     6753    +5     
+ Partials     2254     2253    -1     

Files	Coverage Δ
pytensor/link/numba/dispatch/basic.py	85.68% <89.74%> (-0.06%)	⬇️

... and 3 files with indirect coverage changes

[aseyboldt]

I think we might need a workaround for this issue: numba/numba#9573
If the value on the right hand side of the update is a rank zero array (ie not a scalar) that gets broadcasted, I think the current code will break. We should at least test for that I think. Maybe by manually calling funcify and because of the recent no_cpython_wrapper addition, a manual wrapper function around that?

[ricardoV94]

I think we might need a workaround for this issue: numba/numba#9573 If the value on the right hand side of the update is a rank zero array (ie not a scalar) that gets broadcasted, I think the current code will break. We should at least test for that I think. Maybe by manually calling funcify and because of the recent no_cpython_wrapper addition, a manual wrapper function around that?

The last test case I added in this PR is a tensor scalar on the rhs:

(
            pt.as_tensor(np.arange(3 * 4 * 5).reshape((3, 4, 5))),
            pt.as_tensor(rng.poisson(size=())),  # <- Scalar y
            ([2, 0],),
        ),

This is the fgraph of the last check:

AdvancedIncSubtensor1{inplace,inc} [id A] <Tensor3(int64, shape=(3, 4, 5))> 0
 ├─ <Tensor3(int64, shape=(3, 4, 5))> [id B] <Tensor3(int64, shape=(3, 4, 5))>
 ├─ <Scalar(int64, shape=())> [id C] <Scalar(int64, shape=())>
 └─ [2 0] [id D] <Vector(int64, shape=(2,))>

I don't think we allow pytensor scalars in AdvancedSetSubtensor. They always get converted to TensorVariables:

import pytensor.tensor as pt
import pytensor.scalar as ps

x = pt.vector("x")
y = ps.float64("y")
x[[0, 1]].set(y).dprint(print_type=True)

# AdvancedSetSubtensor [id A] <Vector(float64, shape=(?,))>
#  ├─ x [id B] <Vector(float64, shape=(?,))>
#  ├─ TensorFromScalar [id C] <Scalar(float64, shape=())>
#  │  └─ y [id D] <float64>
#  └─ [0 1] [id E] <Vector(int64, shape=(2,))>

[ricardoV94]

Okay finally got it, the bug is when the lhs is a vector, and the rhs a 0d array. It's fine if the lhs is a higher rank array, which was what I was testing before

[ricardoV94]

@aseyboldt I added a patch for the bug you found and a test that would fail without the patch. Let me know what you think

[ricardoV94]

I have to reassign to a new named variable, otherwise typing fails. Is that expected from Numba?

[aseyboldt]

Sometimes numba seems to require this when the type changes, but I don't know the exact rules either.

[ricardoV94]

Opened an issue, perhaps it's documented behavior, seemed strange though: numba/numba#9587

[aseyboldt]

Looks good

[ricardoV94]

This indexing is always safe because x must be at least 1D (otherwise make_node raises)