From a1e280549257c09a7105a6b1d610706c3e976f4a Mon Sep 17 00:00:00 2001 From: shukebeta Date: Tue, 3 Nov 2026 03:45:35 +1300 Subject: [PATCH 1/4] Reconstruct GradientOp with Adapt and copy 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 --- ext/LinearOperatorGPUArraysExt/GradientOp.jl | 6 +- src/GradientOp.jl | 52 ++++++- test/testAdapt.jl | 154 +++++++++++++++++++ 3 files changed, 207 insertions(+), 5 deletions(-) diff --git a/ext/LinearOperatorGPUArraysExt/GradientOp.jl b/ext/LinearOperatorGPUArraysExt/GradientOp.jl index ea8e79f..948472c 100644 --- a/ext/LinearOperatorGPUArraysExt/GradientOp.jl +++ b/ext/LinearOperatorGPUArraysExt/GradientOp.jl @@ -2,7 +2,7 @@ function LinearOperatorCollection.grad!(res::vecT, img::vecT, shape::NTuple{N,In res = reshape(res, shape .- Tuple(di)) backend = get_backend(res) - @kernel cpu = false inbounds = true function grad_kernel!(res, img, di) + @kernel inbounds = true function grad_kernel!(res, img, di) idx = @index(Global, Cartesian) res[idx] = img[idx] - img[idx + di] end @@ -33,12 +33,12 @@ function LinearOperatorCollection.grad_t!(res::vecT, g::vecT, shape::NTuple{N,In return res end -@kernel cpu = false inbounds = true function grad_t_kernel_1!(res, g, di) +@kernel inbounds = true function grad_t_kernel_1!(res, g, di) idx = @index(Global, Cartesian) res[idx] += g[idx] end -@kernel cpu = false inbounds = true function grad_t_kernel_2!(res, g, di) +@kernel inbounds = true function grad_t_kernel_2!(res, g, di) idx = @index(Global, Cartesian) res[idx + di] -= g[idx] end diff --git a/src/GradientOp.jl b/src/GradientOp.jl index 8752616..f5c8b03 100644 --- a/src/GradientOp.jl +++ b/src/GradientOp.jl @@ -11,12 +11,37 @@ directional gradient operator along the dimensions `dims` for an array of size ` # Optional Keyword argument * `dims` - dimension(s) along which the gradient is applied; default is `1:length(shape)` + * `S` - storage type of the operator, e.g. `Vector{T}` (default) or `CuArray{T, 1}` + +# Moving to another array backend + + Adapt.adapt(CuArray, op) + +returns a new gradient operator for the target array type (e.g. `Array`, `JLArray`, `CuArray` +or `ROCArray`), and `copy(op)` returns a new one with the same storage type. Both rebuild the +operator and its temporary storage from `shape` and the directions taken from `dims` at +construction; a later change of the vector passed as `dims`, or an iterator consumed by the +construction, has no effect. `op` itself is not modified. The directions stay on the host. """ function GradientOp(::Type{T}; shape::NTuple{N,Int}, dims=1:length(shape), kwargs...) where {T <: Number, N} return GradientOpImpl(T, shape, dims; kwargs...) end +struct GradientOpImpl{T, vecT <: AbstractVector, O <: AbstractLinearOperator{T}, D, N, W} <: GradientOp{T} + op::O + shape::NTuple{N, Int} + # `Int` for a single direction, otherwise the tuple of directions taken from `dims` at construction + dims::D + # temporary vector of the adjoint for several directions, `nothing` for a single direction + tmp::W + function GradientOpImpl(op::O, shape::NTuple{N, Int}, dims::D, tmp::W) where {T, O <: AbstractLinearOperator{T}, D, N, W} + return new{T, storage_type(op), O, D, N, W}(op, shape, dims, tmp) + end +end + function GradientOpImpl(T::Type, shape::NTuple{N,Int}, dims; S = Vector{T}) where N + # Iterate `dims` once: replaying the operator must not depend on the caller's container later on + dims = Tuple(dims) dirs = CartesianIndex{N}[] cols = Int64[] for dim in dims @@ -34,11 +59,12 @@ function GradientOpImpl(T::Type, shape::NTuple{N,Int}, dims; S = Vector{T}) wher tmp = S(undef, ncol) - return LinearOperator{T}(nrow, ncol, false, false, + op = LinearOperator{T}(nrow, ncol, false, false, (res,x) -> (grad!(res,x,shape,dirs, dims, dim_ends)), (res,x) -> (grad_t!(res,x,shape,dirs, dims, dim_ends, tmp)), (res,x) -> (grad_t!(res,x,shape,dirs, dims, dim_ends, tmp)), S = S) + return GradientOpImpl(op, shape, dims, tmp) end function GradientOpImpl(T::Type, shape::NTuple{N,Int}, dim::Int; S = Vector{T}) where N @@ -48,11 +74,33 @@ function GradientOpImpl(T::Type, shape::NTuple{N,Int}, dim::Int; S = Vector{T}) δ[dim] = 1 δ = Tuple(δ) dir = CartesianIndex(δ) - return LinearOperator{T}(nrow, ncol, false, false, + op = LinearOperator{T}(nrow, ncol, false, false, (res,x) -> (grad!(res,x,shape,dir)), (res,x) -> (grad_t!(res,x,shape,dir)), (res,x) -> (grad_t!(res,x,shape,dir)), S = S) + return GradientOpImpl(op, shape, dim, nothing) +end + +function Base.getproperty(op::GradientOpImpl, field::Symbol) + if in(field, (:op, :shape, :dims, :tmp)) + return getfield(op, field) + else + return getproperty(getfield(op, :op), field) + end +end +Base.setproperty!(op::GradientOpImpl, field::Symbol, value) = setproperty!(op.op, field, value) + +storage_type(::GradientOpImpl{T, vecT}) where {T, vecT} = vecT + +# Rebuild the operator, its closures and temporary vector from the shape and the directions +_rebuild_gradient(op::GradientOpImpl, S) = GradientOpImpl(eltype(op), op.shape, op.dims; S = S) + +Base.copy(op::GradientOpImpl) = _rebuild_gradient(op, storage_type(op)) + +function Adapt.adapt_structure(to, op::GradientOpImpl) + target_storage = _adapt_target_storage(to, op.Mv, "GradientOp") + return _rebuild_gradient(op, typeof(target_storage)) end function grad!(res::T, img::U, shape, dirs, dims, dim_ends) where {T<:AbstractVector, U<:AbstractVector} diff --git a/test/testAdapt.jl b/test/testAdapt.jl index 8528559..e92f385 100644 --- a/test/testAdapt.jl +++ b/test/testAdapt.jl @@ -54,6 +54,81 @@ function custom_point_transform(x::AbstractFloat) end custom_point_transform(x::Tuple) = map(custom_point_transform, x) +gradient_test_data(::Type{T}, n) where {T <: Integer} = rand(T(-5):T(5), n) +gradient_test_data(::Type{T}, n) where {T} = rand(T, n) + +gradient_matches(actual, expected, ::Type{T}) where {T <: Integer} = Array(actual) == expected +function gradient_matches(actual, expected, ::Type{T}) where {T} + tol = real(T) === Float32 ? (rtol = 5f-5, atol = 5f-6) : (rtol = 1e-11, atol = 1e-12) + return isapprox(Array(actual), expected; tol...) +end + +function gradient_inputs(op) + T = eltype(op) + m, n = size(op) + return (x = gradient_test_data(T, n), y = gradient_test_data(T, m), + x0 = gradient_test_data(T, n), y0 = gradient_test_data(T, m), α = T(2), β = T(3)) +end + +# Host copies of the forward, transpose and adjoint actions, each with 3- and 5-argument `mul!` +function gradient_actions(op, storage, inputs) + (; x, y, x0, y0, α, β) = inputs + results = Any[] + for (A, input, init) in ((op, x, y0), (transpose(op), y, x0), (adjoint(op), y, x0)) + out = storage(zero(init)) + mul!(out, A, storage(input)) + push!(results, Array(out)) + out = storage(copy(init)) + mul!(out, A, storage(input), α, β) + push!(results, Array(out)) + end + return results +end + +gradient_actions_match(op, storage, inputs, expected) = + all(gradient_matches(actual, reference, eltype(op)) + for (actual, reference) in zip(gradient_actions(op, storage, inputs), expected)) + +# Give the source scratch deterministic contents before it is snapshotted +function initialize_gradient_tmp!(op) + op.tmp === nothing || copyto!(op.tmp, gradient_test_data(eltype(op.tmp), length(op.tmp))) + return op +end + +snapshot_gradient(op) = (Mv = Array(op.Mv), Mtu = Array(op.Mtu), + tmp = op.tmp === nothing ? nothing : Array(op.tmp), + counters = (op.nprod, op.ntprod, op.nctprod)) + +function check_gradient_snapshot(op, snapshot) + @test Array(op.Mv) == snapshot.Mv + @test Array(op.Mtu) == snapshot.Mtu + if snapshot.tmp === nothing + @test op.tmp === nothing + else + @test Array(op.tmp) == snapshot.tmp + end + @test (op.nprod, op.ntprod, op.nctprod) == snapshot.counters +end + +function check_gradient_rebuild(rebuilt, source, array_type) + @test rebuilt isa LinearOperatorCollection.GradientOpImpl + @test getfield(rebuilt, :op) !== getfield(source, :op) + check_adapted_storage(rebuilt, array_type) + @test typeof(rebuilt.Mv) <: array_type + if source.tmp === nothing + @test rebuilt.tmp === nothing + else + @test typeof(rebuilt.tmp) <: array_type + @test rebuilt.tmp !== source.tmp + end + @test rebuilt.dims === source.dims + @test rebuilt.shape == source.shape + @test size(rebuilt) == size(source) + @test eltype(rebuilt) == eltype(source) +end + +storage_vector_type(array_type, T) = typeof(array_type(zeros(T, 0))) + @testset "Adapt reconstruction" begin N = 16 @@ -421,6 +496,85 @@ custom_point_transform(x::Tuple) = map(custom_point_transform, x) @test adapted.fftplan !== source_normal.fftplan end end + + @testset "GradientOp rebuilds closures and scratch" begin + cases = (((7,), nothing, (1,)), ((7,), 1, 1), + ((4, 3), nothing, (1, 2)), ((4, 3), 2, 2), ((4, 3), 1:2, (1, 2)), + ((4, 3), (2, 1), (2, 1)), ((4, 3), (2, 1, 2), (2, 1, 2)), ((4, 3), (), ()), + ((3, 1, 2), nothing, (1, 2, 3)), ((3, 1, 2), 2, 2)) + directions = array_type === Array ? ((Array, Array),) : ((Array, array_type), (array_type, Array)) + for (from, to) in directions, T in (Float32, Float64, ComplexF32, ComplexF64, Int32), + (shape, dims, expected_dims) in cases + S = storage_vector_type(from, T) + source = dims === nothing ? GradientOp(T; shape, S) : GradientOp(T; shape, dims, S) + @test source.dims === expected_dims + inputs = gradient_inputs(source) + expected = gradient_actions(source, from, inputs) + initialize_gradient_tmp!(source) + snapshot = snapshot_gradient(source) + + adapted = Adapt.adapt(to, source) + copied = copy(source) + check_gradient_rebuild(adapted, source, to) + check_gradient_rebuild(copied, source, from) + @test gradient_actions_match(adapted, to, inputs, expected) + @test gradient_actions_match(copied, from, inputs, expected) + check_gradient_snapshot(source, snapshot) + @test gradient_actions_match(source, from, inputs, expected) + end + + @test_throws BoundsError GradientOp(Float32; shape=(4, 3), dims=3) + end + + @testset "GradientOp replays the directions taken at construction" begin + shape = (4, 3) + T = Float32 + x = rand(T, prod(shape)) + reference = GradientOp(T; shape, dims=(2, 1)) + y = rand(T, size(reference, 1)) + expected = (reference * x, adjoint(reference) * y) + acts_like_reference(op) = gradient_matches(op * array_type(x), expected[1], T) && + gradient_matches(adjoint(op) * array_type(y), expected[2], T) + + live_dims = [2, 1] + mutated = GradientOp(T; shape, dims=live_dims) + reverse!(live_dims) + @test mutated.dims === (2, 1) + @test acts_like_reference(Adapt.adapt(array_type, mutated)) + @test acts_like_reference(Adapt.adapt(array_type, copy(mutated))) + + consumed = GradientOp(T; shape, dims=Iterators.Stateful([2, 1, 2])) + @test consumed.dims === (2, 1, 2) + x3 = rand(T, prod(shape)) + expected3 = GradientOp(T; shape, dims=(2, 1, 2)) * x3 + @test gradient_matches(Adapt.adapt(array_type, consumed) * array_type(x3), expected3, T) + @test gradient_matches(Adapt.adapt(array_type, copy(consumed)) * array_type(x3), expected3, T) + + factory = createLinearOperator(GradientOp{T}; shape, dims=(2, 1)) + @test acts_like_reference(Adapt.adapt(array_type, factory)) + + weights = rand(T, size(reference, 1)) + compositions = (ProdOp(WeightingOp(weights), reference), normalOperator(reference), + DiagOp(reference, 2)) + for source in compositions + xs = rand(T, size(source, 2)) + forward = source * xs + adapted = Adapt.adapt(array_type, source) + check_adapted_storage(adapted, array_type) + @test gradient_matches(adapted * array_type(xs), forward, T) + # The normal operator defines no adjoint action; it is self-adjoint + if !(source isa LinearOperatorCollection.NormalOpImpl) + ys = rand(T, size(source, 1)) + backward = adjoint(source) * ys + @test gradient_matches(adjoint(adapted) * array_type(ys), backward, T) + end + end + @test compositions[3].ops[1].tmp !== compositions[3].ops[2].tmp + diag = Adapt.adapt(array_type, compositions[3]) + @test all(op isa LinearOperatorCollection.GradientOpImpl for op in diag.ops) + @test diag.ops[1].tmp !== diag.ops[2].tmp + @test all(typeof(op.tmp) <: array_type for op in diag.ops) + end end end From 7779d6e4e10e076977b5534013a8d665e57791b0 Mon Sep 17 00:00:00 2001 From: shukebeta Date: Tue, 3 Nov 2026 03:51:52 +1300 Subject: [PATCH 2/4] Reject element-typed Adapt targets that change the GradientOp element 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 --- src/GradientOp.jl | 6 ++++++ test/testAdapt.jl | 16 ++++++++++++++++ 2 files changed, 22 insertions(+) diff --git a/src/GradientOp.jl b/src/GradientOp.jl index f5c8b03..27897f2 100644 --- a/src/GradientOp.jl +++ b/src/GradientOp.jl @@ -22,6 +22,8 @@ or `ROCArray`), and `copy(op)` returns a new one with the same storage type. Bot operator and its temporary storage from `shape` and the directions taken from `dims` at construction; a later change of the vector passed as `dims`, or an iterator consumed by the construction, has no effect. `op` itself is not modified. The directions stay on the host. +A target with an element type, e.g. `CuArray{Float32}`, must match the element type of `op`; +otherwise an `ArgumentError` is thrown. """ function GradientOp(::Type{T}; shape::NTuple{N,Int}, dims=1:length(shape), kwargs...) where {T <: Number, N} return GradientOpImpl(T, shape, dims; kwargs...) @@ -100,6 +102,10 @@ Base.copy(op::GradientOpImpl) = _rebuild_gradient(op, storage_type(op)) function Adapt.adapt_structure(to, op::GradientOpImpl) target_storage = _adapt_target_storage(to, op.Mv, "GradientOp") + # The rebuild keeps the element type of the operator, so the target storage must share it + eltype(target_storage) == eltype(op.Mv) || throw(ArgumentError( + "unsupported Adapt target $(to) for GradientOp: element type $(eltype(target_storage)) " * + "differs from the operator's $(eltype(op.Mv))")) return _rebuild_gradient(op, typeof(target_storage)) end diff --git a/test/testAdapt.jl b/test/testAdapt.jl index e92f385..e6d5366 100644 --- a/test/testAdapt.jl +++ b/test/testAdapt.jl @@ -526,6 +526,22 @@ storage_vector_type(array_type, T) = typeof(array_type(zeros(T, 0))) @test_throws BoundsError GradientOp(Float32; shape=(4, 3), dims=3) end + @testset "GradientOp element-typed targets" begin + directions = array_type === Array ? ((Array, Array),) : ((Array, array_type), (array_type, Array)) + for (from, to) in directions, dims in ((1, 2), 2) + source = GradientOp(Float64; shape=(4, 3), dims, S=storage_vector_type(from, Float64)) + inputs = gradient_inputs(source) + expected = gradient_actions(source, from, inputs) + + adapted = Adapt.adapt(to{Float64}, source) + check_gradient_rebuild(adapted, source, to) + @test gradient_actions_match(adapted, to, inputs, expected) + + @test_throws ArgumentError Adapt.adapt(to{Float32}, source) + @test_throws ArgumentError Adapt.adapt(to{ComplexF64}, source) + end + end + @testset "GradientOp replays the directions taken at construction" begin shape = (4, 3) T = Float32 From 8d15decba5fd15bfb836c624e265e6fb06172224 Mon Sep 17 00:00:00 2001 From: shukebeta Date: Tue, 3 Nov 2026 03:59:40 +1300 Subject: [PATCH 3/4] Check GradientOp Adapt storage against the operator element type 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 --- src/GradientOp.jl | 11 ++++++----- test/testAdapt.jl | 13 +++++++++++++ 2 files changed, 19 insertions(+), 5 deletions(-) diff --git a/src/GradientOp.jl b/src/GradientOp.jl index 27897f2..4552f25 100644 --- a/src/GradientOp.jl +++ b/src/GradientOp.jl @@ -22,8 +22,9 @@ or `ROCArray`), and `copy(op)` returns a new one with the same storage type. Bot operator and its temporary storage from `shape` and the directions taken from `dims` at construction; a later change of the vector passed as `dims`, or an iterator consumed by the construction, has no effect. `op` itself is not modified. The directions stay on the host. -A target with an element type, e.g. `CuArray{Float32}`, must match the element type of `op`; -otherwise an `ArgumentError` is thrown. +The adapted storage must have the element type of `op`: an element-typed target such as +`CuArray{Float32}` for a `Float64` operator throws an `ArgumentError`, and so does a target +without an element type for an operator built with an `S` of another element type. """ function GradientOp(::Type{T}; shape::NTuple{N,Int}, dims=1:length(shape), kwargs...) where {T <: Number, N} return GradientOpImpl(T, shape, dims; kwargs...) @@ -103,9 +104,9 @@ Base.copy(op::GradientOpImpl) = _rebuild_gradient(op, storage_type(op)) function Adapt.adapt_structure(to, op::GradientOpImpl) target_storage = _adapt_target_storage(to, op.Mv, "GradientOp") # The rebuild keeps the element type of the operator, so the target storage must share it - eltype(target_storage) == eltype(op.Mv) || throw(ArgumentError( - "unsupported Adapt target $(to) for GradientOp: element type $(eltype(target_storage)) " * - "differs from the operator's $(eltype(op.Mv))")) + eltype(target_storage) == eltype(op) || throw(ArgumentError( + "unsupported Adapt target $(to) for GradientOp: storage element type $(eltype(target_storage)) " * + "differs from the operator's element type $(eltype(op))")) return _rebuild_gradient(op, typeof(target_storage)) end diff --git a/test/testAdapt.jl b/test/testAdapt.jl index e6d5366..0e3a76c 100644 --- a/test/testAdapt.jl +++ b/test/testAdapt.jl @@ -540,6 +540,19 @@ storage_vector_type(array_type, T) = typeof(array_type(zeros(T, 0))) @test_throws ArgumentError Adapt.adapt(to{Float32}, source) @test_throws ArgumentError Adapt.adapt(to{ComplexF64}, source) end + + # Storage of another element type than the operator is rejected for any target + for dims in ((1, 2), 2) + mixed = GradientOp(Float64; shape=(4, 3), dims, S=Vector{Float32}) + x = rand(Float64, size(mixed, 2)) + @test_throws ArgumentError Adapt.adapt(array_type, mixed) + @test_throws ArgumentError Adapt.adapt(Array, mixed) + @test copy(mixed) * x == mixed * x + # A target naming the operator's element type rebuilds consistent storage + consistent = Adapt.adapt(array_type{Float64}, mixed) + check_adapted_storage(consistent, array_type{Float64}) + @test gradient_matches(consistent * array_type(x), GradientOp(Float64; shape=(4, 3), dims) * x, Float64) + end end @testset "GradientOp replays the directions taken at construction" begin From 9fe52df2cd71f233973e306628eec42ddf14a3f6 Mon Sep 17 00:00:00 2001 From: shukebeta Date: Tue, 3 Nov 2026 04:04:09 +1300 Subject: [PATCH 4/4] Reject Adapt of a GradientOp whose storage element type differs 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 --- src/GradientOp.jl | 15 ++++++++++----- test/testAdapt.jl | 18 +++++++++--------- 2 files changed, 19 insertions(+), 14 deletions(-) diff --git a/src/GradientOp.jl b/src/GradientOp.jl index 4552f25..85c3d3c 100644 --- a/src/GradientOp.jl +++ b/src/GradientOp.jl @@ -23,8 +23,9 @@ operator and its temporary storage from `shape` and the directions taken from `d construction; a later change of the vector passed as `dims`, or an iterator consumed by the construction, has no effect. `op` itself is not modified. The directions stay on the host. The adapted storage must have the element type of `op`: an element-typed target such as -`CuArray{Float32}` for a `Float64` operator throws an `ArgumentError`, and so does a target -without an element type for an operator built with an `S` of another element type. +`CuArray{Float32}` for a `Float64` operator throws an `ArgumentError`. An operator built with +an `S` of another element type cannot be adapted to any target and throws an `ArgumentError`; +`copy` keeps it as it is. """ function GradientOp(::Type{T}; shape::NTuple{N,Int}, dims=1:length(shape), kwargs...) where {T <: Number, N} return GradientOpImpl(T, shape, dims; kwargs...) @@ -102,11 +103,15 @@ _rebuild_gradient(op::GradientOpImpl, S) = GradientOpImpl(eltype(op), op.shape, Base.copy(op::GradientOpImpl) = _rebuild_gradient(op, storage_type(op)) function Adapt.adapt_structure(to, op::GradientOpImpl) + T = eltype(op) + # A rebuild with other storage would change the precision of the adjoint's temporary vector + eltype(storage_type(op)) == T || throw(ArgumentError( + "cannot adapt GradientOp with storage element type $(eltype(storage_type(op))) " * + "and operator element type $(T); construct it with matching element types")) target_storage = _adapt_target_storage(to, op.Mv, "GradientOp") - # The rebuild keeps the element type of the operator, so the target storage must share it - eltype(target_storage) == eltype(op) || throw(ArgumentError( + eltype(target_storage) == T || throw(ArgumentError( "unsupported Adapt target $(to) for GradientOp: storage element type $(eltype(target_storage)) " * - "differs from the operator's element type $(eltype(op))")) + "differs from the operator's element type $(T)")) return _rebuild_gradient(op, typeof(target_storage)) end diff --git a/test/testAdapt.jl b/test/testAdapt.jl index 0e3a76c..66d151f 100644 --- a/test/testAdapt.jl +++ b/test/testAdapt.jl @@ -541,17 +541,17 @@ storage_vector_type(array_type, T) = typeof(array_type(zeros(T, 0))) @test_throws ArgumentError Adapt.adapt(to{ComplexF64}, source) end - # Storage of another element type than the operator is rejected for any target + # An operator whose storage element type differs from its own is rejected for any target for dims in ((1, 2), 2) mixed = GradientOp(Float64; shape=(4, 3), dims, S=Vector{Float32}) - x = rand(Float64, size(mixed, 2)) - @test_throws ArgumentError Adapt.adapt(array_type, mixed) - @test_throws ArgumentError Adapt.adapt(Array, mixed) - @test copy(mixed) * x == mixed * x - # A target naming the operator's element type rebuilds consistent storage - consistent = Adapt.adapt(array_type{Float64}, mixed) - check_adapted_storage(consistent, array_type{Float64}) - @test gradient_matches(consistent * array_type(x), GradientOp(Float64; shape=(4, 3), dims) * x, Float64) + inputs = gradient_inputs(mixed) + expected = gradient_actions(mixed, Array, inputs) + for to in unique((array_type, Array, array_type{Float64}, array_type{Float32}, Array{Float64})) + @test_throws ArgumentError Adapt.adapt(to, mixed) + end + copied = copy(mixed) + @test LinearOperators.storage_type(copied) == Vector{Float32} + @test all(actual == reference for (actual, reference) in zip(gradient_actions(copied, Array, inputs), expected)) end end