Skip to content

Handle UnionAll Broadcasted in broadcast override inference#3267

Draft
vchuravy wants to merge 1 commit into
mainfrom
vc/fix-broadcast-unionall
Draft

Handle UnionAll Broadcasted in broadcast override inference#3267
vchuravy wants to merge 1 commit into
mainfrom
vc/fix-broadcast-unionall

Conversation

@vchuravy

Copy link
Copy Markdown
Member

Problem

Inference crashes with a FieldError when Enzyme's broadcast mapreduce/foldl override is consulted for a Broadcasted type that widenconst left only partially specified (a UnionAll with free Axes/F/Args parameters):

FieldError: type UnionAll has no field `parameters`, available fields: `var`, `body`
  [2] bc_or_array_or_number_ty(Ty::Type, midnothing::Bool)
      @ Enzyme.Compiler.Interpreter src/compiler/interpreter.jl:1077

bc_or_array_or_number_ty and has_array match Ty <: Broadcast.Broadcasted{...} and then directly read Ty.parameters[4] (the broadcast args tuple type). That access assumes Ty is a concrete DataType, but a UnionAll has no parameters field. This was hit in the wild by the SVD reverse rule in TensorKit, which produces such a partially-specified Broadcasted type.

Fix

Guard both helpers so they bail out (return false) when Ty is not a fully-specified DataType. Returning false is the conservative, correct choice: when the args tuple type is unknown, Enzyme simply skips the broadcast override and falls back to the normal AD path instead of crashing.

Testing

BC = Base.Broadcast.Broadcasted
DAS = Base.Broadcast.DefaultArrayStyle
T1 = BC{DAS{2}, A, F, Args} where {A, F, Args}   # UnionAll, previously crashed
bc_or_array_or_number_ty(T1, false)  # => false (was: FieldError)
has_array(T1, false)                 # => false (was: FieldError)
# concrete Broadcasted still returns true as before

🤖 Generated with Claude Code

`widenconst` can yield a `UnionAll` `Broadcasted` type (free Axes/F/Args
parameters), which has no `.parameters` field, causing a `FieldError`
during inference of the broadcast mapreduce/foldl override. Guard both
helpers so they conservatively skip the override when the type is not a
fully-specified `DataType`.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant