diff --git a/Project.toml b/Project.toml index 36edfbd21..403da1e65 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "ChainRulesCore" uuid = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" -version = "1.26.4" +version = "1.26.5" [deps] LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" diff --git a/src/projection.jl b/src/projection.jl index 90e11ddb4..ca577aef2 100644 --- a/src/projection.jl +++ b/src/projection.jl @@ -221,23 +221,9 @@ function ProjectTo(xs::AbstractArray) end end -function (project::ProjectTo{AbstractArray})(dx::AbstractArray{S,M}) where {S,M} - # First deal with shape. The rule is that we reshape to add or remove trivial dimensions - # like dx = ones(4,1), where x = ones(4), but throw an error on dx = ones(1,4) etc. - dy = if axes(dx) === project.axes - dx - elseif project.axes isa Tuple{Vararg{Base.OneTo}} && axes(dx) isa typeof(project.axes) - # Same number of `OneTo` axes that are not identical: the sizes differ. Resolved from - # the types alone, so inference does not see the type-unstable `reshape` below. - throw(_projection_mismatch(project.axes, size(dx))) - else - for d in 1:max(M, length(project.axes)) - if size(dx, d) != length(get(project.axes, d, 1)) - throw(_projection_mismatch(project.axes, size(dx))) - end - end - reshape(dx, project.axes) - end +function (project::ProjectTo{AbstractArray})(dx::AbstractArray{S}) where {S} + # First deal with shape: + dy = _projection_reshape(project.axes, dx) # Then deal with the elements. One projector if AbstractArray{<:Number}, # or one per element for arrays of anything else, including arrays of arrays: dz = if hasproperty(project, :element) @@ -270,6 +256,25 @@ function (project::ProjectTo{AbstractArray})(dx::Number) # ... so we restore fro return fill(project.element(dx)) end +# Reshape dx to axes_x, adding or removing trivial dimensions like dx = ones(4,1), +# where x = ones(4), but throw an error on dx = ones(1,4) etc. +function _projection_reshape(axes_x::Tuple, dx::AbstractArray) + return if axes(dx) === axes_x + dx + elseif axes_x isa Tuple{Vararg{Base.OneTo}} && axes(dx) isa typeof(axes_x) + # Same number of `OneTo` axes that are not identical: the sizes differ. Resolved from + # the types alone, so inference does not see the type-unstable `reshape` below. + throw(_projection_mismatch(axes_x, size(dx))) + else + for d in 1:max(ndims(dx), length(axes_x)) + if size(dx, d) != length(get(axes_x, d, 1)) + throw(_projection_mismatch(axes_x, size(dx))) + end + end + reshape(dx, axes_x) + end +end + function _projection_mismatch(axes_x::Tuple, size_dx::Tuple) size_x = map(length, axes_x) return DimensionMismatch( @@ -439,7 +444,11 @@ end # Diagonal ProjectTo(x::Diagonal) = ProjectTo{Diagonal}(; diag=ProjectTo(x.diag)) -(project::ProjectTo{Diagonal})(dx::AbstractMatrix) = Diagonal(project.diag(diag(dx))) +function (project::ProjectTo{Diagonal})(dx::AbstractArray) + ax = only(project.diag.axes) + dy = _projection_reshape((ax, ax), dx) + return Diagonal(project.diag(diag(dy))) +end (project::ProjectTo{Diagonal})(dx::Diagonal) = Diagonal(project.diag(dx.diag)) # Symmetric @@ -489,10 +498,12 @@ end # Weird -- not exhaustive! # one strategy is to recurse into the struct: ProjectTo(x::Bidiagonal{T}) where {T<:Number} = generic_projector(x) -function (project::ProjectTo{Bidiagonal})(dx::AbstractMatrix) +function (project::ProjectTo{Bidiagonal})(dx::AbstractArray) + ax = only(project.dv.axes) + dy = _projection_reshape((ax, ax), dx) uplo = LinearAlgebra.sym_uplo(project.uplo) - dv = project.dv(diag(dx)) - ev = project.ev(uplo === :U ? diag(dx, 1) : diag(dx, -1)) + dv = project.dv(diag(dy)) + ev = project.ev(uplo === :U ? diag(dy, 1) : diag(dy, -1)) return Bidiagonal(dv, ev, uplo) end function (project::ProjectTo{Bidiagonal})(dx::Bidiagonal) @@ -507,9 +518,11 @@ function (project::ProjectTo{Bidiagonal})(dx::Bidiagonal) end ProjectTo(x::SymTridiagonal{T}) where {T<:Number} = generic_projector(x) -function (project::ProjectTo{SymTridiagonal})(dx::AbstractMatrix) - dv = project.dv(diag(dx)) - ev = project.ev((diag(dx, 1) .+ diag(dx, -1)) ./ 2) +function (project::ProjectTo{SymTridiagonal})(dx::AbstractArray) + ax = only(project.dv.axes) + dy = _projection_reshape((ax, ax), dx) + dv = project.dv(diag(dy)) + ev = project.ev((diag(dy, 1) .+ diag(dy, -1)) ./ 2) return SymTridiagonal(dv, ev) end (project::ProjectTo{SymTridiagonal})(dx::SymTridiagonal) = generic_projection(project, dx) diff --git a/test/projection.jl b/test/projection.jl index 912131242..95046147a 100644 --- a/test/projection.jl +++ b/test/projection.jl @@ -322,6 +322,18 @@ struct NoSuperType end @test pstri(stri) == real(stri) @test_throws DimensionMismatch pstri(rand(ComplexF32, 3, 2)) + # same shape rules as ProjectTo{AbstractArray}, e.g. trailing singleton dimensions + for p in (pdiag, pbi, pstri) + @test @inferred(p(reshape(1:9, 3, 3, 1))) == p(reshape(1:9, 3, 3)) + @test p(reshape(1:9, 3, 3, 1, 1)) == p(reshape(1:9, 3, 3)) + @test p(rand(ComplexF32, 3, 3, 1)) isa typeof(p(rand(3, 3))) + @test_throws DimensionMismatch p(rand(3, 3, 2)) + @test_throws DimensionMismatch p(rand(3, 1, 3)) + @test_throws DimensionMismatch p(rand(3, 4)) + end + psdiag = ProjectTo(Diagonal(SA[1.0, 2.0, 3.0])) + @test psdiag(reshape(1:9, 3, 3, 1)) == psdiag(reshape(1:9, 3, 3)) + ptri = ProjectTo(Tridiagonal(rand(3, 3))) @test ptri(reshape(1:9, 3, 3)) == [1.0 4.0 0.0; 2.0 5.0 8.0; 0.0 6.0 9.0] @test ptri(ptri(reshape(1:9, 3, 3))) == ptri(reshape(1:9, 3, 3))