Skip to content

Commit

Permalink
Minor broadcasting fixes (#673)
Browse files Browse the repository at this point in the history
* fixes

* a test

* revert some as the broke inference

* ignore contents of at-debug macro

* version
  • Loading branch information
mcabbott authored Sep 4, 2022
1 parent 13ccc86 commit 5855c10
Show file tree
Hide file tree
Showing 4 changed files with 13 additions and 3 deletions.
2 changes: 1 addition & 1 deletion Project.toml
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
name = "ChainRules"
uuid = "082447d4-558c-5d27-93f4-14fc19e9eca2"
version = "1.44.5"
version = "1.44.6"

[deps]
Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
Expand Down
8 changes: 6 additions & 2 deletions src/rulesets/Base/broadcast.jl
Original file line number Diff line number Diff line change
Expand Up @@ -128,13 +128,14 @@ function split_bc_pullbacks(cfg::RCR, f::F, args::Vararg{Any,N}) where {F,N}
rrule_via_ad(cfg, f, a...)
end
function back_generic(dys)
deltas = unzip_broadcast(backs, unthunk(dys)) do back, dy # (could be map, sizes match)
deltas = unzip_broadcast(backs, dys) do back, dy # (could be map, sizes match)
map(unthunk, back(dy))
end
dargs = map(unbroadcast, args, Base.tail(deltas))
df = ProjectTo(f)(sum(first(deltas)))
return (NoTangent(), NoTangent(), df, dargs...)
end
back_generic(dys::AbstractThunk) = back_generic(unthunk(dys))
back_generic(z::AbstractZero) = (TRI_NO..., map(Returns(z), args)...)
return ys3, back_generic
end
Expand Down Expand Up @@ -318,7 +319,7 @@ rrule(::typeof(broadcasted), ::typeof(complex), x::Number) = rrule(complex, x) |

function unbroadcast(x::Base.AbstractArrayOrBroadcasted, dx_raw)
dx = unthunk(dx_raw)
N = ndims(dx)
N = _ndims(dx)
if length(x) == length(dx)
ProjectTo(x)(dx) # handles trivial reshapes, offsets, structured matrices, row vectors
else
Expand All @@ -328,6 +329,9 @@ function unbroadcast(x::Base.AbstractArrayOrBroadcasted, dx_raw)
end
unbroadcast(x::Base.AbstractArrayOrBroadcasted, dx::AbstractZero) = dx

_ndims(x) = ndims(x)
_ndims(::Tuple) = 1

function unbroadcast(x::T, dx_raw) where {T<:Tuple{Vararg{Any,N}}} where {N}
dx = unthunk(dx_raw)
val = if N == length(dx)
Expand Down
5 changes: 5 additions & 0 deletions src/rulesets/Base/nondiff.jl
Original file line number Diff line number Diff line change
Expand Up @@ -477,6 +477,11 @@ end
@non_differentiable Broadcast.result_style(::Any)
@non_differentiable Broadcast.result_style(::Any, ::Any)

@non_differentiable Base.CoreLogging.current_logger_for_env(::Any...)
@non_differentiable Base.CoreLogging._invoked_shouldlog(::Any...)
@non_differentiable Base.CoreLogging.Base.fixup_stdlib_path(::Any)
@non_differentiable Base.CoreLogging.handle_message(::Any...)

@non_differentiable Libc.free(::Any)
@non_differentiable Libc.getpid()
@non_differentiable Libc.strptime(::AbstractString)
Expand Down
1 change: 1 addition & 0 deletions test/rulesets/Base/broadcast.jl
Original file line number Diff line number Diff line change
Expand Up @@ -176,5 +176,6 @@ BT1 = Broadcast.BroadcastStyle(Tuple)

@testset "bugs" begin
@test ChainRules.unbroadcast((1, 2, [3]), [4, 5, [6]]) isa Tangent # earlier, NTuple demanded same type
@test ChainRules.unbroadcast(broadcasted(-, (1, 2), 3), (4, 5)) == (4, 5) # earlier, called ndims(::Tuple)
end
end

2 comments on commit 5855c10

@mcabbott
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@JuliaRegistrator
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Registration pull request created: JuliaRegistries/General/67686

After the above pull request is merged, it is recommended that a tag is created on this repository for the registered package version.

This will be done automatically if the Julia TagBot GitHub Action is installed, or can be done manually through the github interface, or via:

git tag -a v1.44.6 -m "<description of version>" 5855c10bdbe691fc07822752f5b5865b9cea44d3
git push origin v1.44.6

Please sign in to comment.