Skip to content

Replace AbstractArray lazy types with LinearBroadcasted and Mul#156

Merged
mtfishman merged 29 commits intomainfrom
mf/linear-broadcasted
Mar 26, 2026
Merged

Replace AbstractArray lazy types with LinearBroadcasted and Mul#156
mtfishman merged 29 commits intomainfrom
mf/linear-broadcasted

Conversation

@mtfishman
Copy link
Copy Markdown
Member

@mtfishman mtfishman commented Mar 25, 2026

Summary

Replace the old AbstractArray lazy types and macro system with a cleaner, non-AbstractArray design that avoids method ambiguities with ArrayLayouts/BlockArrays.

What changed

  • Replace ScaledArray, ConjArray, AddArray, MulArray (<: AbstractArray) with LinearBroadcasted subtypes (ScaledBroadcasted, ConjBroadcasted, AddBroadcasted) and standalone Mul — none subtype AbstractArray
  • Replace Unicode operators (*ₗ, +ₗ, etc.) and macro system with linearbroadcasted(f, args...) constructor that dispatches on f to create the appropriate subtype
  • Delete LinearBroadcastedStyle and all BC.broadcasted overloads — broadcast integration moves to instantiation-time conversion via tryflattenlinear(bc::Broadcasted), which downstream styles call from Base.copy
  • All LinearBroadcasted materialization funnels through permutedimsopadd!(dest, op, src, perm, α, β) as the single overload point for downstream types
  • Op composition with simplification rules (conj ∘ conj = identity, etc.) modeled on StridedViews
  • Delete entire macro system and old lazy array types (~600 lines removed, net -564 lines)
  • Breaking: bump to v0.8.0

Broadcast integration design

Three functions for downstream opt-in:

  • islinearbroadcast(f, args...) — per-node trait
  • linearbroadcasted(f, args...) — constructor with algebraic simplifications
  • tryflattenlinear(bc::Broadcasted) — recursive conversion, returns nothing on failure

Downstream styles choose strict (error on nonlinear) or opportunistic (fall back to standard broadcasting).

Materialization chain

copy(lb) → copyto!(similar(lb), lb) → add!(dest, lb, true, false) → permutedimsopadd!(dest, identity, src, perm, α, β)
copy(m::Mul) → copyto!(similar(m), m) → mul!(dest, factors...)

permutedimsopadd! dispatches on LinearBroadcasted subtypes to unwrap scaling, conjugation, and addition before hitting the base case.

Remaining work

  • Downstream companion PRs: NamedDimsArrays, GradedArrays (delete per-package lazy types, opt in via tryflattenlinear)
  • OpBroadcasted{F} for general element-wise ops (e.g. Float64.(a))
  • Multiplicative-side op primitives (contractopadd!)
  • TensorOperations extension update for permutedimsopadd! with op
  • Styles in LinearBroadcasted/Mul types (deferred)
  • PermutedDimsArray unwrapping in permutedimsopadd! (TBD — may not be needed)

Test plan

  • CI passes
  • Downstream package compatibility (GradedArrays, NamedDimsArrays)

🤖 Generated with Claude Code

Redesign the lazy type system to avoid AbstractArray subtyping, which
caused method ambiguities with ArrayLayouts and BlockArrays.

- Add LinearBroadcasted abstract type with ScaledBroadcasted,
  ConjBroadcasted, AddBroadcasted concrete subtypes
- Add standalone Mul type for lazy matrix multiplication
- Add LinearBroadcastFunction constructor API (replaces Unicode operators)
- Rename LazyArrayStyle to LinearBroadcastedStyle
- Materialization follows Broadcasted protocol: copy → copyto! → add!/mul!
- Delete macro system and old AbstractArray lazy types
- Rename lazyarrays.jl to linearbroadcasted.jl
- Bump version to 0.8.0 (breaking)

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
@github-actions
Copy link
Copy Markdown
Contributor

github-actions bot commented Mar 25, 2026

Your PR no longer requires formatting changes. Thank you for your contribution!

@codecov
Copy link
Copy Markdown

codecov bot commented Mar 25, 2026

Codecov Report

❌ Patch coverage is 72.61905% with 46 lines in your changes missing coverage. Please review.
✅ Project coverage is 80.32%. Comparing base (9a88120) to head (6d7ab0c).
⚠️ Report is 1 commits behind head on main.

Files with missing lines Patch % Lines
src/linearbroadcasted.jl 69.73% 46 Missing ⚠️
Additional details and impacted files
@@            Coverage Diff             @@
##             main     #156      +/-   ##
==========================================
+ Coverage   72.57%   80.32%   +7.74%     
==========================================
  Files          24       24              
  Lines        1258      935     -323     
==========================================
- Hits          913      751     -162     
+ Misses        345      184     -161     
Flag Coverage Δ
docs 0.00% <0.00%> (ø)

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

mtfishman and others added 22 commits March 25, 2026 00:02
These were needed for AbstractArray{T,N} subtyping and are harder to
define for ITensorBase where N isn't well-defined. eltype/ndims are
now computed from the wrapped data instead.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
- Add generic similar(::LinearBroadcasted) → similar(a, eltype(a)) and
  similar(::LinearBroadcasted, elt) → similar(a, elt, axes(a)) fallbacks
- Each subtype now only defines the 3-arg similar(a, elt, ax)
- Remove redundant true/false defaults from copyto!(dest, ::Mul)

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Remove Val(:conj) hack. Just pass conj(unconj(src)) to add!, which
uses Base's lazy conj wrapper on the underlying array.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
- copyto!(dest, src::LinearBroadcasted) generic instead of per-subtype
- show(io, a::LinearBroadcasted) generic via operation/arguments
- iscall(::LinearBroadcasted) generic instead of per-subtype

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Use LBF (uppercase since it aliases a type) throughout source and tests
to reduce verbosity of LinearBroadcastFunction call sites.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Replace `::LinearBroadcastFunction{typeof(f)}` with `::typeof(LBF(f))`
in method signatures for consistency with the LBF shorthand. Fix1/Fix2
patterns remain as LinearBroadcastFunction{<:...} since those can't be
expressed via typeof(LBF(...)).

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
LinearFunc is less terse than LBF while still shorter than
LinearBroadcastFunction. Also use LinearFunc{<:Base.Fix1{typeof(*)}}
etc. in Fix1/Fix2 dispatch signatures for consistency.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
No shorthand — just use the full name consistently, matching the
Base.BroadcastFunction convention. Keep typeof(LinearBroadcastFunction(f))
in dispatch signatures.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Remove the LinearBroadcastFunction callable struct. Instead, use
linearbroadcasted(f, args...) as a generic constructor that dispatches
on f to create the appropriate LinearBroadcasted subtype. This mirrors
the Base.Broadcast.broadcasted(f, args...) pattern.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Mark as non-breaking (0.7.x) so downstream packages (NamedDimsArrays,
GradedArrays) can Pkg.develop this branch without compat clashes.
Will bump to 0.8.0 before merging.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Remove adjoint, transpose, permuteddims, conj, StridedViews methods
from LinearBroadcasted types and Mul. These were algebraic rewrite
rules carried over from the AbstractArray era. Will bring back as
needed when validating against NamedDimsArrays and GradedArrays PRs.

Each type now only defines: axes, eltype, ndims, similar, operation,
arguments. Plus the materialization chain (copy, copyto!, add!).

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
permutedimsopadd!(dest, op, src, perm, α, β) is the op-parameterized
materialization primitive. Default eagerly applies op. ConjBroadcasted
now materializes through permutedimsopadd!(dest, conj, ...) giving
downstream types a dispatch point to fuse conjugation.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
All LinearBroadcasted materialization now funnels through
permutedimsopadd!(dest, op, src, perm, α, β). add! and
permutedimsadd! are convenience functions that call it with
identity op and/or identity perm.

This gives downstream array types (GradedArrays, etc.) a single
function to implement for full LinearBroadcasted support.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Match Base.Broadcast's pattern of deferring axes checks to
materialization time. The check is redundant with combine_axes
in axes(::AddBroadcasted).

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
linearbroadcasted(f, args...) is the constructor (parallels broadcasted).
to_linear(bc::Broadcasted) is the internal rewriter (parallels flatten).
broadcasted_linear(style, f, args...) validates + converts.

Mixing construction and rewriting in one function was confusing.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Single recursive function validates linearity and converts in one pass,
returning nothing on failure. Eliminates redundant tree traversal.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Move to instantiation-time conversion design. Downstream styles
call tryflattenlinear(bc) from their own copy method instead of
routing through LinearBroadcastedStyle.

- Delete LinearBroadcastedStyle and all constructors
- Delete all BC.broadcasted(::LinearBroadcastedStyle, ...) overloads
- Delete to_broadcasted, broadcasted_linear
- Rename broadcast_is_linear → islinearbroadcast
- Rename _to_linear → tryflattenlinear
- BroadcastStyle for LinearBroadcasted types delegates to wrapped type

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Single implementation of permutedimsopadd! with the broadcasting logic
inlined. Fuse op into the broadcast expression (op.(src′) inside .=)
to avoid intermediate allocation. Branch on β first (NaN avoidance),
then op === identity (skip no-op).

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This allows add! to accept LinearBroadcasted types, which is needed
for copyto!(dest, ::LinearBroadcasted) to work through the add! →
permutedimsopadd! chain.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
mtfishman and others added 4 commits March 25, 2026 22:48
Add runtime branch for ndims == 0 to avoid wrapping in PermutedDimsArray,
which doesn't support getindex() on 0-dimensional BlockSparseArray.
Needed because GradedArray is a type alias and 0-dimensional contraction
results don't match the alias.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Add _to_broadcasted to convert LinearBroadcasted tree back to Broadcasted
(inverse of tryflattenlinear). Uses BC.Broadcasted constructor directly
to avoid style-based dispatch that could re-enter LinearBroadcasted.

Replace per-subtype similar methods with a single generic method that
delegates to similar(_to_broadcasted(a), elt, ax).

Also handle 0-dimensional arrays in permutedimsopadd! and widen add!
signature to accept any src type.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Use generic implementation via operation/arguments instead of
per-subtype _to_broadcasted methods. Remove per-subtype similar
methods in favor of a single generic one.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
…oadcasted, similar(AddBroadcasted), and 0-dim permutedimsopadd!

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
@mtfishman mtfishman changed the title [WIP] Replace AbstractArray lazy types with LinearBroadcasted and Mul Replace AbstractArray lazy types with LinearBroadcasted and Mul Mar 26, 2026
@mtfishman mtfishman marked this pull request as ready for review March 26, 2026 17:19
@mtfishman mtfishman merged commit 4f0cb64 into main Mar 26, 2026
40 of 42 checks passed
@mtfishman mtfishman deleted the mf/linear-broadcasted branch March 26, 2026 18:19
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