Repository navigation
Support Adapt.adapt and copy for GradientOp - #48
Merged
nHackel merged 4 commits intoNov 3, 2026
Merged
Conversation
GradientOp now returns a GradientOpImpl wrapper that keeps the shape, the directions taken from `dims` at construction, and the adjoint scratch vector. Adapt.adapt and copy rebuild the operator, its closures and scratch storage for the target storage type instead of sharing the source's state. `dims` is iterated once, so a later mutation of the caller's vector or a consumed iterator does not change a rebuilt operator. The GPU gradient kernels drop `cpu = false` so they also run on KernelAbstractions' CPU backend (JLArray). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
… type
Rebuilding keeps the operator's element type, so a target such as
Array{Float32} for a Float64 operator would mix precisions in the scratch
and buffers. Throw an ArgumentError instead and document the rule.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Compare the adapted storage with eltype(op) rather than the source storage, so an operator built with an S of another element type is rejected for a target without an element type instead of returning an operator whose storage disagrees with its element type. A target naming the operator's element type still rebuilds consistent storage. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
A rebuild with storage of the operator's element type would drop the rounding through the source's lower-precision temporary vector and change the adjoint, so such operators cannot be adapted to any target. copy keeps them unchanged. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #48 +/- ##
==========================================
+ Coverage 73.01% 81.23% +8.22%
==========================================
Files 15 15
Lines 593 613 +20
==========================================
+ Hits 433 498 +65
+ Misses 160 115 -45 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Adapt.adapt(target, op)andcopy(op)now work forGradientOp, also insideProdOp,normalOperatorandDiagOp. Both return a new operator that is rebuilt from the shape and the directions taken at construction, with a newLinearOperator, closures and temporary vector for the target storage type;opis not modified.GradientOpnow returns aGradientOpImpl <: GradientOp{T}that wraps theLinearOperatorand keepsshape,dimsand the adjoint's temporary vector, like the other wrapper operators. Properties of theLinearOperatorare forwarded.dimsis iterated once at construction and stored as a tuple, or as anIntfor a single direction. Rebuilding does not depend on the caller's container: a vector changed later, or a consumed iterator, has no effect.CuArray{Float32}for aFloat64operator throws anArgumentError, and so does an operator built with anSof another element type, whose adjoint rounds through its temporary vector.copykeeps such an operator as it is.cpu = false, so they also run on CPU array types such asJLArray. Their computation is unchanged.test/testAdapt.jlmovesGradientOpbetweenArrayandJLArray(or the configured GPU array type) forFloat32,Float64,ComplexF32,ComplexF64andInt32, and for single, several, repeated and no directions. It compares forward, transpose and adjoint, with 3- and 5-argumentmul!, against the source. It also checks that the source's buffers, temporary vector and counters are unchanged, and that copies get their own temporary vector. It further covers a mutateddimsvector, a consumed iterator,createLinearOperator,ProdOp,normalOperator,DiagOp, and the element-type errors.The tests pass on the CPU with Julia 1.10.12. With two threads, the only failure is the existing
DiagOp((op, op))FFT test from #44.