Skip to content

Commit ea0da01

Browse files
committed
Attempting to fix broadcast
1 parent 7868448 commit ea0da01

File tree

1 file changed

+12
-4
lines changed

1 file changed

+12
-4
lines changed

src/broadcast.jl

Lines changed: 12 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -414,11 +414,19 @@ function Base.copyto!(
414414
A, C = C, A
415415
end
416416
if !isnothing(A) && !isnothing(C) && !isnothing(α) && !isnothing(β)
417-
## TODO this line fails, because the function (r, t) -> β * r + α * t
418-
## is an anonymous ITensor function which the GPU compiler cannot process
419-
## This code is being called in ITensorsVectorInterfaceExt but checked was missed because
420417

421-
map!((r, t) -> β * r + α * t, T, T, A)
418+
#map!((r, t) -> β * r + α * t, T, T, A)
419+
## I tried to make a closure like this but
420+
## This throws an error still during GPU compiling.
421+
# f1 = +
422+
# f2 = *
423+
# map!((r, t) -> f1(f2(β, r), f2(α, t)), T, T, A)
424+
## The only way I could get this to work was by
425+
## Splitting it into two calls
426+
## The original call made 51 allocations on CPU and this new
427+
## version makes 44 allocations
428+
T .= β .* T
429+
T .+= α .* A
422430
else
423431
bc_bc_α = find_type(Broadcasted, bc_α.args)
424432
if isnothing(α)

0 commit comments

Comments
 (0)