Cholesky type stability

Project.toml

@@ -1,6 +1,6 @@
 name = "Zygote"
 uuid = "e88e6eb3-aa80-5325-afca-941959d7151f"
-version = "0.6.30"
+version = "0.6.31"
 
 [deps]
 AbstractFFTs = "621f4979-c628-5d54-868e-fcf4e3e8185c"

src/lib/array.jl

@@ -741,15 +741,19 @@ end
     return ((uplo=nothing, info=nothing, factors=nothing),)
   end
 end
-@adjoint function literal_getproperty(C::Cholesky, ::Val{:U})
+@adjoint function literal_getproperty(
+  C::Cholesky{T, <:StridedMatrix{T}} where {T<:Real}, ::Val{:U}
+)
   return literal_getproperty(C, Val(:U)), function(Δ)
-    Δ_factors = C.uplo == 'U' ? UpperTriangular(Δ) : LowerTriangular(copy(Δ'))
+    Δ_factors = C.uplo == 'U' ? triu!(collect(Δ)) : tril!(collect(Δ'))
     return ((uplo=nothing, info=nothing, factors=Δ_factors),)
   end
 end
-@adjoint function literal_getproperty(C::Cholesky, ::Val{:L})
+@adjoint function literal_getproperty(
+  C::Cholesky{T, <:StridedMatrix{T}} where {T<:Real}, ::Val{:L}
+)
   return literal_getproperty(C, Val(:L)), function(Δ)
-    Δ_factors = C.uplo == 'L' ? LowerTriangular(Δ) : UpperTriangular(copy(Δ'))
+    Δ_factors = C.uplo == 'L' ? tril!(collect(Δ)) : triu!(collect(Δ'))
     return ((uplo=nothing, info=nothing, factors=Δ_factors),)
   end
 end

test/gradcheck.jl

@@ -801,6 +801,17 @@ end
     @test gradtest(A->logdet(cholesky(A' * A + I)), A)
     @test gradtest(B->cholesky(Symmetric(B)).U, A * A' + I)
     @test gradtest(B->logdet(cholesky(Symmetric(B))), A * A' + I)
+
+    @testset "inference" begin
+      out, pb = _pullback(Context(), C -> C.U, cholesky(Symmetric(A'A + I, :U)))
+      @inferred pb(out)
+      out, pb = _pullback(Context(), C -> C.U, cholesky(Symmetric(A'A + I, :L)))
+      @inferred pb(out)
+      out, pb = _pullback(Context(), C -> C.L, cholesky(Symmetric(A'A + I, :U)))
+      @inferred pb(out)
+      out, pb = _pullback(Context(), C -> C.L, cholesky(Symmetric(A'A + I, :L)))
+      @inferred pb(out)
+    end
   end
   @testset "cholesky - scalar" begin
     rng = MersenneTwister(123456)

test/runtests.jl

@@ -1,5 +1,5 @@
 using Zygote, Test
-using Zygote: gradient, ZygoteRuleConfig
+using Zygote: gradient, ZygoteRuleConfig, _pullback, Context
 using CUDA
 using CUDA: has_cuda