-
Notifications
You must be signed in to change notification settings - Fork 3
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Tapir.value_and_pullback!! and mutation #113
Comments
Eurgh, yeah, this looks like another interface issue. This is what I get for publishing an interface without working through a proper range of examples with it. If you tweak your above example to make use of the rule directly, something like the following: function tapir_rrule!!(f!; y, dx, x, dy, verbose=false)
rrule = build_rrule(f!, y, x)
dy_righttype = convert(tangent_type(typeof(y)), dy)
dx_righttype = convert(tangent_type(typeof(x)), dx)
dx_righttype .= zero(eltype(dx_righttype))
@show dx_righttype === dx
verbose && @info "before: primals" x y
verbose && @info "before: tangents" dx_righttype dy_righttype
# new_y, (new_df!, new_dy, new_dx) = value_and_pullback!!(
# rrule,
# NoTangent(),
# zero_codual(f!),
# CoDual(y, dy_righttype),
# CoDual(x, dx_righttype),
# )
new_y, pb!! = rrule(
zero_codual(f!),
CoDual(y, dy_righttype),
CoDual(x, dx_righttype),
)
display(dy_righttype)
display(dx_righttype)
_, new_dy, new_dx = pb!!(Tapir.tangent(new_y), NoTangent(), dy_righttype, dx_righttype)
verbose && @info "after: primals" x y
verbose && @info "after: tangents" dx_righttype dy_righttype
verbose && @info "after: Tapir output" new_y new_dx new_dy
return y, new_dx
end you'll get the following: julia> tapir_rrule!!(f!; y = zeros(2), dx = zeros(2), x = float.(1:2), dy = float.(3:4), verbose=true)
dx_righttype === dx = true
┌ Info: before: primals
│ x =
│ 2-element Vector{Float64}:
│ 1.0
│ 2.0
│ y =
│ 2-element Vector{Float64}:
│ 0.0
└ 0.0
┌ Info: before: tangents
│ dx_righttype =
│ 2-element Vector{Float64}:
│ 0.0
│ 0.0
│ dy_righttype =
│ 2-element Vector{Float64}:
│ 3.0
└ 4.0
2-element Vector{Float64}:
0.0
0.0
2-element Vector{Float64}:
0.0
0.0
┌ Info: after: primals
│ x =
│ 2-element Vector{Float64}:
│ 1.0
│ 2.0
│ y =
│ 2-element Vector{Float64}:
│ 0.0
└ 0.0
┌ Info: after: tangents
│ dx_righttype =
│ 2-element Vector{Float64}:
│ 0.0
│ 0.0
│ dy_righttype =
│ 2-element Vector{Float64}:
│ 3.0
└ 4.0
┌ Info: after: Tapir output
│ new_y = CoDual{Int64, NoTangent}(0, NoTangent())
│ new_dx =
│ 2-element Vector{Float64}:
│ 0.0
│ 0.0
│ new_dy =
│ 2-element Vector{Float64}:
│ 3.0
└ 4.0
([0.0, 0.0], [0.0, 0.0]) So. Similarly, the state of Conversely, if we do the following: function tapir_rrule!!_correct(f!; y, dx, x, dy, verbose=false)
rrule = build_rrule(f!, y, x)
yc = copy(y)
xc = copy(x)
dy_righttype = zero_tangent(y)
dx_righttype = zero_tangent(x)
verbose && @info "before: primals" x y
verbose && @info "before: tangents" dx_righttype dy_righttype
new_y, pb!! = rrule(
zero_codual(f!),
CoDual(yc, dy_righttype),
CoDual(xc, dx_righttype),
)
display(dy_righttype)
display(dx_righttype)
# Copy state into originals,
y .= yc
x .= xc
# Load tangents into AD memory.
dy_righttype .= dy
dx_righttype .= dx
_, new_dy, new_dx = pb!!(Tapir.tangent(new_y), NoTangent(), dy_righttype, dx_righttype)
verbose && @info "after: primals" x y
verbose && @info "after: tangents" dx_righttype dy_righttype
verbose && @info "after: Tapir output" new_y new_dx new_dy
return y, new_dx
end we see julia> tapir_rrule!!_correct(f!; y = zeros(2), dx = zeros(2), x = float.(1:2), dy = float.(3:4), verbose=true)
┌ Info: before: primals
│ x =
│ 2-element Vector{Float64}:
│ 1.0
│ 2.0
│ y =
│ 2-element Vector{Float64}:
│ 0.0
└ 0.0
┌ Info: before: tangents
│ dx_righttype =
│ 2-element Vector{Float64}:
│ 0.0
│ 0.0
│ dy_righttype =
│ 2-element Vector{Float64}:
│ 0.0
└ 0.0
2-element Vector{Float64}:
0.0
0.0
2-element Vector{Float64}:
0.0
0.0
┌ Info: after: primals
│ x =
│ 2-element Vector{Float64}:
│ 1.0
│ 2.0
│ y =
│ 2-element Vector{Float64}:
│ 1.0
└ 4.0
┌ Info: after: tangents
│ dx_righttype =
│ 2-element Vector{Float64}:
│ 6.0
│ 16.0
│ dy_righttype =
│ 2-element Vector{Float64}:
│ 0.0
└ 0.0
┌ Info: after: Tapir output
│ new_y = CoDual{Int64, NoTangent}(0, NoTangent())
│ new_dx =
│ 2-element Vector{Float64}:
│ ⋮
│ new_dy =
│ 2-element Vector{Float64}:
└ ⋮
([1.0, 4.0], [6.0, 16.0]) The lesson for I'l have a go at this today. Clearly there are some lots of opportunities for subtle bugs in this interface, so I really appreciate you finding some of them @gdalle ! |
Sorry I got confused in the MWE with the function output itself, let's call it |
Why is it the intended behavior for the AD to leave everything the way it was found? |
There are two things here. Primals: the reason that the primals must be left the way they were found is that you need to be able to assume that, upon entering the reverse-pass of an rrule, everything is as you left it at the end of the forwards-pass. (If you weren't to make this assumption, you'd have to make copies of the bits of all of the inputs that you need on the reverse-pass in case they're modified later on, which would often be quite computationally demanding). You can guarantee that this assumption holds by requiring each rrule to undo any changes it makes to primal state. This property holds for derived rules by induction. (Co)tangents: to see why this makes sense, note that This raises the following question: are
you should expect to see that the tangents to Does this help at all? |
It helps a little bit, some parts are still obscure to me but it's okay. My "mutating case" in DifferentiationInterface is still conceptually a function So when I run
|
Ahhh cool. Yeah, we'll probably have to write that case specifically using the |
Cool! Could you maybe write me a prototype manual |
Sure! Something along the lines of the following ought to do it: using Tapir:
build_rrule, tangent_type, set_to_zero!!, zero_codual, CoDual, tangent, increment!!,
NoTangent
function tapir_value_and_pullback!!(f!; y, dx, x, dy, verbose=false)
# Build the rule. Obviously, this should really be done separately and re-used, as
# you've previously pointed out.
rrule = build_rrule(f!, y, x)
# I'm just going to assert that the tangents are of the correct type. Possibly this is
# the best strategy generally, but maybe not?
@assert tangent_type(typeof(y)) == typeof(dy)
@assert tangent_type(typeof(x)) == typeof(dx)
# Assert that f! has no differentiable parameters.
@assert tangent_type(typeof(f!)) == NoTangent
# We want the VJP, not VJP + dx, so I'm going to zero-out `dx`. `set_to_zero!!` has the advantage
# that it will also replace any immutable components of `dx` to zero.
dx = set_to_zero!!(dx)
# Per the above discussion, we want `dy` to correspond to the cotangent of `y` _after_
# running the forwards-pass, so I'm going to take a copy, and zero-out the original.
dy_1 = copy(dy)
dy = set_to_zero!!(dy)
# Mutate a copy of `y`, so that we can run the reverse-pass later on.
y_copy = copy(y)
# Run the forwards-pass.
out, pb!! = rrule(zero_codual(f!), CoDual(y_copy, dy), CoDual(x, dx))
# Verify that the output is non-differentiable.
@assert tangent(out) == NoTangent()
# Set the cotangent of `y` to be equal to the requested value.
dy = increment!!(dy, dy_1)
# Record the state of `y` before running the reverse-pass.
y = copy!(y, y_copy)
# Run the reverse-pass.
_, _, new_dx = pb!!(NoTangent(), NoTangent(), dy, dx)
return y, new_dx
end Let me know if this helps! I think the main way that this could fall over is if |
As discovered in gdalle/DifferentiationInterface.jl#126, Tapir seems to struggle with mutating functions
z = f!(y, x)
(wherez
is ignored). I extracted the MWE from DifferentiationInterface.jl.Setup:
Results:
There might be something wrong in the way I used Tapir though.
The text was updated successfully, but these errors were encountered: