From 9c2bfc9863baa3112d657bc7e7ceaddfc67a3a22 Mon Sep 17 00:00:00 2001
From: Alex Black <blacka101@gmail.com>
Date: Thu, 22 Aug 2019 16:16:03 +1000
Subject: [PATCH] Various fixes (DL4J, ND4J) (#147)

* Import fixes, IsMax dtype calc, small test fix

Signed-off-by: AlexDBlack <blacka101@gmail.com>

* SubsamplingLayer fix

Signed-off-by: AlexDBlack <blacka101@gmail.com>

* DL4J - SpaceToBatch layer updates

Signed-off-by: AlexDBlack <blacka101@gmail.com>
---
 .../gradientcheck/CNNGradientCheckTest.java   | 12 +++++-----
 .../nn/layers/convolution/SpaceToBatch.java   | 22 ++++++++++++-------
 .../subsampling/SubsamplingLayer.java         |  2 +-
 .../converters/ImportClassMapping.java        | 13 ++++++-----
 .../api/ops/impl/transforms/any/IsMax.java    |  6 +++++
 .../opvalidation/TransformOpValidation.java   |  2 +-
 6 files changed, 36 insertions(+), 21 deletions(-)

diff --git a/deeplearning4j/deeplearning4j-core/src/test/java/org/deeplearning4j/gradientcheck/CNNGradientCheckTest.java b/deeplearning4j/deeplearning4j-core/src/test/java/org/deeplearning4j/gradientcheck/CNNGradientCheckTest.java
index 89293829a..c1a873dc8 100644
--- a/deeplearning4j/deeplearning4j-core/src/test/java/org/deeplearning4j/gradientcheck/CNNGradientCheckTest.java
+++ b/deeplearning4j/deeplearning4j-core/src/test/java/org/deeplearning4j/gradientcheck/CNNGradientCheckTest.java
@@ -298,7 +298,7 @@ public class CNNGradientCheckTest extends BaseDL4JTest {
         int inputDepth = 1;
 
         int[] kernel = {2, 2};
-        int[] blocks = {1, 1};
+        int[] blocks = {2, 2};
 
         String[] activations = {"sigmoid", "tanh"};
         SubsamplingLayer.PoolingType[] poolingTypes =
@@ -309,8 +309,8 @@ public class CNNGradientCheckTest extends BaseDL4JTest {
             for (SubsamplingLayer.PoolingType poolingType : poolingTypes) {
                 for (int minibatchSize : minibatchSizes) {
                     INDArray input = Nd4j.rand(minibatchSize, width * height * inputDepth);
-                    INDArray labels = Nd4j.zeros(minibatchSize, nOut);
-                    for (int i = 0; i < minibatchSize; i++) {
+                    INDArray labels = Nd4j.zeros(4 * minibatchSize, nOut);
+                    for (int i = 0; i < 4 * minibatchSize; i++) {
                         labels.putScalar(new int[]{i, i % nOut}, 1.0);
                     }
 
@@ -318,11 +318,11 @@ public class CNNGradientCheckTest extends BaseDL4JTest {
                             new NeuralNetConfiguration.Builder()
                                     .dataType(DataType.DOUBLE)
                                     .updater(new NoOp()).weightInit(new NormalDistribution(0, 1))
-                                    .list().layer(new ConvolutionLayer.Builder(kernel).nIn(inputDepth)
-                                    .nOut(3).build())//output: (5-2+0)/1+1 = 4
+                                    .list()
+                                    .layer(new ConvolutionLayer.Builder(kernel).nIn(inputDepth).nOut(3).build())
                                     .layer(new SpaceToBatchLayer.Builder(blocks).build()) //trivial space to batch
                                     .layer(new OutputLayer.Builder(LossFunctions.LossFunction.MCXENT)
-                                            .activation(Activation.SOFTMAX).nIn(4 * 4 * 3)
+                                            .activation(Activation.SOFTMAX)
                                             .nOut(nOut).build())
                                     .setInputType(InputType.convolutionalFlat(height, width, inputDepth))
                                     .build();
diff --git a/deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/layers/convolution/SpaceToBatch.java b/deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/layers/convolution/SpaceToBatch.java
index c0770161d..eef8fcee7 100644
--- a/deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/layers/convolution/SpaceToBatch.java
+++ b/deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/layers/convolution/SpaceToBatch.java
@@ -100,12 +100,15 @@ public class SpaceToBatch extends AbstractLayer<org.deeplearning4j.nn.conf.layer
 
         Gradient gradient = new DefaultGradient();
 
-        CustomOp op = DynamicCustomOp.builder("batch_to_space")
-                .addInputs(epsilon, getBlocksArray(), getPaddingArray())
-                .addOutputs(outEpsilon)
+        INDArray epsilonNHWC = epsilon.permute(0, 2, 3, 1);
+        INDArray outEpsilonNHWC = outEpsilon.permute(0, 2, 3, 1);
+
+        CustomOp op = DynamicCustomOp.builder("batch_to_space_nd")
+                .addInputs(epsilonNHWC, getBlocksArray(), getPaddingArray())
+                .addOutputs(outEpsilonNHWC)
                 .callInplace(false)
                 .build();
-        Nd4j.getExecutioner().exec(op);
+        Nd4j.exec(op);
 
         outEpsilon = backpropDropOutIfPresent(outEpsilon);
         return new Pair<>(gradient, outEpsilon);
@@ -143,11 +146,14 @@ public class SpaceToBatch extends AbstractLayer<org.deeplearning4j.nn.conf.layer
 
         INDArray out = workspaceMgr.create(ArrayType.ACTIVATIONS, input.dataType(), new long[]{outMiniBatch, depth, outH, outW}, 'c');
 
-        CustomOp op = DynamicCustomOp.builder("space_to_batch")
-                .addInputs(input, getBlocksArray(), getPaddingArray())
-                .addOutputs(out)
+        INDArray inNHWC = input.permute(0, 2, 3, 1);
+        INDArray outNHWC = out.permute(0, 2, 3, 1);
+
+        CustomOp op = DynamicCustomOp.builder("space_to_batch_nd")
+                .addInputs(inNHWC, getBlocksArray(), getPaddingArray())
+                .addOutputs(outNHWC)
                 .build();
-        Nd4j.getExecutioner().exec(op);
+        Nd4j.exec(op);
 
         return out;
     }
diff --git a/deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/layers/convolution/subsampling/SubsamplingLayer.java b/deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/layers/convolution/subsampling/SubsamplingLayer.java
index 9eb4adb82..494d9c374 100644
--- a/deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/layers/convolution/subsampling/SubsamplingLayer.java
+++ b/deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/layers/convolution/subsampling/SubsamplingLayer.java
@@ -172,7 +172,7 @@ public class SubsamplingLayer extends AbstractLayer<org.deeplearning4j.nn.conf.l
                 b = DynamicCustomOp.builder("maxpool2d_bp");
                 break;
             case AVG:
-                b = DynamicCustomOp.builder("maxpool2d_bp");
+                b = DynamicCustomOp.builder("avgpool2d_bp");
                 if(layerConf().isAvgPoolIncludePadInDivisor()){
                     //Mostly this is a legacy case - beta4 and earlier models.
                     extra = 1;    //Divide by "number present" excluding padding
diff --git a/nd4j/nd4j-backends/nd4j-api-parent/nd4j-api/src/main/java/org/nd4j/imports/converters/ImportClassMapping.java b/nd4j/nd4j-backends/nd4j-api-parent/nd4j-api/src/main/java/org/nd4j/imports/converters/ImportClassMapping.java
index e13c45294..92c721474 100644
--- a/nd4j/nd4j-backends/nd4j-api-parent/nd4j-api/src/main/java/org/nd4j/imports/converters/ImportClassMapping.java
+++ b/nd4j/nd4j-backends/nd4j-api-parent/nd4j-api/src/main/java/org/nd4j/imports/converters/ImportClassMapping.java
@@ -350,9 +350,12 @@ public class ImportClassMapping {
             org.nd4j.linalg.api.ops.impl.transforms.custom.ATan2.class,
             org.nd4j.linalg.api.ops.impl.transforms.custom.Assign.class,
             org.nd4j.linalg.api.ops.impl.transforms.custom.BatchToSpace.class,
+            org.nd4j.linalg.api.ops.impl.transforms.custom.BatchToSpaceND.class,
             org.nd4j.linalg.api.ops.impl.transforms.custom.Choose.class,
             org.nd4j.linalg.api.ops.impl.transforms.custom.CumProd.class,
             org.nd4j.linalg.api.ops.impl.transforms.custom.CumSum.class,
+            org.nd4j.linalg.api.ops.impl.transforms.custom.CyclicShiftBits.class,
+            org.nd4j.linalg.api.ops.impl.transforms.custom.CyclicRShiftBits.class,
             org.nd4j.linalg.api.ops.impl.transforms.custom.Dilation2D.class,
             org.nd4j.linalg.api.ops.impl.transforms.custom.DotProductAttention.class,
             org.nd4j.linalg.api.ops.impl.transforms.custom.DotProductAttentionBp.class,
@@ -396,8 +399,11 @@ public class ImportClassMapping {
             org.nd4j.linalg.api.ops.impl.transforms.custom.Reverse.class,
             org.nd4j.linalg.api.ops.impl.transforms.custom.ReverseSequence.class,
             org.nd4j.linalg.api.ops.impl.transforms.custom.ReverseV2.class,
+            org.nd4j.linalg.api.ops.impl.transforms.custom.RShiftBits.class,
+            org.nd4j.linalg.api.ops.impl.transforms.custom.ShiftBits.class,
             org.nd4j.linalg.api.ops.impl.transforms.custom.SoftMax.class,
             org.nd4j.linalg.api.ops.impl.transforms.custom.SpaceToBatch.class,
+            org.nd4j.linalg.api.ops.impl.transforms.custom.SpaceToBatchND.class,
             org.nd4j.linalg.api.ops.impl.transforms.custom.Standardize.class,
             org.nd4j.linalg.api.ops.impl.transforms.custom.StandardizeBp.class,
             org.nd4j.linalg.api.ops.impl.transforms.custom.Svd.class,
@@ -457,6 +463,7 @@ public class ImportClassMapping {
             org.nd4j.linalg.api.ops.impl.transforms.pairwise.arithmetic.bp.DivBpOp.class,
             org.nd4j.linalg.api.ops.impl.transforms.pairwise.arithmetic.bp.FloorDivBpOp.class,
             org.nd4j.linalg.api.ops.impl.transforms.pairwise.arithmetic.bp.FloorModBpOp.class,
+            org.nd4j.linalg.api.ops.impl.transforms.pairwise.arithmetic.bp.ModBpOp.class,
             org.nd4j.linalg.api.ops.impl.transforms.pairwise.arithmetic.bp.MulBpOp.class,
             org.nd4j.linalg.api.ops.impl.transforms.pairwise.arithmetic.bp.RDivBpOp.class,
             org.nd4j.linalg.api.ops.impl.transforms.pairwise.arithmetic.bp.RSubBpOp.class,
@@ -559,11 +566,7 @@ public class ImportClassMapping {
             org.nd4j.linalg.api.ops.random.impl.ProbablisticMerge.class,
             org.nd4j.linalg.api.ops.random.impl.Range.class,
             org.nd4j.linalg.api.ops.random.impl.TruncatedNormalDistribution.class,
-            org.nd4j.linalg.api.ops.random.impl.UniformDistribution.class,
-            org.nd4j.linalg.api.ops.impl.transforms.custom.ShiftBits.class,
-            org.nd4j.linalg.api.ops.impl.transforms.custom.RShiftBits.class,
-            org.nd4j.linalg.api.ops.impl.transforms.custom.CyclicShiftBits.class,
-            org.nd4j.linalg.api.ops.impl.transforms.custom.CyclicRShiftBits.class
+            org.nd4j.linalg.api.ops.random.impl.UniformDistribution.class
 
     );
 
diff --git a/nd4j/nd4j-backends/nd4j-api-parent/nd4j-api/src/main/java/org/nd4j/linalg/api/ops/impl/transforms/any/IsMax.java b/nd4j/nd4j-backends/nd4j-api-parent/nd4j-api/src/main/java/org/nd4j/linalg/api/ops/impl/transforms/any/IsMax.java
index c06fe8026..4a72349d3 100644
--- a/nd4j/nd4j-backends/nd4j-api-parent/nd4j-api/src/main/java/org/nd4j/linalg/api/ops/impl/transforms/any/IsMax.java
+++ b/nd4j/nd4j-backends/nd4j-api-parent/nd4j-api/src/main/java/org/nd4j/linalg/api/ops/impl/transforms/any/IsMax.java
@@ -79,4 +79,10 @@ public class IsMax extends DynamicCustomOp {
     public List<SDVariable> doDiff(List<SDVariable> f1) {
         return Collections.singletonList(f().zerosLike(arg()));
     }
+
+    @Override
+    public List<DataType> calculateOutputDataTypes(List<DataType> inputDataTypes){
+        //Also supports other types if say float array is provided as output array
+        return Collections.singletonList(DataType.BOOL);
+    }
 }
diff --git a/nd4j/nd4j-backends/nd4j-tests/src/test/java/org/nd4j/autodiff/opvalidation/TransformOpValidation.java b/nd4j/nd4j-backends/nd4j-tests/src/test/java/org/nd4j/autodiff/opvalidation/TransformOpValidation.java
index c266aa4bd..0d177027d 100644
--- a/nd4j/nd4j-backends/nd4j-tests/src/test/java/org/nd4j/autodiff/opvalidation/TransformOpValidation.java
+++ b/nd4j/nd4j-backends/nd4j-tests/src/test/java/org/nd4j/autodiff/opvalidation/TransformOpValidation.java
@@ -1249,7 +1249,7 @@ public class TransformOpValidation extends BaseOpValidation {
                 case 2:
                     //TODO: IsMax supports both bool and float out: https://github.com/deeplearning4j/deeplearning4j/issues/6872
                     inArr = Nd4j.create(new double[]{-3,5,0,2});
-                    exp = Nd4j.create(new boolean[]{false,true,false,false}).castTo(DataType.DOUBLE);
+                    exp = Nd4j.create(new boolean[]{false,true,false,false});
                     out = sd.math().isMax(in);
                     break;
                 case 3: