assertion for array inputs

Author: AstitvaAggarwal

DifferentiationInterface/ext/DifferentiationInterfaceMooncakeExt/differentiate_with.jl

@@ -33,15 +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))
+        @assert first(only(tx)) isa rdata_type(typeof(first(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))
+        @assert first(only(tx)) isa rdata_type(typeof(first(primal_x)))
         fdata_arg .+= only(tx)
         return NoRData(), NoRData()
     end

DifferentiationInterface/test/Back/DifferentiateWith/test.jl

@@ -13,7 +13,9 @@ LOGGING = get(ENV, "CI", "false") == "false"
 
 function differentiatewith_scenarios()
     bad_scens =  # these closurified scenarios have mutation and type constraints
-        filter(default_scenarios(; include_normal=false, include_closurified=true)) do scen
+        filter(
+            DIT.default_scenarios(; include_normal=false, include_closurified=true)
+        ) do scen
             DIT.function_place(scen) == :out
         end
     good_scens = map(bad_scens) do scen