Update differentiate_with.jl

Author: AstitvaAggarwal

DifferentiationInterface/ext/DifferentiationInterfaceMooncakeExt/differentiate_with.jl

@@ -9,14 +9,14 @@ function Mooncake.rrule!!(dw::CoDual{<:DI.DifferentiateWith}, x::CoDual{<:Number
     # 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(x))
+        @assert only(tx) isa rdata_type(typeof(primal_x))
         return NoRData(), only(tx)
     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(x))
+        @assert only(tx) isa rdata_type(typeof(primal_x))
         return NoRData(), only(tx)
     end