Fix missing non-optional argument in rrules with backend. (#168)

* Add non-optional arguments in chainrules implementation

`tensorcontract(args..., alfa, backend)` only has a default value for `alfa` when no backend is supplied.

* Add tests for rrules with backend specified

Author: lkdvos

ext/TensorOperationsChainRulesCoreExt.jl

@@ -55,13 +55,13 @@ function ChainRulesCore.rrule(::typeof(TensorOperations.tensoradd!),
             _dα = tensorscalar(tensorcontract(((), ()), A, ((), linearize(pC)),
                                               _conj(conjA), ΔC,
                                               (trivtuple(numind(pC)),
-                                               ()), :N, backend...))
+                                               ()), :N, One(), backend...))
             return projectα(_dα)
         end
         dβ = @thunk begin
             _dβ = tensorscalar(tensorcontract(((), ()), C,
                                               ((), trivtuple(numind(pC))), :C, ΔC,
-                                              (trivtuple(numind(pC)), ()), :N,
+                                              (trivtuple(numind(pC)), ()), :N, One(),
                                               backend...))
             return projectβ(_dβ)
         end
@@ -116,17 +116,17 @@ function ChainRulesCore.rrule(::typeof(TensorOperations.tensorcontract!),
         dα = @thunk begin
             _dα = tensorscalar(tensorcontract(((), ()),
                                               tensorcontract(pC, A, pA, conjA, B, pB,
-                                                             conjB),
+                                                             conjB, One(), backend...),
                                               ((), trivtuple(numind(pC))),
                                               :C, ΔC,
-                                              (trivtuple(numind(pC)), ()), :N,
+                                              (trivtuple(numind(pC)), ()), :N, One(),
                                               backend...))
             return projectα(_dα)
         end
         dβ = @thunk begin
             _dβ = tensorscalar(tensorcontract(((), ()), C,
                                               ((), trivtuple(numind(pC))), :C, ΔC,
-                                              (trivtuple(numind(pC)), ()), :N,
+                                              (trivtuple(numind(pC)), ()), :N, One(),
                                               backend...))
             return projectβ(_dβ)
         end
@@ -172,14 +172,14 @@ function ChainRulesCore.rrule(::typeof(tensortrace!), C, pC::Index2Tuple, A,
                                               tensortrace(pC, A, pA),
                                               ((), trivtuple(numind(pC))),
                                               _conj(conjA), ΔC,
-                                              (trivtuple(numind(pC)), ()), :N,
+                                              (trivtuple(numind(pC)), ()), :N, One(),
                                               backend...))
             return projectα(_dα)
         end
         dβ = @thunk begin
             _dβ = tensorscalar(tensorcontract(((), ()), C,
                                               ((), trivtuple(numind(pC))), :C, ΔC,
-                                              (trivtuple(numind(pC)), ()), :N,
+                                              (trivtuple(numind(pC)), ()), :N, One(),
                                               backend...))
             return projectβ(_dβ)
         end

test/ad.jl

@@ -1,4 +1,5 @@
 using TensorOperations
+using TensorOperations: StridedBLAS, StridedNative
 using Test
 using ChainRulesTestUtils
 
@@ -19,7 +20,12 @@ precision(::Type{<:Union{Float64,Complex{Float64}}}) = 1e-8
     β = rand(T)
     A = rand(T₁, (2, 3, 4, 2, 5))
     C = rand(T₂, size.(Ref(A), pC[1]))
+
     test_rrule(tensortrace!, C, pC, A, pA, :N, α, β; atol, rtol)
+    test_rrule(tensortrace!, C, pC, A, pA, :C, α, β; atol, rtol)
+
+    test_rrule(tensortrace!, C, pC, A, pA, :C, α, β, StridedBLAS(); atol, rtol)
+    test_rrule(tensortrace!, C, pC, A, pA, :N, α, β, StridedNative(); atol, rtol)
 end
 
 @testset "tensoradd! ($T₁, $T₂)" for (T₁, T₂) in ((Float64, Float64), (Float32, Float64),
@@ -35,6 +41,9 @@ end
     β = rand(T)
     test_rrule(tensoradd!, C, pC, A, :N, α, β; atol, rtol)
     test_rrule(tensoradd!, C, pC, A, :C, α, β; atol, rtol)
+
+    test_rrule(tensoradd!, C, pC, A, :N, α, β, StridedBLAS(); atol, rtol)
+    test_rrule(tensoradd!, C, pC, A, :C, α, β, StridedNative(); atol, rtol)
 end
 
 @testset "tensorcontract! ($T₁, $T₂)" for (T₁, T₂) in
@@ -58,6 +67,11 @@ end
     test_rrule(tensorcontract!, C, pC, A, pA, :C, B, pB, :N, α, β; atol, rtol)
     test_rrule(tensorcontract!, C, pC, A, pA, :N, B, pB, :C, α, β; atol, rtol)
     test_rrule(tensorcontract!, C, pC, A, pA, :C, B, pB, :C, α, β; atol, rtol)
+
+    test_rrule(tensorcontract!, C, pC, A, pA, :N, B, pB, :N, α, β, StridedBLAS();
+               atol, rtol)
+    test_rrule(tensorcontract!, C, pC, A, pA, :C, B, pB, :N, α, β, StridedNative();
+               atol, rtol)
 end
 
 @testset "tensorscalar ($T)" for T in (Float32, Float64, ComplexF64)