fix transpose case and add explicit Expand testing

DiffSharp · gbaydin · Oct 26, 2020 · Feb 26, 2020 · Feb 28, 2020 · Feb 28, 2020
commit 171d66289bf3642b65c9086045e448be97f789ef
diff --git a/src/DiffSharp.Backend.None/RawTensorCPU.fs b/src/DiffSharp.Backend.None/RawTensorCPU.fs
@@ -196,7 +196,7 @@ type RawTensorFloat32CPU(values: float32[], shape:int[]) =
         Array.init n (fun i -> Array.init unstackedLength (fun j -> t.Values.[i*unstackedLength+j]))
         |> Array.map (fun v -> upcast RawTensorFloat32CPU(v, unstackedShape))
 
-    override t.TransposeT2() =
+    override t.TransposeT() =
         if t.Dim < 2 then failwith "Expecting at least a 2D tensor"
         let oldShape = t.Shape
         let batch = oldShape.[0..oldShape.Length-3]
@@ -207,7 +207,7 @@ type RawTensorFloat32CPU(values: float32[], shape:int[]) =
         for i = 0 to values.Length-1 do
             let col = i % ncols 
             let row = (i / ncols ) % nrows
-            let j = (i / ncols / nrows) + col*nrows + row
+            let j = (i / ncols / nrows)*ncols*nrows + col*nrows + row
             result.[j] <- values.[i]
         upcast RawTensorFloat32CPU(result, newShape)
 

diff --git a/src/DiffSharp.Core/RawTensor.fs b/src/DiffSharp.Core/RawTensor.fs
@@ -171,7 +171,7 @@ and [<AbstractClass>]
     abstract member NegT : unit -> RawTensor
     abstract member SumT : unit -> RawTensor
     abstract member SumT2Dim0 : unit -> RawTensor
-    abstract member TransposeT2: unit -> RawTensor
+    abstract member TransposeT: unit -> RawTensor
     abstract member SqueezeT: int -> RawTensor
     abstract member UnsqueezeT: int -> RawTensor
     abstract member FlipT: int[] -> RawTensor

diff --git a/src/DiffSharp.Core/Tensor.fs b/src/DiffSharp.Core/Tensor.fs
@@ -577,11 +577,11 @@ type Tensor =
     member t.SumT2Dim0() = Tensor.SumT2Dim0(t)
 
     static member Transpose (a:Tensor) =
-        if a.Dim > 2 then failwithf "Expecting at least a 2d tensor, received Tensor with shape %A" a.Shape
-        let inline fRaw(a:RawTensor) = a.TransposeT2()
+        if a.Dim < 2 then failwithf "Expecting at least a 2d tensor, received Tensor with shape %A" a.Shape
+        let inline fRaw(a:RawTensor) = a.TransposeT()
         let inline fTensor(a) = Tensor.Transpose(a)
         let inline dfTensorFwd(cp,ap,ad) = Tensor.Transpose(ad)
-        let inline dfTensorRev(a) = TransposeT2(a)
+        let inline dfTensorRev(a) = TransposeT(a)
         Tensor.OpUnary(a, fRaw, fTensor, dfTensorFwd, dfTensorRev)
     member t.Transpose() = Tensor.Transpose(t)
 
@@ -964,7 +964,7 @@ type Tensor =
                         | ExpandT(a) -> reset (a::tt)
                         | StackTs(a) -> reset (List.append (a |> List.ofSeq) tt)
                         | UnstackT(a,_) -> reset (a::tt)
-                        | TransposeT2(a) -> reset (a::tt)
+                        | TransposeT(a) -> reset (a::tt)
                         | SqueezeT(a) -> reset (a::tt)
                         | UnsqueezeT(a) -> reset (a::tt)
                         | FlipT(a,_) -> reset (a::tt)
@@ -1245,7 +1245,7 @@ type Tensor =
                             if a.Derivative.Dim = 0 then a.Derivative <- Tensor.ZerosLike(a) + a.Derivative
                             a.Derivative <- Tensor.AddSlice(a.Derivative, Array.init a.Dim (fun j -> if j=0 then i else 0), t.Derivative.Unsqueeze(0))
                             push ((a.Zero(), a) :: tt)
-                        | TransposeT2(a) -> push ((
8000
t.Derivative.Transpose(), a) :: tt)
+                        | TransposeT(a) -> push ((t.Derivative.Transpose(), a) :: tt)
                         | SqueezeT(a) -> push ((t.Derivative.ViewAs(a), a) :: tt)
                         | UnsqueezeT(a) -> push ((t.Derivative.ViewAs(a), a) :: tt)
                         | FlipT(a, dims) -> push ((t.Derivative.Flip(dims), a) :: tt)
@@ -1353,7 +1353,7 @@ and TensorOp =
     | AddTTSlice of Tensor * int[] * Tensor
     | AddTTConstSlice of Tensor
     | AddTConstTSlice of int[] * Tensor
-    | TransposeT2 of Tensor
+    | TransposeT of Tensor
     | SqueezeT of Tensor
     | UnsqueezeT of Tensor
     | FlipT of Tensor * int[]

diff --git a/src/DiffSharp.Tests/TestDerivatives.fs b/src/DiffSharp.Tests/TestDerivatives.fs
@@ -70,7 +70,39 @@ type TestDerivatives () =
     // TODO: add test for AddT2ConstT1
 
     [<Test>]
-    member this.TestDerivativeAddWithBroadcast () =
+    member this.TestDerivativeExpand () =
+
+        let t1 = Tensor.Create([[1.]; [2.]]).ForwardDiff(Tensor.Create([[5.]; [6.]])) // 2x1
+        let t1Expand = t1.Expand([2;2;2]) // 2x2x2 = [[[1.;1]; [2.;2]]; [[1.;1]; [2.;2]]]
+        let fwdz = t1Expand
+        let fwdzd = fwdz.Derivative
+        let fwdzdCorrect = Tensor.Create ([[[5., 5.], [6., 6.]], [[5., 5.], [6., 6.]]])
+
+        (* Python:
+        import torch 
+        t1 = torch.tensor([[1.], [2.]], requires_grad=True)
+        revz = t1.expand([2,2,2])
+        revz.backward(torch.tensor([[[3.,3.], [6.,6.]], [[3.,3.], [6.,6.]]]))
+        t1.grad
+        --> tensor([[12.],[24.]])
+        *)
+        let revy = t1.ReverseDiff()
+        let revz = revy.Expand([2;2;2])
+        let revz_grad = Tensor.Create([[[3.;3.]; [6.;6.]]; [[3.;3.]; [6.;6.]]])
+        revz.Reverse(revz_grad)
+        let revyd = revy.Derivative
+        // Note: The 4x'3' accumulate to the first entry, the 4x'6' accumulate to the second entry
+        let revydCorrect = Tensor.Create [[12.], [24.]]
+        Assert.AreEqual(fwdzd,fwdzdCorrect)
+        Assert.AreEqual(revyd,revydCorrect)
+
+    [<Test>]
+    member this.TestAddWithBroadcastSystematic () =
+
+        // This is a somewhat adhoc extra test to do a whole range of additiosn
+        // with broadcast, mainly to check that not problems occur in taking the
+        // derivatives.
+        //
         // Systematically do all allowed broadcasts into 2x3x4
         // 2x3x4 + 1  (broadcast --> 2x3x4)
         // 2x3x4 + 4  (broadcast --> 2x3x4)
@@ -81,7 +113,7 @@ type TestDerivatives () =
         let t1a = Tensor.Create([ [ [1.; 2.; 3.; 4.]; [5.; 6.; 7.; 8.]; [9.; 10.; 11.; 12.] ];
                                   [ [13.; 14.; 15.; 16.]; [17.; 18.; 19.; 20.]; [21.; 22.; 23.; 24.] ]  ])
 
-        // Get all the interesting shapes that broadcast into t1a
+        // Get all the interesting shapes that expand into t1a
         let shapes = 
             [ for i1 in [0;1;2] do
                 for i2 in [0;1;3] do
@@ -2228,15 +2260,40 @@ type TestDerivatives () =
         let fwdzd = fwdz.Derivative
         let fwdzdCorrect = Tensor.Create([[2.; 10.]; [3.; 20.]; [4.; 30.]])
 
+        Assert.AreEqual(fwdzCorrect, fwdz)
+        Assert.AreEqual(fwdzdCorrect, fwdzd)
+
         let revx = Tensor.Create([[1.; 2.; 3.]; [4.; 5.; 6.]]).ReverseDiff()
         let revz = revx.Transpose()
         let revzCorrect = Tensor.Create([[1.; 4.]; [2.; 5.]; [3.; 6.]])
         revz.Reverse(Tensor.Create([[5.; 5.]; [2.; 5.]; [3.; 7.]]))
         let revxd = revx.Derivative
         let revxdCorrect = Tensor.Create([[5.; 2.; 3.]; [5.; 5.; 7.]])
 
+        Assert.AreEqual(revzCorrect, revz)
+        Assert.AreEqual(revxdCorrect, revxd)
+
+    [<Test>]
+    member this.TestDerivativeTransposeBatch () =
+        // This test is the same as TestDerivativeTransposeT2 except we add a batching expansion to
+        // both input and expected results
+        let t = Tensor.Create([[1.; 2.; 3.]; [4.; 5.; 6.]]).Expand([| 3;2;3 |])
+        let fwdx = t.ForwardDiff(Tensor.Create([[2.; 3.; 4.]; [10.; 20.; 30.]]).Expand([| 3;2;3 |]))
+        let fwdz = fwdx.Transpose()
+        let fwdzCorrect = Tensor.Create([[1.; 4.]; [2.; 5.]; [3.; 6.]]).Expand([| 3;3;2 |])
+        let fwdzd = fwdz.Derivative
+        let fwdzdCorrect = Tensor.Create([[2.; 10.]; [3.; 20.]; [4.; 30.]]).Expand([| 3;3;2 |])
+
         Assert.AreEqual(fwdzCorrect, fwdz)
         Assert.AreEqual(fwdzdCorrect, fwdzd)
+
+        let revx = t.ReverseDiff()
+        let revz = revx.Transpose()
+        let revzCorrect = Tensor.Create([[1.; 4.]; [2.; 5.]; [3.; 6.]]).Expand([| 3;3;2 |])
+        revz.Reverse(Tensor.Create([[5.; 5.]; [2.; 5.]; [3.; 7.]]).Expand([| 3;3;2 |]))
+        let revxd = revx.Derivative
+        let revxdCorrect = Tensor.Create([[5.; 2.; 3.]; [5.; 5.; 7.]]).Expand([| 3;2;3 |])
+
         Assert.AreEqual(revzCorrect, revz)
         Assert.AreEqual(revxdCorrect, revxd)
 

diff --git a/src/DiffSharp.Tests/TestTensor.fs b/src/DiffSharp.Tests/TestTensor.fs
@@ -1255,6 +1255,19 @@ type TestTensor () =
         Assert.AreEqual(t1TransposeCorrect, t1Transpose)
         Assert.AreEqual(t2TransposeTransposeCorrect, t2TransposeTranspose)
 
+    [<Test>]
+    member this.TestTensorTransposeBatch () =
+        let t1 = Tensor.Create([[1.; 2.; 3.]; [4.; 5.; 6.]]).Expand([|3;2;3|])
+        let t1Transpose = t1.Transpose()
+        let t1TransposeCorrect = Tensor.Create([[1.; 4.]; [2.; 5.]; [3.; 6.]]).Expand([|3;3;2|])
+
+        let t2 = Tensor.Create([[1.; 2.]; [3.; 4.]]).Expand([|3;2;2|])
+        let t2TransposeTranspose = t2.Transpose().Transpose()
+        let t2TransposeTransposeCorrect = t2
+
+        Assert.AreEqual(t1TransposeCorrect, t1Transpose)
+        Assert.AreEqual(t2TransposeTransposeCorrect, t2TransposeTranspose)
+
     [<Test>]
     member this.TestTensorSignT () =
         let t1 = Tensor.Create([-1.; -2.; 0.; 3.])