Replace bound letters when diff:ing Powers

Author: asterycs

src/forward.jl

@@ -40,9 +40,11 @@ function diff(arg::UnaryOperation{Cos}, wrt::Monomial)
 end
 
 function diff(arg::Power, wrt::Monomial)
+    inner = replace_bound_letters(arg.base)
+
     return BinaryOperation{Mult}(
         BinaryOperation{Mult}(arg.exponent, Power(arg.base, arg.exponent - 1)),
-        diff(arg.base, wrt),
+        diff(inner, wrt),
     )
 end
 
@@ -57,6 +59,18 @@ function diff(arg::BinaryOperation{Op}, wrt::Monomial) where {Op<:AdditiveOperat
     return BinaryOperation{Op}(diff(arg.arg1, wrt), diff(arg.arg2, wrt))
 end
 
+function replace_bound_letters(arg::Tensor)
+    letters = unique(get_letters(get_indices(arg)))
+    free_letters = unique(get_letters(get_free_indices(arg)))
+    bound_letters = setdiff(letters, free_letters)
+    next_letter = get_next_letter(arg)
+
+    letter_map =
+        Dict(bound_letters[li] => next_letter + li for li ∈ eachindex(bound_letters))
+
+    return replace_letters(arg, letter_map)
+end
+
 function collect_factors(arg::BinaryOperation{Mult})
     return [collect_factors(arg.arg1); collect_factors(arg.arg2)]
 end

test/StdTest.jl

@@ -351,6 +351,7 @@ end
     @test to_std_string(gradient(sum(2 * cos(A * x + y)), x)) == "(-2)Aᵀsin(Ax + y)"
     @test to_std_string(gradient(sum(x .^ 2), x)) == "2x"
     @test to_std_string(gradient(sum(x .^ 3), x)) == "3x²"
+    @test to_std_string(gradient(sum(x)^2, x)) == "2sum(xᵀ)vec(1)"
     @test to_std_string(gradient(sum((x + y) .^ 2), x)) == "2(x + y)"
     @test to_std_string(gradient(sum((x .* y) .^ 2), x)) == "2x ⊙ y ⊙ y"
     @test to_std_string(gradient(sum((A * x - y) .^ 2), x)) == "2Aᵀ(Ax - y)"