Skip to content

Commit

Permalink
Merge pull request #1462 from FluxML/bc/rm-boundscheck
Browse files Browse the repository at this point in the history
Handle `Expr(:boundscheck)`
  • Loading branch information
ToucheSir authored Jan 8, 2024
2 parents 9df7226 + bc78d43 commit 392a1f9
Show file tree
Hide file tree
Showing 2 changed files with 13 additions and 0 deletions.
4 changes: 4 additions & 0 deletions src/compiler/reverse.jl
Original file line number Diff line number Diff line change
Expand Up @@ -132,6 +132,10 @@ function instrument(ir::IR)
elseif isexpr(ex, :(=))
@assert ex.args[1] isa GlobalRef
pr[v] = xcall(Zygote, :global_set, QuoteNode(ex.args[1]), ex.args[2])
elseif isexpr(ex, :boundscheck)
# Expr(:boundscheck) now appears in common Julia code paths, so we need to handle it.
# For correctness sake, fix to true like https://github.com/dfdx/Umlaut.jl/issues/34.
pr[v] = true
else
ex = instrument_new!(pr, v, ex)
ex = instrument_literals!(pr, v, ex)
Expand Down
9 changes: 9 additions & 0 deletions test/compiler.jl
Original file line number Diff line number Diff line change
Expand Up @@ -226,6 +226,15 @@ end
# issue 897
@test gradient(x -> sum(norm, collect(eachcol(x))), ones(3, 400))[1] fill(0.5773502691896258, 3, 400)

# Tests adapted from https://github.com/dfdx/Umlaut.jl/pull/35
@eval _has_boundscheck(x) = ifelse($(Expr(:boundscheck)), 2x, x)

@testset "Meta Expr handling" begin
y, (dx,) = withgradient(_has_boundscheck, 1)
@test y == 2
@test dx == 2
end

# issue 1118 & 1380
function f_1380(x)
if rand(Bool)
Expand Down

0 comments on commit 392a1f9

Please sign in to comment.