WIP: Overlay for Enzyme.jacobian#2839
Conversation
| x isa Annotation && | ||
| error( | ||
| "Reactant jacobian overlay v1 expects raw values or Const-wrapped values for non-differentiated arguments.", | ||
| ) |
There was a problem hiding this comment.
[JuliaFormatter] reported by reviewdog 🐶
| x isa Annotation && | |
| error( | |
| "Reactant jacobian overlay v1 expects raw values or Const-wrapped values for non-differentiated arguments.", | |
| ) | |
| x isa Annotation && error( | |
| "Reactant jacobian overlay v1 expects raw values or Const-wrapped values for non-differentiated arguments.", | |
| ) |
| idx == 0 && | ||
| error( | ||
| "Reactant jacobian overlay v1 requires exactly one differentiable argument; all arguments were Const.", | ||
| ) |
There was a problem hiding this comment.
[JuliaFormatter] reported by reviewdog 🐶
| idx == 0 && | |
| error( | |
| "Reactant jacobian overlay v1 requires exactly one differentiable argument; all arguments were Const.", | |
| ) | |
| idx == 0 && error( | |
| "Reactant jacobian overlay v1 requires exactly one differentiable argument; all arguments were Const.", | |
| ) |
|
|
||
| function _jacobian_parse_chunk(chunk) | ||
| chunk === nothing && return nothing | ||
| chunk isa Val || error("Reactant jacobian overlay v1 expects `chunk` to be `nothing` or `Val{N}()`.") |
There was a problem hiding this comment.
[JuliaFormatter] reported by reviewdog 🐶
| chunk isa Val || error("Reactant jacobian overlay v1 expects `chunk` to be `nothing` or `Val{N}()`.") | |
| chunk isa Val || | |
| error("Reactant jacobian overlay v1 expects `chunk` to be `nothing` or `Val{N}()`.") |
| error( | ||
| "Reactant jacobian overlay v1 requires explicit `n_outs` in reverse mode.", | ||
| ) |
There was a problem hiding this comment.
[JuliaFormatter] reported by reviewdog 🐶
| error( | |
| "Reactant jacobian overlay v1 requires explicit `n_outs` in reverse mode.", | |
| ) | |
| error("Reactant jacobian overlay v1 requires explicit `n_outs` in reverse mode.") |
| (d isa Integer && d >= 0) || | ||
| error("Reactant jacobian overlay v1 requires non-negative integer `n_outs` dimensions.") |
There was a problem hiding this comment.
[JuliaFormatter] reported by reviewdog 🐶
| (d isa Integer && d >= 0) || | |
| error("Reactant jacobian overlay v1 requires non-negative integer `n_outs` dimensions.") | |
| (d isa Integer && d >= 0) || error( | |
| "Reactant jacobian overlay v1 requires non-negative integer `n_outs` dimensions.", | |
| ) |
| isempty(kwargs) || | ||
| error( | ||
| "Reactant jacobian overlay v1 only supports `chunk` and `shadows` keywords in forward mode.", | ||
| ) | ||
| shadows === nothing || | ||
| error( | ||
| "Reactant jacobian overlay v1 does not support explicit `shadows`; omit the keyword to use internal Jacobian seeds.", | ||
| ) |
There was a problem hiding this comment.
[JuliaFormatter] reported by reviewdog 🐶
| isempty(kwargs) || | |
| error( | |
| "Reactant jacobian overlay v1 only supports `chunk` and `shadows` keywords in forward mode.", | |
| ) | |
| shadows === nothing || | |
| error( | |
| "Reactant jacobian overlay v1 does not support explicit `shadows`; omit the keyword to use internal Jacobian seeds.", | |
| ) | |
| isempty(kwargs) || error( | |
| "Reactant jacobian overlay v1 only supports `chunk` and `shadows` keywords in forward mode.", | |
| ) | |
| shadows === nothing || error( | |
| "Reactant jacobian overlay v1 does not support explicit `shadows`; omit the keyword to use internal Jacobian seeds.", | |
| ) |
| n_out_elems == 1 || | ||
| error( | ||
| "Reactant jacobian overlay v1 currently supports only scalar-output functions in reverse mode (`prod(n_outs) == 1`).", | ||
| ) |
There was a problem hiding this comment.
[JuliaFormatter] reported by reviewdog 🐶
| n_out_elems == 1 || | |
| error( | |
| "Reactant jacobian overlay v1 currently supports only scalar-output functions in reverse mode (`prod(n_outs) == 1`).", | |
| ) | |
| n_out_elems == 1 || error( | |
| "Reactant jacobian overlay v1 currently supports only scalar-output functions in reverse mode (`prod(n_outs) == 1`).", | |
| ) |
| (c === nothing || c == 1) || | ||
| error( | ||
| "Reactant jacobian overlay v1 reverse mode currently supports only `chunk=nothing` or `chunk=Val(1)`.", | ||
| ) |
There was a problem hiding this comment.
[JuliaFormatter] reported by reviewdog 🐶
| (c === nothing || c == 1) || | |
| error( | |
| "Reactant jacobian overlay v1 reverse mode currently supports only `chunk=nothing` or `chunk=Val(1)`.", | |
| ) | |
| (c === nothing || c == 1) || error( | |
| "Reactant jacobian overlay v1 reverse mode currently supports only `chunk=nothing` or `chunk=Val(1)`.", | |
| ) |
| active_val isa Union{AbstractArray,AbstractFloat} || | ||
| error( | ||
| "Reactant jacobian overlay v1 currently supports differentiating only `AbstractArray` or `AbstractFloat` arguments.", | ||
| ) |
There was a problem hiding this comment.
[JuliaFormatter] reported by reviewdog 🐶
| active_val isa Union{AbstractArray,AbstractFloat} || | |
| error( | |
| "Reactant jacobian overlay v1 currently supports differentiating only `AbstractArray` or `AbstractFloat` arguments.", | |
| ) | |
| active_val isa Union{AbstractArray,AbstractFloat} || error( | |
| "Reactant jacobian overlay v1 currently supports differentiating only `AbstractArray` or `AbstractFloat` arguments.", | |
| ) |
| jacobian_forward_with_shadows(x) = | ||
| Enzyme.jacobian(Forward, jacobian_loss, x; shadows=(nothing,)) |
There was a problem hiding this comment.
[JuliaFormatter] reported by reviewdog 🐶
| jacobian_forward_with_shadows(x) = | |
| Enzyme.jacobian(Forward, jacobian_loss, x; shadows=(nothing,)) | |
| function jacobian_forward_with_shadows(x) | |
| return Enzyme.jacobian(Forward, jacobian_loss, x; shadows=(nothing,)) | |
| end |
No description provided.