Skip to content

WIP: Overlay for Enzyme.jacobian#2839

Draft
vimarsh6739 wants to merge 2 commits intomainfrom
vim/jacobian
Draft

WIP: Overlay for Enzyme.jacobian#2839
vimarsh6739 wants to merge 2 commits intomainfrom
vim/jacobian

Conversation

@vimarsh6739
Copy link
Copy Markdown
Member

No description provided.

@vimarsh6739 vimarsh6739 changed the title WIP: Overlay for Enzyme.jacobian and Enzyme.gradient WIP: Overlay for Enzyme.jacobian Apr 24, 2026
Comment thread src/Enzyme.jl
Comment on lines +840 to +843
x isa Annotation &&
error(
"Reactant jacobian overlay v1 expects raw values or Const-wrapped values for non-differentiated arguments.",
)
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

[JuliaFormatter] reported by reviewdog 🐶

Suggested change
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.",
)

Comment thread src/Enzyme.jl
Comment on lines +859 to +862
idx == 0 &&
error(
"Reactant jacobian overlay v1 requires exactly one differentiable argument; all arguments were Const.",
)
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

[JuliaFormatter] reported by reviewdog 🐶

Suggested change
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.",
)

Comment thread src/Enzyme.jl

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}()`.")
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

[JuliaFormatter] reported by reviewdog 🐶

Suggested change
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}()`.")

Comment thread src/Enzyme.jl
Comment on lines +877 to +879
error(
"Reactant jacobian overlay v1 requires explicit `n_outs` in reverse mode.",
)
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

[JuliaFormatter] reported by reviewdog 🐶

Suggested change
error(
"Reactant jacobian overlay v1 requires explicit `n_outs` in reverse mode.",
)
error("Reactant jacobian overlay v1 requires explicit `n_outs` in reverse mode.")

Comment thread src/Enzyme.jl
Comment on lines +886 to +887
(d isa Integer && d >= 0) ||
error("Reactant jacobian overlay v1 requires non-negative integer `n_outs` dimensions.")
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

[JuliaFormatter] reported by reviewdog 🐶

Suggested change
(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.",
)

Comment thread src/Enzyme.jl
Comment on lines +946 to +953
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.",
)
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

[JuliaFormatter] reported by reviewdog 🐶

Suggested change
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.",
)

Comment thread src/Enzyme.jl
Comment on lines +1021 to +1024
n_out_elems == 1 ||
error(
"Reactant jacobian overlay v1 currently supports only scalar-output functions in reverse mode (`prod(n_outs) == 1`).",
)
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

[JuliaFormatter] reported by reviewdog 🐶

Suggested change
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`).",
)

Comment thread src/Enzyme.jl
Comment on lines +1027 to +1030
(c === nothing || c == 1) ||
error(
"Reactant jacobian overlay v1 reverse mode currently supports only `chunk=nothing` or `chunk=Val(1)`.",
)
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

[JuliaFormatter] reported by reviewdog 🐶

Suggested change
(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)`.",
)

Comment thread src/Enzyme.jl
Comment on lines +1043 to +1046
active_val isa Union{AbstractArray,AbstractFloat} ||
error(
"Reactant jacobian overlay v1 currently supports differentiating only `AbstractArray` or `AbstractFloat` arguments.",
)
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

[JuliaFormatter] reported by reviewdog 🐶

Suggested change
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.",
)

Comment thread test/core/autodiff.jl
Comment on lines +125 to +126
jacobian_forward_with_shadows(x) =
Enzyme.jacobian(Forward, jacobian_loss, x; shadows=(nothing,))
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

[JuliaFormatter] reported by reviewdog 🐶

Suggested change
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

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant