test: broadcasting preserves nested types

Author: DhairyaLGandhi

test/adjoints.jl

@@ -92,3 +92,7 @@ loss(x)
       VectorOfArray([collect((3i):(3i + 3)) for i in 1:5])
 @test Zygote.gradient(loss10, x)[1] == ForwardDiff.gradient(loss10, x)
 @test Zygote.gradient(loss11, x)[1] == ForwardDiff.gradient(loss11, x)
+
+x = ArrayPartition(ArrayPartition(rand(3,4), rand(3,4)), rand(2))
+g = Zygote.gradient(norm, x)[1]
+@test g isa typeof(x)

test/basic_indexing.jl

@@ -280,3 +280,9 @@ x = VectorOfArray(StructArray{SVector{1, Float64}}(ntuple(_ -> [1.0, 2.0], 1)))
 y = 2 * x
 @. x = y
 @test all(all.(y .== x))
+
+
+x = ArrayPartition(ArrayPartition(rand(3,4), rand(3,4)), rand(2))
+@test (x .* 1.2) isa ArrayPartition{<:Any, <:ArrayPartition}
+
+g = Zygote.gradient(norm, x)[1]