Fix numba failure on gradient of chained batched Convolve1d - #2375
Fix numba failure on gradient of chained batched Convolve1d#2375juanitorduz wants to merge 2 commits into
Conversation
|
The 8 failing CI jobs here are pre-existing on Fix proposed in #2376. Once that merges I'll update this branch with |
|
I think failing test are coming from #2376 |
e587dcb to
9b32c97
Compare
…form _vectorize_node_perform wrapped every core input value in np.asarray, storing 0d ndarrays. For a core input with ScalarType (e.g. the boolean full_mode flag of Convolve1d, signature "()"), the core C thunk's c_extract rejects a 0d array with "ValueError: Scalar check failed". The bug was latent for every linker but only surfaced when the lazily built gufunc (impl=None) picked the C thunk for the core op, as happens in the numba object-mode fallback which calls perform on an unprepared node. Storage now honors the type's own contract (ScalarType.filter returns numpy scalars), with per-input converters precomputed outside the loop.
… shapes introduce_explicit_core_shape_blockwise bailed whenever any Blockwise appeared in applys_between(node.inputs, core_shapes). Because that traversal also yields the owners of the blocker variables, a Blockwise whose core shape reads Shape_i of an input produced by another Blockwise was permanently excluded from BlockwiseWithCoreShape wrapping and fell back to object mode. This broke the gradient of chained convolve1d calls with a kernel batched by vectorize_graph (issue pymc-devs#2360). Per ShapeFeature.get_non_recursive_shape's contract, the shape expressions read only the node's own inputs; the only case that cannot be introduced is the Shape_i(output) fallback of Blockwise.infer_shape, where the core shape requires evaluating the node itself. Bail only on that self-reference. Closes pymc-devs#2360
9b32c97 to
d1693aa
Compare
|
I'm less sure on this one, want @ricardoV94 to take a look. My gut says this is addressing the symptom and not the cause. Details about different linkers shouldn't be leaking into the implementation of Blockwise. If one linker is doing something that doesn't conform to the others, we should fix it. |
|
Doesn't Blockwise explicitly require tensor inputs? Numba has a thing for downcasting 0d arrays to python scalars, it's a known limitation. I think the bot missed this as the source of the failure. If that's the case, I'd stick only with the core shape fix (which I haven't reviewed yet). We can't be defensive about numba scalar downcasting everywhere. Fortunately it only hurts at the boundaries (obj mode or function output) |
| for node in applys_between(node.inputs, core_shapes) | ||
| ): | ||
| # If Blockwise shows up in the shape graph we can't introduce the core shape | ||
| if set(node.outputs) & set(ancestors(core_shapes, blockers=node.inputs)): |
There was a problem hiding this comment.
This check is too soft. If a core shape rebuilds the same Blockwise but fresh or a new one you'd accept it.
There was a problem hiding this comment.
If the issue was indeed "Because that traversal also yields the owners of the blocker variables" just patch that. It's a boundary precision question.
If it's not just that the explanation is still lacking
3 participants
Description
Taking the gradient of chained
convolve1dcalls, where the kernel has been batched withvectorize_graphand the signal left unbatched (the adstock pattern used in pymc-marketing), failed under the numba backend: theBlockwise{Convolve1d}fell back to object mode with aUserWarning, and evaluation then raisedValueError: Scalar check failed (npy_bool). The same graph runs fine with thepy/cvm/c|py/JAX linkers.Two independent bugs were involved, fixed in one commit each:
1. Over-broad bail-out in
introduce_explicit_core_shape_blockwise(pytensor/tensor/rewriting/numba.py)The rewrite skipped any node for which a
Blockwiseappeared inapplys_between(node.inputs, core_shapes). Because that traversal also yields the owners of the blocker variables, a Blockwise whose core shape readsShape_iof an input produced by another Blockwise was permanently excluded fromBlockwiseWithCoreShapewrapping — which is exactly the shape of the chained-convolution gradient (each grad conv consumes the previous conv's output of unknown static length). PerShapeFeature.get_non_recursive_shape's contract, the shape expressions read only the node's own inputs; the only case that genuinely cannot be introduced is theShape_i(output)fallback ofBlockwise.infer_shape, where computing the core shape would require evaluating the node itself. The guard now bails only on that self-reference. With this change the issue's repro compiles without any object-mode fallback, for both static and dynamic signal shapes, and the gradient matches thecvmresult.2. 0d arrays stored into
ScalarTypecore-input storage inBlockwise.perform(pytensor/tensor/blockwise.py)_vectorize_node_performwrapped every core input value innp.asarray(...). For a core input withScalarType— the booleanfull_modeflag ofConvolve1d, signature()— the core C thunk'sc_extractrejects a 0d array (Scalar check failed (npy_bool)). The bug is latent for every linker but only surfaces when the lazily-built gufunc (impl=None) picks the C thunk, as happens in the numba object-mode fallback which callsperformon an unprepared node (thecvm/pylinkers end up rebuilding the gufunc with the Python impl viaprepare_node, which is why they worked). Storage now honors the type's own contract (ScalarType.filterreturns numpy scalars), with per-input converters precomputed outside the hot loop — the input-side counterpart of the scalar-output handling from #1846. This keeps the object-mode fallback correct for the cases that legitimately remain there (e.g. a boolean mode that varies across batch dimensions).The fixes are deliberately backend/graph-level rather than Convolve-specific, following the direction from #1522 and #1998 (keep the mode symbolic; no conv-specific mode rewrites; static shapes are a nice-to-have recovered at dispatch time).
Notes for possible follow-ups (not addressed here):
applys_between's docstring says it excludes the owners of the input variables, but the implementation yields them (viaancestorsyielding blockers); this PR works around it locally rather than changing shared traversal behavior.Blockwise(Convolve1d)(Blockwise{Convolve1d}never infers static output shape #1998, related JAX issue JAX linker: Cannot compile Convolve1d without static mode #2317) would additionally need constant-mode extraction inAbstractConvolveNd.make_nodepluspropagate_unbatched_core_inputs=TrueinBlockwise.make_node.Tests: a numba end-to-end regression test of the issue's graph (object-mode fallback path and fully-rewritten path), unit tests for the corrected guard (including the self-referential case that must still bail), and unit tests for
Blockwise.performwith broadcast and batched booleanScalarTypecore inputs exercising the core C thunk.Related Issue
Checklist
Type of change
🤖 Generated with Claude Code