Remove index replacement for powers

Author: asterycs

src/forward.jl

@@ -44,10 +44,8 @@ function diff(arg::UnaryOperation{Cos}, wrt::Variable)
 end
 
 function diff(arg::Power, wrt::Variable)
-    outer = replace_bound_letters(arg.base, wrt)
-
     return BinaryOperation{Mult}(
-        BinaryOperation{Mult}(arg.exponent, Power(outer, arg.exponent - 1)),
+        BinaryOperation{Mult}(arg.exponent, Power(arg.base, arg.exponent - 1)),
         diff(arg.base, wrt),
     )
 end
@@ -63,18 +61,6 @@ function diff(arg::BinaryOperation{Op}, wrt::Variable) where {Op<:AdditiveOperat
     return BinaryOperation{Op}(diff(arg.arg1, wrt), diff(arg.arg2, wrt))
 end
 
-function replace_bound_letters(arg::Tensor, letters_to_skip::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, letters_to_skip...)
-
-    letter_map =
-        Dict(bound_letters[li] => next_letter + li - 1 for li ∈ eachindex(bound_letters))
-
-    return replace_letters(arg, letter_map)
-end
-
 function collect_factors(arg::BinaryOperation{Mult})
     return Value[collect_factors(arg.arg1); collect_factors(arg.arg2)]
 end

test/StdStrTest.jl

@@ -30,8 +30,8 @@
     @test to_std(gradient(sum(2 * cos(A * x + y)), x)) == "(-2)Aᵀsin(Ax + y)"
     @test to_std(gradient(sum(x .^ 2), x)) == "2x"
     @test to_std(gradient(sum(x .^ 3), x)) == "3x^2"
-    @test to_std(gradient(sum(x)^2, x)) == "2sum(xᵀ)vec(1)"
-    @test to_std(gradient(sum(x .^ 2)^2, x)) == "2sum(xᵀ^2)2x"
+    @test to_std(gradient(sum(x)^2, x)) == "2sum(xᵀ)Iᵀvec(1)" # TODO: Simplify and remove Iᵀ
+    @test to_std(gradient(sum(x .^ 2)^2, x)) == "4sum(xᵀ^2)x"
     @test to_std(gradient(sum((x + y) .^ 2), x)) == "2(x + y)"
     @test to_std(gradient(sum((x .* y) .^ 2), x)) == "2(x ⊙ y ⊙ y)"
     @test to_std(gradient(sum((A * x - y) .^ 2), x)) == "2Aᵀ(Ax - y)"