diff --git a/README.md b/README.md index bf041dfc..3264bbfa 100644 --- a/README.md +++ b/README.md @@ -45,9 +45,9 @@ pB = ArrayPartition(b) pA .* pB # Now all standard array stuff works! # or do: -x0 = rand(3,3) -v0 = rand(3,3) -a0 = rand(3,3) +x0 = rand(3, 3) +v0 = rand(3, 3) +a0 = rand(3, 3) u0 = ArrayPartition(x0, v0, a0) u0.x[1] == x0 # true diff --git a/ext/RecursiveArrayToolsKernelAbstractionsExt.jl b/ext/RecursiveArrayToolsKernelAbstractionsExt.jl index f4e5f9e6..673d60f2 100644 --- a/ext/RecursiveArrayToolsKernelAbstractionsExt.jl +++ b/ext/RecursiveArrayToolsKernelAbstractionsExt.jl @@ -3,7 +3,6 @@ module RecursiveArrayToolsKernelAbstractionsExt import RecursiveArrayTools: VectorOfArray import KernelAbstractions - function KernelAbstractions.get_backend(x::VectorOfArray) u = parent(x) if length(u) == 0 diff --git a/ext/RecursiveArrayToolsZygoteExt.jl b/ext/RecursiveArrayToolsZygoteExt.jl index 7502a71f..e1987dd8 100644 --- a/ext/RecursiveArrayToolsZygoteExt.jl +++ b/ext/RecursiveArrayToolsZygoteExt.jl @@ -5,7 +5,6 @@ using RecursiveArrayTools using Zygote using Zygote: FillArrays, ChainRulesCore, literal_getproperty, @adjoint - # Define a new species of projection operator for this type: # ChainRulesCore.ProjectTo(x::VectorOfArray) = ChainRulesCore.ProjectTo{VectorOfArray}() @@ -216,10 +215,12 @@ z̄ -> (nothing, conj.(z̄)) @adjoint Broadcast.broadcasted(::typeof(real), x::AbstractVectorOfArray) = real.(x), z̄ -> (nothing, real.(z̄)) -@adjoint Broadcast.broadcasted(::typeof(imag), x::AbstractVectorOfArray) = imag.(x), +@adjoint Broadcast.broadcasted( + ::typeof(imag), x::AbstractVectorOfArray) = imag.(x), z̄ -> (nothing, im .* real.(z̄)) -@adjoint Broadcast.broadcasted(::typeof(abs2), x::AbstractVectorOfArray) = abs2.(x), +@adjoint Broadcast.broadcasted(::typeof(abs2), + x::AbstractVectorOfArray) = abs2.(x), z̄ -> (nothing, 2 .* real.(z̄) .* x) @adjoint function Broadcast.broadcasted( @@ -260,7 +261,9 @@ end end end -@adjoint Broadcast.broadcasted(::Type{T}, x::AbstractVectorOfArray) where {T <: Number} = T.(x), +@adjoint Broadcast.broadcasted(::Type{T}, + x::AbstractVectorOfArray) where {T <: + Number} = T.(x), ȳ -> (nothing, Zygote._project(x, ȳ)) function Zygote.unbroadcast(x::AbstractVectorOfArray, x̄) @@ -273,11 +276,16 @@ function Zygote.unbroadcast(x::AbstractVectorOfArray, x̄) end end -@adjoint Broadcast.broadcasted(::Broadcast.AbstractArrayStyle, f::F, a::AbstractVectorOfArray, b) where {F} = _broadcast_generic( +@adjoint Broadcast.broadcasted( + ::Broadcast.AbstractArrayStyle, f::F, a::AbstractVectorOfArray, + b) where {F} = _broadcast_generic( __context__, f, a, b) -@adjoint Broadcast.broadcasted(::Broadcast.AbstractArrayStyle, f::F, a, b::AbstractVectorOfArray) where {F} = _broadcast_generic( +@adjoint Broadcast.broadcasted(::Broadcast.AbstractArrayStyle, f::F, a, + b::AbstractVectorOfArray) where {F} = _broadcast_generic( __context__, f, a, b) -@adjoint Broadcast.broadcasted(::Broadcast.AbstractArrayStyle, f::F, a::AbstractVectorOfArray, b::AbstractVectorOfArray) where {F} = _broadcast_generic( +@adjoint Broadcast.broadcasted( + ::Broadcast.AbstractArrayStyle, f::F, a::AbstractVectorOfArray, + b::AbstractVectorOfArray) where {F} = _broadcast_generic( __context__, f, a, b) @inline function _broadcast_generic(__context__, f::F, args...) where {F} diff --git a/src/RecursiveArrayTools.jl b/src/RecursiveArrayTools.jl index cb2f4255..32e72c82 100644 --- a/src/RecursiveArrayTools.jl +++ b/src/RecursiveArrayTools.jl @@ -28,14 +28,7 @@ An AbstractVectorOfArray subtype should match the following behaviors. !!! note - In 2023 the linear indexing `A[i]`` was deprecated. It previously had the behavior that - `A[i] = A.u[i]`. However, this is incompatible with standard `AbstractArray` interfaces, - Since if `A = VectorOfArray([[1,2],[3,4]])` and `A` is supposed to act like `[1 3; 2 4]`, - then there is a difference `A[1] = [1,2]` for the VectorOfArray while `A[1] = 1` for the - matrix. This causes many issues if `AbstractVectorOfArray <: AbstractArray`. Thus we - plan in 2026 to complete the deprecation and thus have a breaking update where `A[i]` - matches the linear indexing of an `AbstractArray`, and then making - `AbstractVectorOfArray <: AbstractArray`. Until then, `AbstractVectorOfArray` due to + In 2023 the linear indexing `A[i]`` was deprecated. It previously had the behavior that `A[i] = A.u[i]`. However, this is incompatible with standard `AbstractArray`interfaces, Since if`A = VectorOfArray([[1,2],[3,4]])`and`A`is supposed to act like`[1 3; 2 4]`, then there is a difference `A[1] = [1,2]`for the VectorOfArray while`A[1] = 1`for the matrix. This causes many issues if`AbstractVectorOfArray <: AbstractArray`. Thus we plan in 2026 to complete the deprecation and thus have a breaking update where `A[i]`matches the linear indexing of an`AbstractArray`, and then making `AbstractVectorOfArray <: AbstractArray`. Until then, `AbstractVectorOfArray` due to this interface break but manually implements an AbstractArray-like interface for future compatibility. @@ -43,7 +36,7 @@ An AbstractVectorOfArray subtype should match the following behaviors. An AbstractVectorOfArray has the following fields: -* `u` which holds the Vector of values at each timestep + - `u` which holds the Vector of values at each timestep ## Array Interface @@ -125,7 +118,7 @@ additional properties: An AbstractDiffEqArray adds the following fields: -* `t` which holds the times of each timestep. + - `t` which holds the times of each timestep. """ abstract type AbstractDiffEqArray{T, N, A} <: AbstractVectorOfArray{T, N, A} end diff --git a/src/named_array_partition.jl b/src/named_array_partition.jl index 1db2b981..a58c9a7b 100644 --- a/src/named_array_partition.jl +++ b/src/named_array_partition.jl @@ -115,7 +115,9 @@ end # hook into ArrayPartition broadcasting routines @inline RecursiveArrayTools.npartitions(x::NamedArrayPartition) = npartitions(ArrayPartition(x)) -@inline RecursiveArrayTools.unpack(bc::Broadcast.Broadcasted{Broadcast.ArrayStyle{NamedArrayPartition}}, i) = Broadcast.Broadcasted( +@inline RecursiveArrayTools.unpack( + bc::Broadcast.Broadcasted{Broadcast.ArrayStyle{NamedArrayPartition}}, + i) = Broadcast.Broadcasted( bc.f, RecursiveArrayTools.unpack_args(i, bc.args)) @inline RecursiveArrayTools.unpack(x::NamedArrayPartition, i) = unpack(ArrayPartition(x), i) @@ -123,7 +125,10 @@ function Base.copy(A::NamedArrayPartition{T, S, NT}) where {T, S, NT} NamedArrayPartition{T, S, NT}(copy(ArrayPartition(A)), getfield(A, :names_to_indices)) end -@inline NamedArrayPartition(f::F, N, names_to_indices) where {F <: Function} = NamedArrayPartition( +@inline NamedArrayPartition(f::F, + N, + names_to_indices) where {F <: + Function} = NamedArrayPartition( ArrayPartition(ntuple(f, Val(N))), names_to_indices) @inline function Base.copy(bc::Broadcast.Broadcasted{Broadcast.ArrayStyle{NamedArrayPartition}}) diff --git a/src/utils.jl b/src/utils.jl index 8c34c316..908f55ad 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -110,6 +110,7 @@ function recursivefill!(bs::AbstractVectorOfArray{T, N}, a::T2) where {T <: StaticArraysCore.StaticArray, T2 <: StaticArraysCore.StaticArray, N} @inbounds for b in bs, i in eachindex(b) + b[i] = copy(a) end end @@ -126,6 +127,7 @@ function recursivefill!(bs::AbstractVectorOfArray{T, N}, a::T2) where {T <: StaticArraysCore.SArray, T2 <: Union{Number, Bool}, N} @inbounds for b in bs, i in eachindex(b) + b[i] = fill(a, typeof(b[i])) end end diff --git a/test/adjoints.jl b/test/adjoints.jl index 2fd48fd1..af2abd42 100644 --- a/test/adjoints.jl +++ b/test/adjoints.jl @@ -13,6 +13,7 @@ function loss3(x) y = VectorOfArray([x .* i for i in 1:5]) tmp = 0.0 for i in 1:5, j in 1:5 + tmp += y[i, j] end tmp @@ -22,6 +23,7 @@ function loss4(x) y = DiffEqArray([x .* i for i in 1:5], 1:5) tmp = 0.0 for i in 1:5, j in 1:5 + tmp += y[i, j] end tmp diff --git a/test/gpu/vectorofarray_gpu.jl b/test/gpu/vectorofarray_gpu.jl index 68ea44d7..dfb63292 100644 --- a/test/gpu/vectorofarray_gpu.jl +++ b/test/gpu/vectorofarray_gpu.jl @@ -43,12 +43,12 @@ va_cu = convert(AbstractArray, va) @test size(va_cu) == size(x) a = VectorOfArray([ones(2) for i in 1:3]) -_a = Adapt.adapt(CuArray,a) +_a = Adapt.adapt(CuArray, a) @test _a isa VectorOfArray @test _a.u isa Vector{<:CuArray} -b = DiffEqArray([ones(2) for i in 1:3],ones(2)) -_b = Adapt.adapt(CuArray,b) +b = DiffEqArray([ones(2) for i in 1:3], ones(2)) +_b = Adapt.adapt(CuArray, b) @test _b isa DiffEqArray @test _b.u isa Vector{<:CuArray} @test _b.t isa CuArray diff --git a/test/partitions_test.jl b/test/partitions_test.jl index d444a4ff..3c0c2232 100644 --- a/test/partitions_test.jl +++ b/test/partitions_test.jl @@ -286,7 +286,8 @@ end @testset "Copy and zero with type changing array" begin # Motivating use case for this is ArrayPartitions of Arrow arrays which are mmap:ed and change type when copied struct TypeChangingArray{T, N} <: AbstractArray{T, N} end - Base.copy(::TypeChangingArray{T, N}) where {T, N} = Array{T, N}(undef, + Base.copy(::TypeChangingArray{ + T, N}) where {T, N} = Array{T, N}(undef, ntuple(_ -> 0, N)) Base.zero(::TypeChangingArray{T, N}) where {T, N} = zeros(T, ntuple(_ -> 0, N))