typecheck for array rule.

Author: AstitvaAggarwal

DifferentiationInterface/ext/DifferentiationInterfaceMooncakeExt/differentiate_with.jl

@@ -33,13 +33,15 @@ function Mooncake.rrule!!(dw::CoDual{<:DI.DifferentiateWith}, x::CoDual{<:Abstra
     # output is a vector, so we need to use the vector pullback
     function pullback_array!!(dy::NoRData)
         tx = DI.pullback(f, backend, primal_x, (fdata(y.dx),))
+        @assert only(tx) isa rdata_type(typeof(primal_x))
         fdata_arg .+= only(tx)
         return NoRData(), dy
     end
 
     # output is a scalar, so we can use the scalar pullback
     function pullback_scalar!!(dy::Number)
         tx = DI.pullback(f, backend, primal_x, (dy,))
+        @assert only(tx) isa rdata_type(typeof(primal_x))
         fdata_arg .+= only(tx)
         return NoRData(), NoRData()
     end