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

1 participant