Added kron and swap to Algodiff operations

[tachukao]

The main changes in this PR are:

1. can now do Ndarray transpose in Algodiff, by exposing the ?axis argument
2. use transpose to implement swap in Algodiff
3. implement kron product using transpose and matrix product in Algodiff (implementation thanks to @ghennequin )

[jzstark]

So I follow the implementation of kron to build a Ndarray version of it:

module N = Dense.Ndarray.S
module M = Dense.Matrix.S

let kron a b =
  let na, ma =
    let s = N.shape a in
    s.(0), s.(1)
  in
  let nb, mb =
    let s = N.shape b in
    s.(0), s.(1)
  in
  let a = N.reshape a [| -1; 1 |] in
  let b = N.reshape a [| 1; -1 |] in
  let c = M.dot a b in
  let c = N.reshape c [| na; ma; nb; mb |] in
  let c = N.transpose ~axis:[| 0; 2; 1; 3 |] c in
  N.reshape c [| Stdlib.(na * nb); Stdlib.(ma * mb) |]

The only difference is that I replace *@ with M.dot. Then I try the test in unit_dense_matrix.ml:

let a = N.sequential [|2;2|];;
let b = Dense.Matrix.S.eye 2;;

Looks like that kron and M.kron give different results:

kron b a ;;

- : N.arr =

   C0 C1 C2 C3
R0  1  0  0  0
R1  0  1  0  0
R2  0  0  1  0
R3  0  0  0  1

M.kron b a 

- : Dense.Matrix.S.mat =

   C0 C1 C2 C3
R0  0  0  1  0
R1  0  0  0  1
R2  2  0  3  0
R3  0  2  0  3

Perhaps need to check the result in Algodiff module to see if I got anything wrong in using the AD implementation to build the Ndarray version kron.

[tachukao]

thanks @jzstark! the latest commit should fix this bug

[jzstark]

Cool! Thanks for the quick fix. The rest looks good to me.