Merge pull request #1670 from AayushSabharwal/as/fix-getindex-der

fix: fix derivatives of indexed array expressions

Author: AayushSabharwal

src/diff.jl

@@ -263,16 +263,25 @@ function executediff(D, arg, simplify=false; throw_no_derivative=false)
             end
         end
     elseif op === getindex
+        arr = arguments(arg)[1]
         inner_args = arguments(arguments(arg)[1])
+        idxs = @views arguments(arg)[2:end]
         c = 0
-        for a in inner_args
+        # We know `D.x` is in `arg`, so the derivative is not identically zero.
+        # `arg` cannot be `D.x` since, that would have also early exited. 
+        for (i, a) in enumerate(inner_args)
+            der = derivative_idx(arr, i)
             if isequal(a, D.x)
-                return D(arg)
+                der isa NoDeriv && return D(arg)
+                c += der[idxs...]
+                continue
+            elseif der isa NoDeriv
+                c += Differential(a)(arg) * executediff(D, a)
             else
-                c += Differential(a)(arg) * D(a)
+                c += der[idxs...] * executediff(D, a)
             end
         end
-        return expand_derivatives(c)
+        return c
     elseif op === ifelse
         args = arguments(arg)
         O = op(args[1],

test/diff.jl

@@ -634,4 +634,35 @@ end
     @test iszero(Symbolics.unwrap.(Symbolics.gradient(f, vp) .- Symbolics.gradient(f, p)))
     @test iszero(Symbolics.unwrap.(Symbolics.hessian(f, vp) .- Symbolics.hessian(f, p)))
     @test iszero(Symbolics.unwrap.(Symbolics.jacobian([f], vp) .- Symbolics.jacobian([f], p)))
-end
+end
+
+@testset "Derivatives of indexed array expressions" begin
+    @register_array_symbolic foobar(x, y) begin
+        size = (2,)
+        eltype = Real
+        ndims = 1
+    end
+    @register_array_symbolic dfoobar1(x, y) begin
+        size = (2,)
+        eltype = Real
+        ndims = 1
+    end
+    @register_array_symbolic dfoobar2(x, y) begin
+        size = (2,)
+        eltype = Real
+        ndims = 1
+    end
+    Symbolics.derivative(::typeof(foobar), args::NTuple{2, Any}, ::Val{1}) = dfoobar1(args...)
+    Symbolics.derivative(::typeof(foobar), args::NTuple{2, Any}, ::Val{2}) = dfoobar2(args...)
+    
+    @variables x y
+    ex = foobar(x + 2y, y)
+    @test ex isa Symbolics.Arr
+    @test size(ex) == (2,)
+
+    der = Symbolics.derivative(ex[1], x)
+    @test isequal(der, dfoobar1(x + 2y, y)[1])
+    
+    der = Symbolics.derivative(ex[1], y)
+    @test isequal(der, dfoobar1(x + 2y, y)[1] * 2 + dfoobar2(x + 2y, y)[1])
+end