DL4J/ND4J: Do pass on integer casts (#15)

* Int cast fixes.

* Revert "Int cast fixes."

This reverts commit aa36e8ca

* Int casts

* Int cast

* Int casts

* Get rid of int casts. Dropping deprecated aggregate ops.

* java scatterUpdate changes

Signed-off-by: raver119 <raver119@gmail.com>

* c++ scatterUpdate changes

Signed-off-by: raver119 <raver119@gmail.com>

* Remove aggregated ops.

* Restored test

* Tests restored.

* Minor fixes

arbiter/arbiter-deeplearning4j/src/test/java/org/deeplearning4j/arbiter/TestUtils.java

@@ -124,7 +124,6 @@ public class TestUtils {
     public static INDArray randomOneHot(long examples, long nOut, Random rng){
         INDArray arr = Nd4j.create(examples, nOut);
         for( int i=0; i<examples; i++ ){
-            // FIXME: int cast
             arr.putScalar(i, rng.nextInt((int) nOut), 1.0);
         }
         return arr;

deeplearning4j/deeplearning4j-core/src/test/java/org/deeplearning4j/TestUtils.java

@@ -132,7 +132,6 @@ public class TestUtils {
     public static INDArray randomOneHot(long examples, long nOut, Random rng){
         INDArray arr = Nd4j.create(examples, nOut);
         for( int i=0; i<examples; i++ ){
-            // FIXME: int cast
             arr.putScalar(i, rng.nextInt((int) nOut), 1.0);
         }
         return arr;

deeplearning4j/deeplearning4j-core/src/test/java/org/deeplearning4j/datasets/iterator/AsyncDataSetIteratorTest.java

@@ -187,7 +187,6 @@ public class AsyncDataSetIteratorTest extends BaseDL4JTest {
         }
     }
 
-
     @Test
     public void testVariableTimeSeries2() throws Exception {
         AsyncDataSetIterator adsi =

deeplearning4j/deeplearning4j-core/src/test/java/org/deeplearning4j/gradientcheck/LossFunctionGradientCheck.java

@@ -464,13 +464,11 @@ public class LossFunctionGradientCheck extends BaseDL4JTest {
                 ret[1] = Nd4j.zeros(labelsShape);
                 if (labelsShape.length == 2) {
                     for (int i = 0; i < labelsShape[0]; i++) {
-                        // FIXME: int cast
                         ret[1].putScalar(i, r.nextInt((int) labelsShape[1]), 1.0);
                     }
                 } else if (labelsShape.length == 3) {
                     for (int i = 0; i < labelsShape[0]; i++) {
                         for (int j = 0; j < labelsShape[2]; j++) {
-                            // FIXME: int cast
                             ret[1].putScalar(i, r.nextInt((int) labelsShape[1]), j, 1.0);
                         }
                     }
@@ -484,13 +482,11 @@ public class LossFunctionGradientCheck extends BaseDL4JTest {
                 ret[1] = Nd4j.ones(labelsShape);
                 if (labelsShape.length == 2) {
                     for (int i = 0; i < labelsShape[0]; i++) {
-                        // FIXME: int cast
                         ret[1].putScalar(i, r.nextInt((int) labelsShape[1]), -1.0);
                     }
                 } else if (labelsShape.length == 3) {
                     for (int i = 0; i < labelsShape[0]; i++) {
                         for (int j = 0; j < labelsShape[2]; j++) {
-                            // FIXME: int cast
                             ret[1].putScalar(i, r.nextInt((int) labelsShape[1]), j, -1.0);
                         }
                     }

deeplearning4j/deeplearning4j-core/src/test/java/org/deeplearning4j/nn/conf/graph/ShiftVertexTest.java

@@ -176,8 +176,7 @@ public class ShiftVertexTest extends BaseDL4JTest {
         manual_weights.put("output_b", c);
 
         // First things first, let's calculate the score.
-        // FIXME: int cast
+        long batchsz = input.shape()[0];
-        int batchsz = (int) input.shape()[0];
         INDArray z = input.castTo(W.dataType()).mmul(W).add(b.repmat(batchsz, 1));
         INDArray a = a1.getActivation(z.dup(), true).add(sf); // activation modifies it's input!!
         INDArray q = a.mmul(V).add(c.repmat(batchsz, 1));

deeplearning4j/deeplearning4j-core/src/test/java/org/deeplearning4j/nn/multilayer/BackPropMLPTest.java

@@ -36,6 +36,7 @@ import org.nd4j.linalg.api.ops.impl.transforms.strict.SigmoidDerivative;
 import org.nd4j.linalg.api.ops.impl.transforms.strict.TanhDerivative;
 import org.nd4j.linalg.dataset.DataSet;
 import org.nd4j.linalg.dataset.api.iterator.DataSetIterator;
+import org.nd4j.linalg.exception.ND4JArraySizeException;
 import org.nd4j.linalg.factory.Nd4j;
 import org.nd4j.linalg.learning.config.Sgd;
 import org.nd4j.linalg.lossfunctions.LossFunctions.LossFunction;
@@ -340,7 +341,8 @@ public class BackPropMLPTest extends BaseDL4JTest {
 
     public static float[] asFloat(INDArray arr) {
         long len = arr.length();
-        // FIXME: int cast
+        if (len > Integer.MAX_VALUE)
+            throw new ND4JArraySizeException();
         float[] f = new float[(int) len];
         NdIndexIterator iterator = new NdIndexIterator('c', arr.shape());
         for (int i = 0; i < len; i++) {

deeplearning4j/deeplearning4j-core/src/test/java/org/deeplearning4j/nn/multilayer/MultiLayerTest.java

@@ -320,7 +320,6 @@ public class MultiLayerTest extends BaseDL4JTest {
     public static float[] asFloat(INDArray arr) {
         long len = arr.length();
 
-        // FIXME: int cast
         float[] f = new float[(int) len];
         for (int i = 0; i < len; i++)
             f[i] = arr.getFloat(i);

deeplearning4j/deeplearning4j-core/src/test/java/org/deeplearning4j/nn/updater/TestUpdaters.java

@@ -331,7 +331,6 @@ public class TestUpdaters extends BaseDL4JTest {
         double calculatedByHandMScalar = 0.2;
         double[] expectedM = Nd4j.ones(1, numParams).mul(calculatedByHandMScalar).data().asDouble();
 
-        // FIXME: int cast
         double[] actualM = Arrays.copyOfRange(nadamUpdater.getM().data().asDouble(), 0, (int) numParams);
         for (int i = 0; i < actualM.length; i++) {
             actualM[i] = Math.round(actualM[i] * 1e2) / 1e2;

deeplearning4j/deeplearning4j-core/src/test/java/org/deeplearning4j/optimize/solver/TestOptimizers.java

@@ -48,6 +48,7 @@ import org.nd4j.linalg.api.rng.DefaultRandom;
 import org.nd4j.linalg.api.rng.Random;
 import org.nd4j.linalg.dataset.DataSet;
 import org.nd4j.linalg.dataset.api.iterator.DataSetIterator;
+import org.nd4j.linalg.exception.ND4JArraySizeException;
 import org.nd4j.linalg.factory.Nd4j;
 import org.nd4j.linalg.indexing.conditions.Condition;
 import org.nd4j.linalg.learning.config.AdaGrad;
@@ -664,7 +665,9 @@ public class TestOptimizers extends BaseDL4JTest {
             double xlm1 = parameters.getDouble(nDims - 2);
             double gl = 200 * (xl - xlm1 * xlm1);
 
-            // FIXME: int cast
+            if (nDims - 1 > Integer.MAX_VALUE) {
+                throw new ND4JArraySizeException();
+            }
             gradient.put(0,  (int)nDims - 1, gl);
             Gradient g = new DefaultGradient();
             g.gradientForVariable().put("W", gradient);
@@ -865,8 +868,7 @@ public class TestOptimizers extends BaseDL4JTest {
 
         @Override
         public long numParams() {
-            // FIXME: int cast
+            return parameters.length();
-            return (int) parameters.length();
         }
 
         @Override

deeplearning4j/deeplearning4j-data/deeplearning4j-datavec-iterators/src/main/java/org/deeplearning4j/datasets/datavec/RecordReaderMultiDataSetIterator.java

@@ -286,7 +286,7 @@ public class RecordReaderMultiDataSetIterator implements MultiDataSetIterator, S
                 for (INDArray w : exampleData) {
                     val n = w.size(0);
 
-                    // FIXME: int cast
+                    if (Math.min(minExamples, n) < Integer.MAX_VALUE)
                         minExamples = (int) Math.min(minExamples, n);
                 }
             }

deeplearning4j/deeplearning4j-data/deeplearning4j-datavec-iterators/src/main/java/org/deeplearning4j/datasets/datavec/SequenceRecordReaderDataSetIterator.java

@@ -366,7 +366,6 @@ public class SequenceRecordReaderDataSetIterator implements DataSetIterator {
         DataSet ds = mdsToDataSet(mds);
 
         if (totalOutcomes == -1) {
-            // FIXME: int cast
             inputColumns = (int) ds.getFeatures().size(1);
             totalOutcomes = ds.getLabels() == null ? -1 : (int) ds.getLabels().size(1);
         }
@@ -394,7 +393,6 @@ public class SequenceRecordReaderDataSetIterator implements DataSetIterator {
         stored = next();
         useStored = true;
 
-        // FIXME: int cast
         inputColumns = (int) stored.getFeatures().size(1);
         totalOutcomes = (int) stored.getLabels().size(1);
     }

deeplearning4j/deeplearning4j-data/deeplearning4j-utility-iterators/src/main/java/org/deeplearning4j/datasets/iterator/AbstractDataSetIterator.java

@@ -172,7 +172,6 @@ public abstract class AbstractDataSetIterator<T> implements DataSetIterator {
                     Pair<T, T> pair = iterator.next();
                     if (numFeatures < 1) {
                         if (pair.getFirst() instanceof INDArray) {
-                            // FIXME: int cast
                             numFeatures = (int) ((INDArray) pair.getFirst()).length();
                             numLabels = (int) ((INDArray) pair.getSecond()).length();
                         } else if (pair.getFirst() instanceof float[]) {

deeplearning4j/deeplearning4j-data/deeplearning4j-utility-iterators/src/main/java/org/deeplearning4j/datasets/iterator/IteratorDataSetIterator.java

@@ -95,7 +95,6 @@ public class IteratorDataSetIterator implements DataSetIterator {
             //Set columns etc for later use
             DataSet temp = list.get(0);
 
-            // FIXME: int cast
             inputColumns = (int) temp.getFeatures().size(1);
             totalOutcomes = temp.getLabels() == null ? 0 : (int) temp.getLabels().size(1); //May be null for layerwise pretraining
         }

deeplearning4j/deeplearning4j-data/deeplearning4j-utility-iterators/src/main/java/org/deeplearning4j/datasets/iterator/IteratorMultiDataSetIterator.java

@@ -73,8 +73,7 @@ public class IteratorMultiDataSetIterator implements MultiDataSetIterator {
                 next = iterator.next();
             }
 
-            // FIXME: int cast
+            long nExamples = next.getFeatures(0).size(0);
-            int nExamples = (int) next.getFeatures(0).size(0);
             if (countSoFar + nExamples <= batchSize) {
                 //Add the entire MultiDataSet as-is
                 list.add(next);
@@ -140,7 +139,7 @@ public class IteratorMultiDataSetIterator implements MultiDataSetIterator {
         return out;
     }
 
-    private static INDArray getRange(INDArray arr, int exampleFrom, int exampleToExclusive) {
+    private static INDArray getRange(INDArray arr, long exampleFrom, long exampleToExclusive) {
         if (arr == null)
             return null;
 

deeplearning4j/deeplearning4j-data/deeplearning4j-utility-iterators/src/main/java/org/deeplearning4j/datasets/iterator/file/BaseFileIterator.java

@@ -134,7 +134,7 @@ public abstract class BaseFileIterator<T, P> implements Iterator<T> {
         List<T> remainder = new ArrayList<>();
         int soFar = 0;
         for (T t : toMerge) {
-            int size = sizeOf(t);
+            long size = sizeOf(t);
 
             if (soFar + size <= batchSize) {
                 correctNum.add(t);
@@ -190,7 +190,7 @@ public abstract class BaseFileIterator<T, P> implements Iterator<T> {
 
     protected abstract T load(File f);
 
-    protected abstract int sizeOf(T of);
+    protected abstract long sizeOf(T of);
 
     protected abstract List<T> split(T toSplit);
 

deeplearning4j/deeplearning4j-data/deeplearning4j-utility-iterators/src/main/java/org/deeplearning4j/datasets/iterator/file/FileDataSetIterator.java

@@ -151,7 +151,7 @@ public class FileDataSetIterator extends BaseFileIterator<DataSet, DataSetPrePro
     }
 
     @Override
-    protected int sizeOf(DataSet of) {
+    protected long sizeOf(DataSet of) {
         return of.numExamples();
     }
 

deeplearning4j/deeplearning4j-data/deeplearning4j-utility-iterators/src/main/java/org/deeplearning4j/datasets/iterator/file/FileMultiDataSetIterator.java

@@ -151,9 +151,8 @@ public class FileMultiDataSetIterator extends BaseFileIterator<MultiDataSet, Mul
     }
 
     @Override
-    protected int sizeOf(MultiDataSet of) {
+    protected long sizeOf(MultiDataSet of) {
-        // FIXME: int cast
+        return  of.getFeatures(0).size(0);
-        return  (int) of.getFeatures(0).size(0);
     }
 
     @Override

deeplearning4j/deeplearning4j-manifold/deeplearning4j-tsne/src/main/java/org/deeplearning4j/plot/BarnesHutTsne.java

@@ -665,8 +665,7 @@ public class BarnesHutTsne implements Model {
 
             if (useAdaGrad) {
                 if (adaGrad == null) {
-                    // FIXME: int cast
+                    adaGrad = new AdaGrad(gradient.shape(), learningRate);
-                    adaGrad = new AdaGrad(ArrayUtil.toInts(gradient.shape()), learningRate);
                     adaGrad.setStateViewArray(Nd4j.zeros(gradient.shape()).reshape(1, gradChange.length()),
                             gradChange.shape(), gradient.ordering(), true);
                 }

deeplearning4j/deeplearning4j-modelimport/src/main/java/org/deeplearning4j/nn/modelimport/keras/layers/core/KerasReshape.java

@@ -18,6 +18,7 @@ package org.deeplearning4j.nn.modelimport.keras.layers.core;
 
 
 import lombok.val;
+import org.apache.commons.lang3.ArrayUtils;
 import org.deeplearning4j.nn.conf.InputPreProcessor;
 import org.deeplearning4j.nn.conf.inputs.InputType;
 import org.deeplearning4j.nn.modelimport.keras.KerasLayer;
@@ -51,6 +52,13 @@ public class KerasReshape extends KerasLayer {
         this(layerConfig, true);
     }
 
+    private long[] listToLongArray(List<Integer> list) {
+        long[] retVal = new long[list.size()];
+        for (int i = 0; i < list.size(); ++i) {
+            retVal[i] = list.get(i);
+        }
+        return retVal;
+    }
     /**
      * Constructor from parsed Keras layer configuration dictionary.
      *
@@ -67,9 +75,7 @@ public class KerasReshape extends KerasLayer {
         if (innerConfig.containsKey(targetShape)) {
             @SuppressWarnings("unchecked")
             List<Integer> targetShapeList = (List<Integer>) innerConfig.get(targetShape);
-
+            this.targetShape = listToLongArray(targetShapeList);
-            // FIXME: int cast
-            this.targetShape = ArrayUtil.toLongArray(ArrayUtil.toArray(targetShapeList));
         }
 
     }

deeplearning4j/deeplearning4j-modelimport/src/test/java/org/deeplearning4j/nn/modelimport/keras/e2e/KerasModelEndToEndTest.java

@@ -690,13 +690,11 @@ public class KerasModelEndToEndTest extends BaseDL4JTest {
                 INDArray testLabels = Nd4j.create(predictionsDl4j.shape());
                 if (testLabels.rank() == 2) {
                     for (int i = 0; i < testLabels.size(0); i++) {
-                        // FIXME: int cast
                         testLabels.putScalar(i, r.nextInt((int) testLabels.size(1)), 1.0);
                     }
                 } else if (testLabels.rank() == 3) {
                     for (int i = 0; i < testLabels.size(0); i++) {
                         for (int j = 0; j < testLabels.size(1); j++) {
-                            // FIXME: int cast
                             testLabels.putScalar(i, j, r.nextInt((int) testLabels.size(1)), 1.0);
                         }
                     }

deeplearning4j/deeplearning4j-nlp-parent/deeplearning4j-nlp/src/main/java/org/deeplearning4j/models/embeddings/inmemory/InMemoryLookupTable.java

@@ -104,7 +104,7 @@ public class InMemoryLookupTable<T extends SequenceElement> implements WeightLoo
     }
 
     protected void initAdaGrad() {
-        int[] shape = new int[] {vocab.numWords() + 1, vectorLength};
+        long[] shape = new long[] {vocab.numWords() + 1, vectorLength};
         int length = ArrayUtil.prod(shape);
         adaGrad = new AdaGrad(shape, lr.get());
         adaGrad.setStateViewArray(Nd4j.zeros(shape).reshape(1, length), shape, Nd4j.order(), true);
@@ -124,8 +124,7 @@ public class InMemoryLookupTable<T extends SequenceElement> implements WeightLoo
         if (adaGrad == null)
             initAdaGrad();
 
-        // FIXME: int cast
+        return adaGrad.getGradient(gradient, column, syn0.shape());
-        return adaGrad.getGradient(gradient, column, ArrayUtil.toInts(syn0.shape()));
     }
 
     @Override
@@ -370,7 +369,6 @@ public class InMemoryLookupTable<T extends SequenceElement> implements WeightLoo
                 else {
                     nextRandom.set(nextRandom.get() * 25214903917L + 11);
 
-                    // FIXME: int cast
                     int idx = (int) Math.abs((int) (nextRandom.get() >> 16) % table.length());
 
                     target = table.getInt(idx);

deeplearning4j/deeplearning4j-nlp-parent/deeplearning4j-nlp/src/main/java/org/deeplearning4j/models/embeddings/learning/impl/elements/CBOW.java

@@ -33,7 +33,6 @@ import org.deeplearning4j.models.word2vec.wordstore.VocabCache;
 import org.nd4j.linalg.api.buffer.DataType;
 import org.nd4j.linalg.api.ndarray.INDArray;
 import org.nd4j.linalg.api.ops.aggregates.Aggregate;
-import org.nd4j.linalg.api.ops.aggregates.impl.AggregateCBOW;
 import org.nd4j.linalg.api.ops.impl.nlp.CbowRound;
 import org.nd4j.linalg.factory.Nd4j;
 import org.nd4j.linalg.util.DeviceLocalNDArray;

deeplearning4j/deeplearning4j-nlp-parent/deeplearning4j-nlp/src/main/java/org/deeplearning4j/models/embeddings/learning/impl/elements/GloVe.java

@@ -104,11 +104,10 @@ public class GloVe<T extends SequenceElement> implements ElementsLearningAlgorit
 
 
 
-        weightAdaGrad = new AdaGrad(new int[] {this.vocabCache.numWords() + 1, vectorLength}, learningRate);
+        weightAdaGrad = new AdaGrad(new long[] {this.vocabCache.numWords() + 1, vectorLength}, learningRate);
         bias = Nd4j.create(syn0.rows());
 
-        // FIXME: int cast
+        biasAdaGrad = new AdaGrad(bias.shape(), this.learningRate);
-        biasAdaGrad = new AdaGrad(ArrayUtil.toInts(bias.shape()), this.learningRate);
 
         //  maxmemory = Runtime.getRuntime().maxMemory() - (vocabCache.numWords() * vectorLength * 2 * 8);
 
@@ -237,15 +236,13 @@ public class GloVe<T extends SequenceElement> implements ElementsLearningAlgorit
     private void update(T element1, INDArray wordVector, INDArray contextVector, double gradient) {
         //gradient for word vectors
         INDArray grad1 = contextVector.mul(gradient);
-        // FIXME: int cast
+        INDArray update = weightAdaGrad.getGradient(grad1, element1.getIndex(), syn0.shape());
-        INDArray update = weightAdaGrad.getGradient(grad1, element1.getIndex(), ArrayUtil.toInts(syn0.shape()));
 
         //update vector
         wordVector.subi(update);
 
         double w1Bias = bias.getDouble(element1.getIndex());
-        // FIXME: int cast
+        double biasGradient = biasAdaGrad.getGradient(gradient, element1.getIndex(), bias.shape());
-        double biasGradient = biasAdaGrad.getGradient(gradient, element1.getIndex(), ArrayUtil.toInts(bias.shape()));
         double update2 = w1Bias - biasGradient;
         bias.putScalar(element1.getIndex(), update2);
     }

deeplearning4j/deeplearning4j-nlp-parent/deeplearning4j-nlp/src/main/java/org/deeplearning4j/models/embeddings/reader/impl/BasicModelUtils.java

@@ -358,7 +358,6 @@ public class BasicModelUtils<T extends SequenceElement> implements ModelUtils<T>
             INDArray sort = sorted[0];
             List<String> ret = new ArrayList<>();
 
-            // FIXME: int cast
             if (top > sort.length())
                 top = (int) sort.length();
             //there will be a redundant word

deeplearning4j/deeplearning4j-nlp-parent/deeplearning4j-nlp/src/main/java/org/deeplearning4j/models/glove/GloveWeightLookupTable.java

@@ -72,7 +72,7 @@ public class GloveWeightLookupTable<T extends SequenceElement> extends InMemoryL
             putVector(Word2Vec.DEFAULT_UNK, randUnk);
         }
         if (weightAdaGrad == null || reset) {
-            weightAdaGrad = new AdaGrad(new int[] {vocab.numWords() + 1, vectorLength}, lr.get());
+            weightAdaGrad = new AdaGrad(new long[]{vocab.numWords() + 1, vectorLength}, lr.get());
         }
 
 
@@ -81,7 +81,7 @@ public class GloveWeightLookupTable<T extends SequenceElement> extends InMemoryL
             bias = Nd4j.create(syn0.rows());
 
         if (biasAdaGrad == null || reset) {
-            biasAdaGrad = new AdaGrad(ArrayUtil.toInts(bias.shape()), lr.get());
+            biasAdaGrad = new AdaGrad(bias.shape(), lr.get());
         }
 
 
@@ -140,13 +140,13 @@ public class GloveWeightLookupTable<T extends SequenceElement> extends InMemoryL
     private void update(T w1, INDArray wordVector, INDArray contextVector, double gradient) {
         //gradient for word vectors
         INDArray grad1 = contextVector.mul(gradient);
-        INDArray update = weightAdaGrad.getGradient(grad1, w1.getIndex(), ArrayUtil.toInts(syn0.shape()));
+        INDArray update = weightAdaGrad.getGradient(grad1, w1.getIndex(), syn0.shape());
 
         //update vector
         wordVector.subi(update);
 
         double w1Bias = bias.getDouble(w1.getIndex());
-        double biasGradient = biasAdaGrad.getGradient(gradient, w1.getIndex(), ArrayUtil.toInts(bias.shape()));
+        double biasGradient = biasAdaGrad.getGradient(gradient, w1.getIndex(), bias.shape());
         double update2 = w1Bias - biasGradient;
         bias.putScalar(w1.getIndex(), update2);
     }

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/gradientcheck/GradientCheckUtil.java

@@ -20,6 +20,7 @@ import lombok.extern.slf4j.Slf4j;
 import lombok.val;
 import org.deeplearning4j.nn.api.Model;
 import org.nd4j.linalg.api.buffer.DataType;
+import org.nd4j.linalg.exception.ND4JArraySizeException;
 import org.nd4j.linalg.function.Consumer;
 import org.nd4j.linalg.lossfunctions.impl.LossBinaryXENT;
 import org.nd4j.linalg.primitives.Pair;
@@ -293,7 +294,8 @@ public class GradientCheckUtil {
                     ss = n;
                 }
 
-                // FIXME: int cast
+                if (ss > Integer.MAX_VALUE)
+                    throw new ND4JArraySizeException();
                 stepSizeForParam.put(paramNames.get(i), (int) ss);
             }
         }

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/conf/graph/ElementWiseVertex.java

@@ -140,10 +140,9 @@ public class ElementWiseVertex extends GraphVertex {
             //CNN inputs... also check that the channels, width and heights match:
             InputType.InputTypeConvolutional firstConv = (InputType.InputTypeConvolutional) first;
 
-            // FIXME: int cast
+            val fd = firstConv.getChannels();
-            val fd = (int) firstConv.getChannels();
+            val fw = firstConv.getWidth();
-            val fw = (int) firstConv.getWidth();
+            val fh = firstConv.getHeight();
-            val fh = (int) firstConv.getHeight();
 
             for (int i = 1; i < vertexInputs.length; i++) {
                 if (vertexInputs[i].getType() != InputType.Type.CNN) {
@@ -155,10 +154,9 @@ public class ElementWiseVertex extends GraphVertex {
 
                 InputType.InputTypeConvolutional otherConv = (InputType.InputTypeConvolutional) vertexInputs[i];
 
-                // FIXME: int cast
+                val od = otherConv.getChannels();
-                val od = (int) otherConv.getChannels();
+                val ow = otherConv.getWidth();
-                val ow = (int) otherConv.getWidth();
+                val oh = otherConv.getHeight();
-                val oh = (int) otherConv.getHeight();
 
                 if (fd != od || fw != ow || fh != oh) {
                     throw new InvalidInputTypeException(

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/conf/graph/MergeVertex.java

@@ -94,13 +94,12 @@ public class MergeVertex extends GraphVertex {
             // CNN3D inputs: check that the channels, width and height match:
             InputType.InputTypeConvolutional3D firstConv = (InputType.InputTypeConvolutional3D) first;
 
-            // FIXME: int cast
+            val fd = firstConv.getDepth();
-            val fd = (int) firstConv.getDepth();
+            val fw = firstConv.getWidth();
-            val fw = (int) firstConv.getWidth();
+            val fh = firstConv.getHeight();
-            val fh = (int) firstConv.getHeight();
+            val fc = firstConv.getChannels();
-            val fc = (int) firstConv.getChannels();
 
-            int depthSum = fc;
+            long depthSum = fc;
             InputType.InputTypeConvolutional3D otherConv = null;
             for (int i = 1; i < vertexInputs.length; i++) {
                 if (vertexInputs[i].getType() != InputType.Type.CNN3D) {
@@ -109,10 +108,10 @@ public class MergeVertex extends GraphVertex {
                 }
 
                 otherConv = (InputType.InputTypeConvolutional3D) vertexInputs[i];
-                val od = (int) otherConv.getDepth();
+                val od = otherConv.getDepth();
-                val ow = (int) otherConv.getWidth();
+                val ow = otherConv.getWidth();
-                val oh = (int) otherConv.getHeight();
+                val oh = otherConv.getHeight();
-                val oc = (int) otherConv.getChannels();
+                val oc = otherConv.getChannels();
 
                 if (fd != od || fw != ow || fh != oh) {
                     throw new InvalidInputTypeException("Invalid input: MergeVertex cannot merge CNN3D activations of different width/heights:" + "first [channels,width,height] = [" + fd + "," + fw + "," + fh
@@ -177,12 +176,11 @@ public class MergeVertex extends GraphVertex {
             //CNN inputs... also check that the channels, width and heights match:
             InputType.InputTypeConvolutional firstConv = (InputType.InputTypeConvolutional) first;
 
-            // FIXME: int cast
+            val fd = firstConv.getChannels();
-            val fd = (int) firstConv.getChannels();
+            val fw = firstConv.getWidth();
-            val fw = (int) firstConv.getWidth();
+            val fh = firstConv.getHeight();
-            val fh = (int) firstConv.getHeight();
 
-            int depthSum = fd;
+            long depthSum = fd;
 
             for (int i = 1; i < vertexInputs.length; i++) {
                 if (vertexInputs[i].getType() != InputType.Type.CNN) {
@@ -194,10 +192,9 @@ public class MergeVertex extends GraphVertex {
 
                 InputType.InputTypeConvolutional otherConv = (InputType.InputTypeConvolutional) vertexInputs[i];
 
-                // FIXME: int cast
+                val od = otherConv.getChannels();
-                val od = (int) otherConv.getChannels();
+                val ow = otherConv.getWidth();
-                val ow = (int) otherConv.getWidth();
+                val oh = otherConv.getHeight();
-                val oh = (int) otherConv.getHeight();
 
                 if (fw != ow || fh != oh) {
                     throw new InvalidInputTypeException(

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/conf/graph/PoolHelperVertex.java

@@ -131,12 +131,11 @@ public class PoolHelperVertex extends GraphVertex {
             //CNN inputs... also check that the channels, width and heights match:
             InputType.InputTypeConvolutional firstConv = (InputType.InputTypeConvolutional) first;
 
-            // FIXME: int cast
+            val fd = firstConv.getChannels();
-            val fd = (int) firstConv.getChannels();
+            val fw = firstConv.getWidth();
-            val fw = (int) firstConv.getWidth();
+            val fh = firstConv.getHeight();
-            val fh = (int) firstConv.getHeight();
 
-            int depthSum = fd;
+            long depthSum = fd;
 
             for (int i = 1; i < vertexInputs.length; i++) {
                 if (vertexInputs[i].getType() != InputType.Type.CNN) {
@@ -148,10 +147,9 @@ public class PoolHelperVertex extends GraphVertex {
 
                 InputType.InputTypeConvolutional otherConv = (InputType.InputTypeConvolutional) vertexInputs[i];
 
-                // FIXME: int cast
+                long od = otherConv.getChannels();
-                int od = (int) otherConv.getChannels();
+                long ow = otherConv.getWidth();
-                int ow = (int) otherConv.getWidth();
+                long oh = otherConv.getHeight();
-                int oh = (int) otherConv.getHeight();
 
                 if (fw != ow || fh != oh) {
                     throw new InvalidInputTypeException(

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/conf/graph/UnstackVertex.java

@@ -150,12 +150,11 @@ public class UnstackVertex extends GraphVertex {
             //CNN inputs... also check that the channels, width and heights match:
             InputType.InputTypeConvolutional firstConv = (InputType.InputTypeConvolutional) first;
 
-            // FIXME: int cast
+            val fd = firstConv.getChannels();
-            val fd = (int) firstConv.getChannels();
+            val fw = firstConv.getWidth();
-            val fw = (int) firstConv.getWidth();
+            val fh = firstConv.getHeight();
-            val fh = (int) firstConv.getHeight();
 
-            int depthSum = fd;
+            long depthSum = fd;
 
             for (int i = 1; i < vertexInputs.length; i++) {
                 if (vertexInputs[i].getType() != InputType.Type.CNN) {
@@ -167,10 +166,9 @@ public class UnstackVertex extends GraphVertex {
 
                 InputType.InputTypeConvolutional otherConv = (InputType.InputTypeConvolutional) vertexInputs[i];
 
-                // FIXME: int cast
+                val od = otherConv.getChannels();
-                val od = (int) otherConv.getChannels();
+                val ow = otherConv.getWidth();
-                val ow = (int) otherConv.getWidth();
+                val oh = otherConv.getHeight();
-                val oh = (int) otherConv.getHeight();
 
                 if (fw != ow || fh != oh) {
                     throw new InvalidInputTypeException(

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/conf/inputs/InputType.java

@@ -402,18 +402,17 @@ public abstract class InputType implements Serializable {
         //Note: ConvolutionalFlat and FeedForward look identical... but either should work OK if using something
         // like FeedForwardToCnnPreProcessor
 
-        // FIXME: int cast
         switch (inputArray.rank()) {
             case 2:
-                return InputType.feedForward((int) inputArray.size(1));
+                return InputType.feedForward(inputArray.size(1));
             case 3:
-                return InputType.recurrent((int) inputArray.size(1), (int) inputArray.size(2));
+                return InputType.recurrent(inputArray.size(1), (int) inputArray.size(2));
             case 4:
                 //Order: [minibatch, channels, height, width] -> [h, w, c]
-                return InputType.convolutional((int) inputArray.size(2), (int) inputArray.size(3), (int) inputArray.size(1));
+                return InputType.convolutional(inputArray.size(2), (int) inputArray.size(3), (int) inputArray.size(1));
             case 5:
                 //Order: [minibatch, channels, depth, height, width] -> [d, h, w, c]
-                return InputType.convolutional3D((int) inputArray.size(2), (int) inputArray.size(3),
+                return InputType.convolutional3D(inputArray.size(2), (int) inputArray.size(3),
                         (int) inputArray.size(4), (int) inputArray.size(1));
             default:
                 throw new IllegalArgumentException(

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/conf/layers/Cnn3DLossLayer.java

@@ -152,17 +152,18 @@ public class Cnn3DLossLayer extends FeedForwardLayer {
         }
 
         @Override
-        public void setNIn(int nIn){
+        public void setNIn(long nIn){
             throw new UnsupportedOperationException(
                     "Cnn3DLossLayer has no parameters, thus nIn will always equal nOut.");
         }
 
         @Override
-        public void setNOut(int nOut){
+        public void setNOut(long nOut){
             throw new UnsupportedOperationException(
                     "Cnn3DLossLayer has no parameters, thus nIn will always equal nOut.");
         }
 
+
         @Override
         @SuppressWarnings("unchecked")
         public Cnn3DLossLayer build() {

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/conf/layers/CnnLossLayer.java

@@ -145,13 +145,13 @@ public class CnnLossLayer extends FeedForwardLayer {
         }
 
         @Override
-        public void setNIn(int nIn){
+        public void setNIn(long nIn){
             throw new UnsupportedOperationException(
                     "This layer has no parameters, thus nIn will always equal nOut.");
         }
 
         @Override
-        public void setNOut(int nOut){
+        public void setNOut(long nOut){
             throw new UnsupportedOperationException(
                     "This layer has no parameters, thus nIn will always equal nOut.");
         }

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/conf/layers/Convolution1DLayer.java

@@ -88,7 +88,7 @@ public class Convolution1DLayer extends ConvolutionLayer {
             //Probably: user did InputType.recurrent(x) without specifying sequence length
             outLength = -1;
         } else {
-            outLength = Convolution1DUtils.getOutputSize((int) inputTsLength, kernelSize[0], stride[0], padding[0],
+            outLength = Convolution1DUtils.getOutputSize(inputTsLength, kernelSize[0], stride[0], padding[0],
                             convolutionMode, dilation[0]);
         }
         return InputType.recurrent(nOut, outLength);

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/conf/layers/FeedForwardLayer.java

@@ -117,14 +117,14 @@ public abstract class FeedForwardLayer extends BaseLayer {
          * this is the input channels, otherwise is the previous layer size.
          *
          */
-        protected int nIn = 0;
+        protected long nIn = 0;
 
         /**
          * Number of inputs for the layer (usually the size of the last layer). <br> Note that for Convolutional layers,
          * this is the input channels, otherwise is the previous layer size.
          *
          */
-        protected int nOut = 0;
+        protected long nOut = 0;
 
         /**
          * Number of inputs for the layer (usually the size of the last layer). <br> Note that for Convolutional layers,
@@ -144,8 +144,7 @@ public abstract class FeedForwardLayer extends BaseLayer {
          * @param nIn Number of inputs for the layer
          */
         public T nIn(long nIn) {
-            // FIXME: int cast
+            this.setNIn(nIn);
-            this.setNIn((int) nIn);
             return (T) this;
         }
 

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/conf/layers/InputTypeUtil.java

@@ -41,12 +41,9 @@ public class InputTypeUtil {
                     Class<?> layerClass) {
         InputType.InputTypeConvolutional i = (InputType.InputTypeConvolutional) inputType;
 
-        // FIXME: int cast
+        val hIn = i.getHeight();
-        val hIn = (int) i.getHeight();
+        val wIn = i.getWidth();
-        val wIn = (int) i.getWidth();
 
-        val inHeight = (int) i.getHeight();
-        val inWidth = (int) i.getWidth();
         int padH = (padding == null ? 0 : padding[0]); //May be null for ConvolutionMode.Same
         int padW = (padding == null ? 0 : padding[1]);
         int kH = kernelSize[0];
@@ -69,13 +66,13 @@ public class InputTypeUtil {
         }
 
         if (convolutionMode == ConvolutionMode.Same) {
-            int hOut = stride[0] * hIn;
+            long hOut = stride[0] * hIn;
-            int wOut = stride[1] * wIn;
+            long wOut = stride[1] * wIn;
             return InputType.convolutional(hOut, wOut, outputDepth);
         }
 
-        int hOut = sH * (hIn - 1) + kH - 2 * padH;
+        long hOut = sH * (hIn - 1) + kH - 2 * padH;
-        int wOut = sW * (wIn - 1) + kW - 2 * padW;
+        long wOut = sW * (wIn - 1) + kW - 2 * padW;
 
         return InputType.convolutional(hOut, wOut, outputDepth);
     }
@@ -91,10 +88,9 @@ public class InputTypeUtil {
 
         InputType.InputTypeConvolutional3D i = (InputType.InputTypeConvolutional3D) inputType;
 
-        // FIXME: int cast
+        long inDepth = i.getDepth();
-        val inDepth = (int) i.getDepth();
+        long inHeight = i.getHeight();
-        val inHeight = (int) i.getHeight();
+        long inWidth = i.getWidth();
-        val inWidth = (int) i.getWidth();
 
         int padD = (padding == null ? 0 : padding[0]);
         int padH = (padding == null ? 0 : padding[1]);
@@ -211,9 +207,9 @@ public class InputTypeUtil {
             return InputType.convolutional3D(outD, outH, outW, outputChannels);
         }
 
-        int dOut = (inDepth - kD + 2 * padD) / sD + 1;
+        long dOut = (inDepth - kD + 2 * padD) / sD + 1;
-        int hOut = (inHeight - kH + 2 * padH) / sH + 1;
+        long hOut = (inHeight - kH + 2 * padH) / sH + 1;
-        int wOut = (inWidth - kW + 2 * padW) / sW + 1;
+        long wOut = (inWidth - kW + 2 * padW) / sW + 1;
         return InputType.convolutional3D(dOut, hOut, wOut, outputChannels);
     }
 
@@ -296,9 +292,8 @@ public class InputTypeUtil {
 
         InputType.InputTypeConvolutional i = (InputType.InputTypeConvolutional) inputType;
 
-        // FIXME: int cast
+        long inHeight = i.getHeight();
-        val inHeight = (int) i.getHeight();
+        long inWidth = i.getWidth();
-        val inWidth = (int) i.getWidth();
         int padH = (padding == null ? 0 : padding[0]); //May be null for ConvolutionMode.Same
         int padW = (padding == null ? 0 : padding[1]);
         int kH = kernelSize[0];
@@ -379,8 +374,8 @@ public class InputTypeUtil {
             return InputType.convolutional(outH, outW, outputDepth);
         }
 
-        int hOut = (inHeight - kH + 2 * padH) / sH + 1;
+        long hOut = (inHeight - kH + 2 * padH) / sH + 1;
-        int wOut = (inWidth - kW + 2 * padW) / sW + 1;
+        long wOut = (inWidth - kW + 2 * padW) / sW + 1;
         return InputType.convolutional(hOut, wOut, outputDepth);
     }
 

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/conf/layers/RnnLossLayer.java

@@ -142,13 +142,13 @@ public class RnnLossLayer extends FeedForwardLayer {
         }
 
         @Override
-        public void setNIn(int nIn){
+        public void setNIn(long nIn){
             throw new UnsupportedOperationException(
                     "This layer has no parameters, thus nIn will always equal nOut.");
         }
 
         @Override
-        public void setNOut(int nOut){
+        public void setNOut(long nOut){
             throw new UnsupportedOperationException(
                     "This layer has no parameters, thus nIn will always equal nOut.");
         }

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/conf/layers/Subsampling1DLayer.java

@@ -82,12 +82,12 @@ public class Subsampling1DLayer extends SubsamplingLayer {
         }
         InputType.InputTypeRecurrent r = (InputType.InputTypeRecurrent) inputType;
         long inputTsLength = r.getTimeSeriesLength();
-        int outLength;
+        long outLength;
         if (inputTsLength < 0) {
             //Probably: user did InputType.recurrent(x) without specifying sequence length
             outLength = -1;
         } else {
-            outLength = Convolution1DUtils.getOutputSize((int) inputTsLength, kernelSize[0], stride[0], padding[0],
+            outLength = Convolution1DUtils.getOutputSize(inputTsLength, kernelSize[0], stride[0], padding[0],
                             convolutionMode, dilation[0]);
         }
         return InputType.recurrent(r.getSize(), outLength);

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/conf/layers/Subsampling3DLayer.java

@@ -32,6 +32,7 @@ import org.deeplearning4j.util.ValidationUtils;
 import org.nd4j.base.Preconditions;
 import org.nd4j.linalg.api.buffer.DataType;
 import org.nd4j.linalg.api.ndarray.INDArray;
+import org.nd4j.linalg.exception.ND4JArraySizeException;
 import org.nd4j.linalg.learning.regularization.Regularization;
 
 import java.util.Collection;
@@ -138,9 +139,11 @@ public class Subsampling3DLayer extends NoParamLayer {
                             + "\"): Expected CNN input, got " + inputType);
         }
 
-        // FIXME: int cast
+        long inChannels = ((InputType.InputTypeConvolutional3D) inputType).getChannels();
+        if (inChannels > Integer.MAX_VALUE)
+            throw new ND4JArraySizeException();
         return InputTypeUtil.getOutputTypeCnn3DLayers(inputType, kernelSize, stride, padding, new int[] {1, 1, 1}, // no dilation
-                        convolutionMode, (int) ((InputType.InputTypeConvolutional3D) inputType).getChannels(),
+                        convolutionMode, (int) inChannels,
                         layerIndex, getLayerName(), Subsampling3DLayer.class);
     }
 

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/conf/layers/Upsampling3D.java

@@ -83,11 +83,10 @@ public class Upsampling3D extends BaseUpsamplingLayer {
         }
         InputType.InputTypeConvolutional3D i = (InputType.InputTypeConvolutional3D) inputType;
 
-        // FIXME: int cast
+        long inHeight = (int) i.getHeight();
-        int inHeight = (int) i.getHeight();
+        long inWidth = (int) i.getWidth();
-        int inWidth = (int) i.getWidth();
+        long inDepth = (int) i.getDepth();
-        int inDepth = (int) i.getDepth();
+        long inChannels = (int) i.getChannels();
-        int inChannels = (int) i.getChannels();
 
         return InputType.convolutional3D(size[0] * inDepth, size[1] * inHeight, size[2] * inWidth, inChannels);
     }

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/conf/ocnn/OCNNOutputLayer.java

@@ -259,7 +259,7 @@ public class OCNNOutputLayer extends BaseOutputLayer {
         }
 
         @Override
-        public void setNOut(int nOut){
+        public void setNOut(long nOut){
             throw new UnsupportedOperationException(
                     "Unable to specify number of outputs with ocnn. Outputs are fixed to 1.");
         }

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/graph/ComputationGraph.java

@@ -79,6 +79,7 @@ import org.nd4j.linalg.dataset.api.DataSetUtil;
 import org.nd4j.linalg.dataset.api.MultiDataSet;
 import org.nd4j.linalg.dataset.api.iterator.DataSetIterator;
 import org.nd4j.linalg.dataset.api.iterator.MultiDataSetIterator;
+import org.nd4j.linalg.exception.ND4JArraySizeException;
 import org.nd4j.linalg.factory.Nd4j;
 import org.nd4j.linalg.heartbeat.Heartbeat;
 import org.nd4j.linalg.heartbeat.reports.Environment;
@@ -3329,7 +3330,6 @@ public class ComputationGraph implements Serializable, Model, NeuralNetwork {
         //In 99+% of cases, the input and labels dimension 0 size should be identical
         //The only real exceptions: space to batch, and batch to space layers
         //In those cases, we should base it on the labels size, as this impacts gradient calculation
-        // FIXME: int cast
         return labels == null || labels[0] == null ? (int) inputs[0].size(0) : (int)labels[0].size(0);
     }
 
@@ -3653,7 +3653,8 @@ public class ComputationGraph implements Serializable, Model, NeuralNetwork {
             if (endTimeIdx > timeSeriesLength)
                 endTimeIdx = timeSeriesLength;
 
-            // FIXME: int cast
+            if (startTimeIdx > Integer.MAX_VALUE)
+                throw new ND4JArraySizeException();
             List<INDArray[]> list = getSubsetsForTbptt((int) startTimeIdx, endTimeIdx, inputs, labels, featureMasks, labelMasks);
 
             setInputs(list.get(0));
@@ -3799,7 +3800,8 @@ public class ComputationGraph implements Serializable, Model, NeuralNetwork {
                         }
                     }
 
-                    // FIXME: int cast
+                    if (minibatchSize > Integer.MAX_VALUE)
+                        throw new ND4JArraySizeException();
                     Pair<INDArray, MaskState> outPair =
                             current.feedForwardMaskArrays(inputMasks, maskState, (int)minibatchSize);
                     map.put(topologicalOrder[i], outPair);
@@ -4664,7 +4666,7 @@ public class ComputationGraph implements Serializable, Model, NeuralNetwork {
      * @param layer Index of the layer to get the size of. Must be in range 0 to nLayers-1 inclusive
      * @return Size of the layer
      */
-    public int layerSize(int layer) {
+    public long layerSize(int layer) {
         if (layer < 0 || layer > layers.length) {
             throw new IllegalArgumentException("Invalid layer index: " + layer + ". Layer index must be between 0 and "
                     + (layers.length - 1) + " inclusive");
@@ -4683,7 +4685,7 @@ public class ComputationGraph implements Serializable, Model, NeuralNetwork {
      * @param layer Index of the layer to get the size of. Must be in range 0 to nLayers-1 inclusive
      * @return Size of the layer
      */
-    public int layerInputSize(int layer) {
+    public long layerInputSize(int layer) {
         if (layer < 0 || layer > layers.length) {
             throw new IllegalArgumentException("Invalid layer index: " + layer + ". Layer index must be between 0 and "
                     + (layers.length - 1) + " inclusive");
@@ -4701,7 +4703,7 @@ public class ComputationGraph implements Serializable, Model, NeuralNetwork {
      * @param layerName Name of the layer to get the size of
      * @return Size of the layer
      */
-    public int layerSize(String layerName) {
+    public long layerSize(String layerName) {
         Layer l = getLayer(layerName);
         if(l == null){
             throw new IllegalArgumentException("No layer with name \"" + layerName + "\" exists");
@@ -4712,8 +4714,7 @@ public class ComputationGraph implements Serializable, Model, NeuralNetwork {
         }
         FeedForwardLayer ffl = (FeedForwardLayer) conf;
 
-        // FIXME: int cast
+        return ffl.getNOut();
-        return (int) ffl.getNOut();
     }
 
     /**
@@ -4727,7 +4728,7 @@ public class ComputationGraph implements Serializable, Model, NeuralNetwork {
      * @param layerName Name of the layer to get the size of
      * @return Size of the layer
      */
-    public int layerInputSize(String layerName) {
+    public long layerInputSize(String layerName) {
         Layer l = getLayer(layerName);
         if(l == null){
             throw new IllegalArgumentException("No layer with name \"" + layerName + "\" exists");
@@ -4738,8 +4739,7 @@ public class ComputationGraph implements Serializable, Model, NeuralNetwork {
         }
         FeedForwardLayer ffl = (FeedForwardLayer) conf;
 
-        // FIXME: int cast
+        return ffl.getNIn();
-        return (int) ffl.getNIn();
     }
 
     /**

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/graph/vertex/impl/UnstackVertex.java

@@ -43,10 +43,10 @@ import java.util.Arrays;
  * @author Justin Long (crockpotveggies)
  */
 public class UnstackVertex extends BaseGraphVertex {
-    private int from;
+    private long from;
     private int stackSize;
     private long forwardShape[];
-    private int step;
+    private long step;
 
     public UnstackVertex(ComputationGraph graph, String name, int vertexIndex, int from, int stackSize, DataType dataType) {
         this(graph, name, vertexIndex, null, null, from, stackSize, dataType);
@@ -77,10 +77,9 @@ public class UnstackVertex extends BaseGraphVertex {
         // once we know the inputs, save the shape and interval size for doBackward
         this.forwardShape = Arrays.copyOf(inputs[0].shape(), inputs[0].rank());
 
-        // FIXME: int cast
+        this.step = inputs[0].size(0) / stackSize;
-        this.step = (int) inputs[0].size(0) / stackSize;
+        long start = from * step;
-        int start = from * step;
+        long end = (from + 1) * step;
-        int end = (from + 1) * step;
 
         INDArray ret;
         switch (inputs[0].rank()) { //TODO remove the dups here if/when possible (gradient checks must pass)
@@ -108,8 +107,8 @@ public class UnstackVertex extends BaseGraphVertex {
             throw new IllegalStateException("Cannot do backward pass: error not set");
 
         INDArray out = workspaceMgr.create(ArrayType.ACTIVATION_GRAD, inputs[0].dataType(), forwardShape);
-        int start = from * step;
+        long start = from * step;
-        int end = (from + 1) * step;
+        long end = (from + 1) * step;
 
         switch (forwardShape.length) {
             case 2:
@@ -154,8 +153,8 @@ public class UnstackVertex extends BaseGraphVertex {
         }
 
         //Mask arrays are either 1d (column vector) or 2d...
-        int start = from * minibatchSize;
+        long start = from * minibatchSize;
-        int end = (from + 1) * minibatchSize;
+        long end = (from + 1) * minibatchSize;
         INDArray outMask = maskArrays[0].get(NDArrayIndex.interval(start, end), NDArrayIndex.all());
         return new Pair<>(outMask, currentMaskState);
     }

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/graph/vertex/impl/rnn/LastTimeStepVertex.java

@@ -87,9 +87,8 @@ public class LastTimeStepVertex extends BaseGraphVertex {
 
         INDArray out;
         if (mask == null) {
-            // FIXME: int cast
             //No mask array -> extract same (last) column for all
-            int lastTS = (int) inputs[0].size(2) - 1;
+            long lastTS = inputs[0].size(2) - 1;
             out = inputs[0].get(NDArrayIndex.all(), NDArrayIndex.all(), NDArrayIndex.point(lastTS));
             out = workspaceMgr.dup(ArrayType.ACTIVATIONS, out);
             fwdPassTimeSteps = null; //Null -> last time step for all examples
@@ -99,8 +98,7 @@ public class LastTimeStepVertex extends BaseGraphVertex {
 
             //Want the index of the last non-zero entry in the mask array.
             //Check a little here by using mulRowVector([0,1,2,3,...]) and argmax
-            // FIXME: int cast
+            long maxTsLength = fwdPassShape[2];
-            int maxTsLength = (int) fwdPassShape[2];
             INDArray row = Nd4j.linspace(0, maxTsLength - 1, maxTsLength, mask.dataType());
             INDArray temp = mask.mulRowVector(row);
             INDArray lastElementIdx = Nd4j.argMax(temp, 1);

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/layers/AbstractLayer.java

@@ -346,7 +346,6 @@ public abstract class AbstractLayer<LayerConfT extends org.deeplearning4j.nn.con
 
     @Override
     public int getInputMiniBatchSize() {
-        // FIXME: int cast
         return (int) input.size(0);
     }
 

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/layers/BaseOutputLayer.java

@@ -229,7 +229,6 @@ public abstract class BaseOutputLayer<LayerConfT extends org.deeplearning4j.nn.c
      */
     @Override
     public int numLabels() {
-        // FIXME: int cast
         return (int) labels.size(1);
     }
 

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/layers/LossLayer.java

@@ -236,7 +236,6 @@ public class LossLayer extends BaseLayer<org.deeplearning4j.nn.conf.layers.LossL
      */
     @Override
     public int numLabels() {
-        // FIXME: int cast
         return (int) labels.size(1);
     }
 

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/layers/convolution/Cnn3DLossLayer.java

@@ -86,19 +86,18 @@ public class Cnn3DLossLayer extends BaseLayer<org.deeplearning4j.nn.conf.layers.
         INDArray delta2d = lossFunction.computeGradient(labels2d, input2d.dup(input2d.ordering()), layerConf().getActivationFn(), maskReshaped);
         delta2d = workspaceMgr.leverageTo(ArrayType.ACTIVATION_GRAD, delta2d);
 
-        // FIXME: int cast
+        long n = input.size(0);
-        int n = (int)input.size(0);
+        long d, h, w, c;
-        int d, h, w, c;
         if(layerConf().getDataFormat() == Convolution3D.DataFormat.NDHWC){
-            d = (int)input.size(1);
+            d = input.size(1);
-            h = (int)input.size(2);
+            h = input.size(2);
-            w = (int)input.size(3);
+            w = input.size(3);
-            c = (int)input.size(4);
+            c = input.size(4);
         } else {
-            d = (int)input.size(2);
+            d = input.size(2);
-            h = (int)input.size(3);
+            h = input.size(3);
-            w = (int)input.size(4);
+            w = input.size(4);
-            c = (int)input.size(1);
+            c = input.size(1);
         }
         INDArray delta5d = ConvolutionUtils.reshape2dTo5d(layerConf().getDataFormat(), delta2d, n, d, h, w, c, workspaceMgr, ArrayType.ACTIVATION_GRAD);
 
@@ -130,7 +129,6 @@ public class Cnn3DLossLayer extends BaseLayer<org.deeplearning4j.nn.conf.layers.
 
     @Override
     public int numLabels() {
-        // FIXME: int cast
         return (int) labels.size(1);
     }
 
@@ -180,10 +178,8 @@ public class Cnn3DLossLayer extends BaseLayer<org.deeplearning4j.nn.conf.layers.
         INDArray input2d = ConvolutionUtils.reshape5dTo2d(layerConf().getDataFormat(), in, workspaceMgr, ArrayType.ACTIVATIONS);
         INDArray out2d = layerConf().getActivationFn().getActivation(input2d, training);
 
-        // FIXME: int cast
+        long n = input.size(0);
-
+        long d, h, w, c;
-        int n = (int)input.size(0);
-        int d, h, w, c;
         if(layerConf().getDataFormat() == Convolution3D.DataFormat.NDHWC){
             d = (int)input.size(1);
             h = (int)input.size(2);
@@ -262,19 +258,18 @@ public class Cnn3DLossLayer extends BaseLayer<org.deeplearning4j.nn.conf.layers.
         val newShape = input.shape().clone();
         newShape[1] = 1;
 
-        // FIXME
+        long n = input.size(0);
-        int n = (int)input.size(0);
+        long d, h, w, c;
-        int d, h, w, c;
         if(layerConf().getDataFormat() == Convolution3D.DataFormat.NDHWC){
-            d = (int)input.size(1);
+            d = input.size(1);
-            h = (int)input.size(2);
+            h = input.size(2);
-            w = (int)input.size(3);
+            w = input.size(3);
-            c = (int)input.size(4);
+            c = input.size(4);
         } else {
-            d = (int)input.size(2);
+            d = input.size(2);
-            h = (int)input.size(3);
+            h = input.size(3);
-            w = (int)input.size(4);
+            w = input.size(4);
-            c = (int)input.size(1);
+            c = input.size(1);
         }
         INDArray scoreArrayTs = ConvolutionUtils.reshape2dTo5d(layerConf().getDataFormat(), scoreArray, n, d, h, w, c, workspaceMgr, ArrayType.FF_WORKING_MEM);
         INDArray summedScores = scoreArrayTs.sum(1,2,3,4);

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/layers/convolution/CnnLossLayer.java

@@ -88,8 +88,7 @@ public class CnnLossLayer extends BaseLayer<org.deeplearning4j.nn.conf.layers.Cn
         INDArray delta2d = lossFunction.computeGradient(labels2d, input2d.dup(input2d.ordering()), layerConf().getActivationFn(), maskReshaped);
         delta2d = workspaceMgr.leverageTo(ArrayType.ACTIVATION_GRAD, delta2d);
 
-        // FIXME: int cast
+        INDArray delta4d = ConvolutionUtils.reshape2dTo4d(delta2d, input.shape(), workspaceMgr, ArrayType.ACTIVATION_GRAD);
-        INDArray delta4d = ConvolutionUtils.reshape2dTo4d(delta2d, ArrayUtil.toInts(input.shape()), workspaceMgr, ArrayType.ACTIVATION_GRAD);
 
         // grab the empty gradient
         Gradient gradient = new DefaultGradient();
@@ -119,7 +118,6 @@ public class CnnLossLayer extends BaseLayer<org.deeplearning4j.nn.conf.layers.Cn
 
     @Override
     public int numLabels() {
-        // FIXME: int cast
         return (int) labels.size(1);
     }
 
@@ -169,8 +167,7 @@ public class CnnLossLayer extends BaseLayer<org.deeplearning4j.nn.conf.layers.Cn
         INDArray input2d = ConvolutionUtils.reshape4dTo2d(in, workspaceMgr, ArrayType.ACTIVATIONS);
         INDArray out2d = layerConf().getActivationFn().getActivation(input2d, training);
 
-        // FIXME: int cast
+        return ConvolutionUtils.reshape2dTo4d(out2d, input.shape(), workspaceMgr, ArrayType.ACTIVATIONS);
-        return ConvolutionUtils.reshape2dTo4d(out2d, ArrayUtil.toInts(input.shape()), workspaceMgr, ArrayType.ACTIVATIONS);
     }
 
     @Override
@@ -236,8 +233,7 @@ public class CnnLossLayer extends BaseLayer<org.deeplearning4j.nn.conf.layers.Cn
         val newShape = input.shape().clone();
         newShape[1] = 1;
 
-        // FIXME
+        INDArray scoreArrayTs = ConvolutionUtils.reshape2dTo4d(scoreArray, newShape, workspaceMgr, ArrayType.FF_WORKING_MEM);
-        INDArray scoreArrayTs = ConvolutionUtils.reshape2dTo4d(scoreArray, ArrayUtil.toInts(newShape), workspaceMgr, ArrayType.FF_WORKING_MEM);
         INDArray summedScores = scoreArrayTs.sum(1,2,3).reshape(scoreArrayTs.size(0), 1);
 
         if (fullNetRegTerm != 0.0) {

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/layers/convolution/Convolution3DLayer.java

@@ -71,8 +71,7 @@ public class Convolution3DLayer extends ConvolutionLayer {
 
         boolean isNCDHW = layerConfig.getDataFormat() == Convolution3D.DataFormat.NCDHW;
 
-        // FIXME: int cast
+        long miniBatch = input.size(0);
-        int miniBatch = (int) input.size(0);
         int inD = (int) (isNCDHW ? input.size(2) : input.size(1));
         int inH = (int) (isNCDHW ? input.size(3) : input.size(2));
         int inW = (int) (isNCDHW ? input.size(4) : input.size(3));
@@ -189,8 +188,7 @@ public class Convolution3DLayer extends ConvolutionLayer {
                     + " " + layerId());
         }
 
-        // FIXME: int cast
+        long miniBatch = input.size(0);
-        int miniBatch = (int) input.size(0);
         int inputChannels = (int) (isNCDHW ? input.size(1) : input.size(4));
         int inD =(int) (isNCDHW ? input.size(2) : input.size(1));
         int inH = (int) (isNCDHW ? input.size(3) : input.size(2));

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/layers/convolution/ConvolutionLayer.java

@@ -35,6 +35,7 @@ import org.nd4j.linalg.api.memory.MemoryWorkspace;
 import org.nd4j.linalg.api.ndarray.INDArray;
 import org.nd4j.linalg.api.shape.Shape;
 import org.nd4j.linalg.convolution.Convolution;
+import org.nd4j.linalg.exception.ND4JArraySizeException;
 import org.nd4j.linalg.exception.ND4JOpProfilerException;
 import org.nd4j.linalg.factory.Nd4j;
 import org.nd4j.linalg.primitives.Pair;
@@ -113,13 +114,12 @@ public class ConvolutionLayer extends BaseLayer<org.deeplearning4j.nn.conf.layer
         if(epsilon.dataType() != dataType)
             epsilon = epsilon.castTo(dataType);
 
-        // FIXME: int cast
+        long miniBatch = input.size(0);
-        int miniBatch = (int) input.size(0);
         int inH = (int) input.size(2);
         int inW = (int) input.size(3);
 
-        int outDepth = (int) weights.size(0);
+        long outDepth = weights.size(0);
-        int inDepth = (int) weights.size(1);
+        long inDepth = weights.size(1);
         int kH = (int) weights.size(2);
         int kW = (int) weights.size(3);
 
@@ -143,7 +143,7 @@ public class ConvolutionLayer extends BaseLayer<org.deeplearning4j.nn.conf.layer
         INDArray biasGradView = gradientViews.get(ConvolutionParamInitializer.BIAS_KEY);
         INDArray weightGradView = gradientViews.get(ConvolutionParamInitializer.WEIGHT_KEY); //4d, c order. Shape: [outDepth,inDepth,kH,kW]
         INDArray weightGradView2df = Shape
-                        .newShapeNoCopy(weightGradView, new int[] {outDepth, inDepth * kH * kW}, false).transpose();
+                        .newShapeNoCopy(weightGradView, new long[]{outDepth, inDepth * kH * kW}, false).transpose();
 
 
 
@@ -204,7 +204,7 @@ public class ConvolutionLayer extends BaseLayer<org.deeplearning4j.nn.conf.layer
 
         //Note: due to the permute in preOut, and the fact that we essentially do a preOut.muli(epsilon), this reshape
         // should be zero-copy; only possible exception being sometimes with the "identity" activation case
-        INDArray delta2d = delta.reshape('c', new int[] {outDepth, miniBatch * outH * outW}); //Shape.newShapeNoCopy(delta,new int[]{outDepth,miniBatch*outH*outW},false);
+        INDArray delta2d = delta.reshape('c', new long[] {outDepth, miniBatch * outH * outW}); //Shape.newShapeNoCopy(delta,new int[]{outDepth,miniBatch*outH*outW},false);
 
         //Do im2col, but with order [miniB,outH,outW,depthIn,kH,kW]; but need to input [miniBatch,channels,kH,kW,outH,outW] given the current im2col implementation
         //To get this: create an array of the order we want, permute it to the order required by im2col implementation, and then do im2col on that
@@ -231,7 +231,7 @@ public class ConvolutionLayer extends BaseLayer<org.deeplearning4j.nn.conf.layer
         //Calculate epsilons for layer below, in 2d format (note: this is in 'image patch' format before col2im reduction)
         //Note: cc -> f mmul here, then reshape to 6d in f order
         INDArray epsNext2d = w2d.mmul(delta2d); //TODO can we reuse im2col array instead of allocating new result array?
-        INDArray eps6d = Shape.newShapeNoCopy(epsNext2d, new int[] {kW, kH, inDepth, outW, outH, miniBatch}, true);
+        INDArray eps6d = Shape.newShapeNoCopy(epsNext2d, new long[] {kW, kH, inDepth, outW, outH, miniBatch}, true);
 
         //Calculate epsilonNext by doing im2col reduction.
         //Current col2im implementation expects input with order: [miniBatch,channels,kH,kW,outH,outW]
@@ -282,7 +282,7 @@ public class ConvolutionLayer extends BaseLayer<org.deeplearning4j.nn.conf.layer
         }
     }
 
-    protected void validateInputDepth(int inDepth) {
+    protected void validateInputDepth(long inDepth) {
         if (input.size(1) != inDepth) {
             String layerName = conf.getLayer().getLayerName();
             if (layerName == null)
@@ -313,14 +313,13 @@ public class ConvolutionLayer extends BaseLayer<org.deeplearning4j.nn.conf.layer
 
         INDArray input = this.input.castTo(dataType);
 
-        // FIXME: int cast
+        long miniBatch = input.size(0);
-        int miniBatch = (int) input.size(0);
+        long outDepth = weights.size(0);
-        int outDepth = (int) weights.size(0);
+        long inDepth = weights.size(1);
-        int inDepth = (int) weights.size(1);
         validateInputDepth(inDepth);
 
-        int kH = (int) weights.size(2);
+        long kH = weights.size(2);
-        int kW = (int) weights.size(3);
+        long kW = weights.size(3);
 
         int[] dilation = layerConf().getDilation();
         int[] kernel = layerConf().getKernelSize();
@@ -331,7 +330,8 @@ public class ConvolutionLayer extends BaseLayer<org.deeplearning4j.nn.conf.layer
         if (convolutionMode == ConvolutionMode.Same) {
             outSize = ConvolutionUtils.getOutputSize(input, kernel, strides, null, convolutionMode, dilation); //Also performs validation
 
-            // FIXME: int cast
+            if (input.size(2) > Integer.MAX_VALUE ||  input.size(3) > Integer.MAX_VALUE)
+                throw new ND4JArraySizeException();
             pad = ConvolutionUtils.getSameModeTopLeftPadding(outSize, new int[] {(int) input.size(2), (int) input.size(3)}, kernel,
                             strides, dilation );
         } else {
@@ -397,10 +397,12 @@ public class ConvolutionLayer extends BaseLayer<org.deeplearning4j.nn.conf.layer
         INDArray col = Nd4j.createUninitialized(weights.dataType(), new long[] {miniBatch, outH, outW, inDepth, kH, kW}, 'c');
         INDArray col2 = col.permute(0, 3, 4, 5, 1, 2);
         INDArray im2ColIn = input.castTo(col2.dataType());      //No op if already (for example) float
-        Convolution.im2col(im2ColIn, kH, kW, strides[0], strides[1], pad[0], pad[1], dilation[0], dilation[1],
+        if (kH > Integer.MAX_VALUE || kW > Integer.MAX_VALUE)
+            throw new ND4JArraySizeException();
+        Convolution.im2col(im2ColIn, (int)kH, (int)kW, strides[0], strides[1], pad[0], pad[1], dilation[0], dilation[1],
                         convolutionMode == ConvolutionMode.Same, col2);
 
-        INDArray im2col2d = Shape.newShapeNoCopy(col, new int[] {miniBatch * outH * outW, inDepth * kH * kW}, false);
+        INDArray im2col2d = Shape.newShapeNoCopy(col, new long[] {miniBatch * outH * outW, inDepth * kH * kW}, false);
 
         //Current order of weights: [depthOut,depthIn,kH,kW], c order
         //Permute to give [kW,kH,depthIn,depthOut], f order
@@ -418,7 +420,7 @@ public class ConvolutionLayer extends BaseLayer<org.deeplearning4j.nn.conf.layer
         }
 
         //Now, reshape to [outW,outH,miniBatch,outDepth], and permute to have correct output order: [miniBath,outDepth,outH,outW];
-        z = Shape.newShapeNoCopy(z, new int[] {outW, outH, miniBatch, outDepth}, true);
+        z = Shape.newShapeNoCopy(z, new long[] {outW, outH, miniBatch, outDepth}, true);
         z = z.permute(2, 3, 1, 0);
 
         if (training && cacheMode != CacheMode.NONE && workspaceMgr.hasConfiguration(ArrayType.FF_CACHE) && workspaceMgr.isWorkspaceOpen(ArrayType.FF_CACHE)) {

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/layers/convolution/Deconvolution2DLayer.java

@@ -171,9 +171,8 @@ public class Deconvolution2DLayer extends ConvolutionLayer {
                     + " " + layerId());
         }
 
-        // FIXME: int cast
+        long inDepth = weights.size(0);
-        int inDepth = (int) weights.size(0);
+        long outDepth = weights.size(1);
-        int outDepth = (int) weights.size(1);
 
         if (input.size(1) != inDepth && input.size(3) == inDepth) {
             //TODO AB 2019/10/25 this is an ugly "pseudo-NHWC support" hack that needs to be removed ASAD
@@ -199,7 +198,6 @@ public class Deconvolution2DLayer extends ConvolutionLayer {
         int[] pad;
         int[] outSize;
         if (convolutionMode == ConvolutionMode.Same) {
-            // FIXME: int cast
             outSize = ConvolutionUtils.getDeconvolutionOutputSize(input, kernel, strides, null, convolutionMode, dilation); //Also performs validation
             pad = ConvolutionUtils.getSameModeTopLeftPadding(outSize, new int[] {(int) input.size(2), (int) input.size(3)}, kernel,
                     strides, dilation );
@@ -208,8 +206,8 @@ public class Deconvolution2DLayer extends ConvolutionLayer {
             outSize = ConvolutionUtils.getDeconvolutionOutputSize(input, kernel, strides, pad, convolutionMode, dilation); //Also performs validation
         }
 
-        int outH = outSize[0];
+        long outH = outSize[0];
-        int outW = outSize[1];
+        long outW = outSize[1];
 
 
         val miniBatch = input.size(0);

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/layers/convolution/DepthwiseConvolution2DLayer.java

@@ -32,6 +32,7 @@ import org.nd4j.linalg.api.buffer.DataType;
 import org.nd4j.linalg.api.ndarray.INDArray;
 import org.nd4j.linalg.api.ops.CustomOp;
 import org.nd4j.linalg.api.ops.DynamicCustomOp;
+import org.nd4j.linalg.exception.ND4JArraySizeException;
 import org.nd4j.linalg.factory.Nd4j;
 import org.nd4j.linalg.primitives.Pair;
 
@@ -75,12 +76,11 @@ public class DepthwiseConvolution2DLayer extends ConvolutionLayer {
 
         INDArray input = this.input.castTo(dataType);   //No-op if correct type
 
-        // FIXME: int cast
+        long miniBatch = input.size(0);
-        int miniBatch = (int) input.size(0);
         int inH = (int)input.size(2);
         int inW = (int)input.size(3);
 
-        int inDepth = (int) depthWiseWeights.size(2);
+        long inDepth = depthWiseWeights.size(2);
         int kH = (int) depthWiseWeights.size(0);
         int kW = (int) depthWiseWeights.size(1);
 
@@ -169,10 +169,9 @@ public class DepthwiseConvolution2DLayer extends ConvolutionLayer {
 
         INDArray input = this.input.castTo(dataType);   //no-op if correct dtype
 
-        // FIXME: int cast
+        long inDepth = depthWiseWeights.size(2);
-        int inDepth = (int) depthWiseWeights.size(2);
+        long depthMultiplier = depthWiseWeights.size(3);
-        int depthMultiplier = (int) depthWiseWeights.size(3);
+        long outDepth = depthMultiplier * inDepth;
-        int outDepth = depthMultiplier * inDepth;
 
         if (input.size(1) != inDepth) {
             String layerName = conf.getLayer().getLayerName();
@@ -197,7 +196,9 @@ public class DepthwiseConvolution2DLayer extends ConvolutionLayer {
         if (convolutionMode == ConvolutionMode.Same) {
             outSize = ConvolutionUtils.getOutputSize(input, kernel, strides, null, convolutionMode, dilation);
 
-            // FIXME: int cast
+            if (input.size(2) > Integer.MAX_VALUE || input.size(3) > Integer.MAX_VALUE) {
+                throw new ND4JArraySizeException();
+            }
             pad = ConvolutionUtils.getSameModeTopLeftPadding(
                     outSize, new int[]{(int) input.size(2), (int) input.size(3)}, kernel, strides, dilation);
         } else {
@@ -205,8 +206,8 @@ public class DepthwiseConvolution2DLayer extends ConvolutionLayer {
             outSize = ConvolutionUtils.getOutputSize(input, kernel, strides, pad, convolutionMode, dilation);
         }
 
-        int outH = outSize[0];
+        long outH = outSize[0];
-        int outW = outSize[1];
+        long outW = outSize[1];
 
         val miniBatch = input.size(0);
         INDArray output = workspaceMgr.create(

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/layers/convolution/SeparableConvolution2DLayer.java

@@ -33,6 +33,7 @@ import org.nd4j.linalg.api.memory.MemoryWorkspace;
 import org.nd4j.linalg.api.ndarray.INDArray;
 import org.nd4j.linalg.api.ops.CustomOp;
 import org.nd4j.linalg.api.ops.DynamicCustomOp;
+import org.nd4j.linalg.exception.ND4JArraySizeException;
 import org.nd4j.linalg.factory.Nd4j;
 import org.nd4j.linalg.primitives.Pair;
 import org.deeplearning4j.nn.workspace.LayerWorkspaceMgr;
@@ -90,8 +91,7 @@ public class SeparableConvolution2DLayer extends ConvolutionLayer {
 
         INDArray input = this.input.castTo(dataType);
 
-        // FIXME: int cast
+        long miniBatch = input.size(0);
-        int miniBatch = (int) input.size(0);
         int inH = (int)input.size(2);
         int inW = (int)input.size(3);
 
@@ -194,9 +194,8 @@ public class SeparableConvolution2DLayer extends ConvolutionLayer {
                     + " " + layerId());
         }
 
-        // FIXME: int cast
+        long inDepth = depthWiseWeights.size(1);
-        int inDepth = (int) depthWiseWeights.size(1);
+        long outDepth = pointWiseWeights.size(0);
-        int outDepth = (int) pointWiseWeights.size(0);
 
         if (input.size(1) != inDepth) {
             String layerName = conf.getLayer().getLayerName();
@@ -220,7 +219,9 @@ public class SeparableConvolution2DLayer extends ConvolutionLayer {
         if (convolutionMode == ConvolutionMode.Same) {
             outSize = ConvolutionUtils.getOutputSize(input, kernel, strides, null, convolutionMode, dilation); //Also performs validation
 
-            // FIXME: int cast
+            if (input.size(2) > Integer.MAX_VALUE || input.size(3) > Integer.MAX_VALUE) {
+                throw new ND4JArraySizeException();
+            }
             pad = ConvolutionUtils.getSameModeTopLeftPadding(outSize, new int[] {(int) input.size(2), (int) input.size(3)}, kernel,
                     strides, dilation );
         } else {

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/layers/convolution/SpaceToDepth.java

@@ -75,11 +75,10 @@ public class SpaceToDepth extends AbstractLayer<org.deeplearning4j.nn.conf.layer
     public Pair<Gradient, INDArray> backpropGradient(INDArray epsilon, LayerWorkspaceMgr workspaceMgr) {
         assertInputSet(true);
 
-        // FIXME: int cast
+        long miniBatch = input.size(0);
-        int miniBatch = (int) input.size(0);
+        long inDepth = input.size(1);
-        int inDepth = (int) input.size(1);
+        long inH = input.size(2);
-        int inH = (int) input.size(2);
+        long inW = input.size(3);
-        int inW = (int) input.size(3);
 
         INDArray input = this.input.castTo(dataType);    //No-op if already correct type
 
@@ -122,17 +121,16 @@ public class SpaceToDepth extends AbstractLayer<org.deeplearning4j.nn.conf.layer
             return preOutput;
         }
 
-        // FIXME: int cast
+        long miniBatch = input.size(0);
-        int miniBatch = (int) input.size(0);
+        long depth = input.size(1);
-        int depth = (int) input.size(1);
+        long inH = input.size(2);
-        int inH = (int) input.size(2);
+        long inW = input.size(3);
-        int inW = (int) input.size(3);
 
         int blockSize = getBlockSize();
 
-        int outH = inH / blockSize;
+        long outH = inH / blockSize;
-        int outW = inW / blockSize;
+        long outW = inW / blockSize;
-        int outDepth = depth * blockSize * blockSize;
+        long outDepth = depth * blockSize * blockSize;
 
         INDArray out = workspaceMgr.create(ArrayType.ACTIVATIONS, input.dataType(), new long[]{1, miniBatch * outDepth * outH * outW}, 'c');
         INDArray reshapedOut;

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/layers/convolution/subsampling/Subsampling3DLayer.java

@@ -71,9 +71,8 @@ public class Subsampling3DLayer extends AbstractLayer<org.deeplearning4j.nn.conf
 
         boolean isNCDHW = layerConf().getDataFormat() == Convolution3D.DataFormat.NCDHW;
 
-        // FIXME: int cast
+        long miniBatch = input.size(0);
-        int miniBatch = (int) input.size(0);
+        long inChannels = isNCDHW ? input.size(1) : input.size(4);
-        int inChannels = (int) (isNCDHW ? input.size(1) : input.size(4));
         int inD = (int) (isNCDHW ? input.size(2) : input.size(1));
         int inH = (int) (isNCDHW ? input.size(3) : input.size(2));
         int inW = (int) (isNCDHW ? input.size(4) : input.size(3));
@@ -148,9 +147,8 @@ public class Subsampling3DLayer extends AbstractLayer<org.deeplearning4j.nn.conf
             }
         }
 
-        // FIXME: int cast
+        long miniBatch = input.size(0);
-        int miniBatch = (int) input.size(0);
+        long inChannels = isNCDHW ? input.size(1) : input.size(4);
-        int inChannels = (int) (isNCDHW ? input.size(1) : input.size(4));
         int inD = (int) (isNCDHW ? input.size(2) : input.size(1));
         int inH = (int) (isNCDHW ? input.size(3) : input.size(2));
         int inW = (int) (isNCDHW ? input.size(4) : input.size(3));
@@ -170,9 +168,9 @@ public class Subsampling3DLayer extends AbstractLayer<org.deeplearning4j.nn.conf
             outSize = Convolution3DUtils.get3DOutputSize(
                     input, kernel, strides, pad, convolutionMode, dilation, isNCDHW);
         }
-        int outD = outSize[0];
+        long outD = outSize[0];
-        int outH = outSize[1];
+        long outH = outSize[1];
-        int outW = outSize[2];
+        long outW = outSize[2];
 
         String opName = layerConf().getPoolingType() == PoolingType.MAX ? "maxpool3dnew" : "avgpool3dnew";
 

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/layers/convolution/subsampling/SubsamplingLayer.java

@@ -108,9 +108,6 @@ public class SubsamplingLayer extends AbstractLayer<org.deeplearning4j.nn.conf.l
         if(epsilon.dataType() != dataType)
             epsilon = epsilon.castTo(dataType);
 
-        // FIXME: int cast
-        int miniBatch = (int) input.size(0);
-        int inDepth = (int) input.size(1);
         int inH = (int)input.size(2);
         int inW = (int)input.size(3);
 
@@ -158,9 +155,6 @@ public class SubsamplingLayer extends AbstractLayer<org.deeplearning4j.nn.conf.l
 
         //subsampling doesn't have weights and thus gradients are not calculated for this layer
         //only scale and reshape epsilon
-        // FIXME: int cast
-        int inputHeight = (int) input().size(-2);
-        int inputWidth = (int) input().size(-1);
         Gradient retGradient = new DefaultGradient();
 
 
@@ -231,9 +225,8 @@ public class SubsamplingLayer extends AbstractLayer<org.deeplearning4j.nn.conf.l
 
         INDArray input = this.input.castTo(dataType);
 
-        // FIXME: int cast
+        long miniBatch = input.size(0);
-        int miniBatch = (int) input.size(0);
+        long inDepth = input.size(1);
-        int inDepth = (int) input.size(1);
         int inH = (int)input.size(2);
         int inW = (int)input.size(3);
 
@@ -250,8 +243,8 @@ public class SubsamplingLayer extends AbstractLayer<org.deeplearning4j.nn.conf.l
             pad = layerConf().getPadding();
             outSize = ConvolutionUtils.getOutputSize(input, kernel, strides, pad, convolutionMode, dilation); //Also performs validation
         }
-        int outH = outSize[0];
+        long outH = outSize[0];
-        int outW = outSize[1];
+        long outW = outSize[1];
 
 
         if (helper != null && (helperCountFail == 0 || !layerConf().isCudnnAllowFallback())) {
@@ -278,9 +271,6 @@ public class SubsamplingLayer extends AbstractLayer<org.deeplearning4j.nn.conf.l
             }
         }
 
-
-
-
         INDArray output = workspaceMgr.createUninitialized(ArrayType.ACTIVATIONS, input.dataType(), new long[]{miniBatch, inDepth, outH, outW}, 'c');
         DynamicCustomOp.DynamicCustomOpsBuilder b;
         int extra = 0;

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/layers/convolution/upsampling/Upsampling1D.java

@@ -65,11 +65,10 @@ public class Upsampling1D extends Upsampling2D {
         INDArray originalInput = input;
         input = input.castTo(dataType).reshape(input.size(0), input.size(1), input.size(2), 1);
 
-        // FIXME: int cast
+        long miniBatch = input.size(0);
-        int miniBatch = (int) input.size(0);
+        long inDepth = input.size(1);
-        int inDepth = (int) input.size(1);
+        long inH = input.size(2);
-        int inH = (int) input.size(2);
+        long inW = input.size(3);
-        int inW = (int) input.size(3);
 
 
         INDArray outEpsilon = workspaceMgr.createUninitialized(ArrayType.ACTIVATION_GRAD, input.dataType(), miniBatch * inDepth * inH * inW);

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/layers/convolution/upsampling/Upsampling2D.java

@@ -62,11 +62,10 @@ public class Upsampling2D extends AbstractLayer<org.deeplearning4j.nn.conf.layer
     public Pair<Gradient, INDArray> backpropGradient(INDArray epsilon, LayerWorkspaceMgr workspaceMgr) {
         assertInputSet(true);
 
-        // FIXME: int cast
+        long miniBatch = (int) input.size(0);
-        int miniBatch = (int) input.size(0);
+        long inDepth = (int) input.size(1);
-        int inDepth = (int) input.size(1);
+        long inH = (int) input.size(2);
-        int inH = (int) input.size(2);
+        long inW = (int) input.size(3);
-        int inW = (int) input.size(3);
 
         INDArray reshapedEpsilon =  workspaceMgr.createUninitialized(ArrayType.ACTIVATION_GRAD, epsilon.dataType(), new long[]{miniBatch, inDepth, inH, inW}, 'c');
 
@@ -106,15 +105,14 @@ public class Upsampling2D extends AbstractLayer<org.deeplearning4j.nn.conf.layer
             return preOutput;
         }
 
-        // FIXME: int cast
+        long miniBatch = (int) input.size(0);
-        int miniBatch = (int) input.size(0);
+        long inDepth = (int) input.size(1);
-        int inDepth = (int) input.size(1);
+        long inH = (int) input.size(2);
-        int inH = (int) input.size(2);
+        long inW = (int) input.size(3);
-        int inW = (int) input.size(3);
 
         int[] size = getSize();
-        int outH = inH * size[0];
+        int outH = (int)inH * size[0];
-        int outW = inW * size[1];
+        int outW = (int)inW * size[1];
 
         INDArray reshapedOutput = workspaceMgr.createUninitialized(ArrayType.ACTIVATIONS, input.dataType(), new long[]{miniBatch, inDepth, outH, outW}, 'c');
 

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/layers/convolution/upsampling/Upsampling3D.java

@@ -68,22 +68,21 @@ public class Upsampling3D extends AbstractLayer<org.deeplearning4j.nn.conf.layer
         assertInputSet(true);
 
         boolean ncdhw = layerConf().getDataFormat() == org.deeplearning4j.nn.conf.layers.Convolution3D.DataFormat.NCDHW;
-        // FIXME: int cast
         // Assumes NCDHW order
-        int miniBatch = (int) input.size(0);
+        long miniBatch = input.size(0);
-        int inChannels, inD, inH, inW;
+        long inChannels, inD, inH, inW;
         int[] intArgs;
         if(ncdhw){
-            inChannels = (int) input.size(1);
+            inChannels = input.size(1);
-            inD = (int) input.size(2);
+            inD = input.size(2);
-            inH = (int) input.size(3);
+            inH = input.size(3);
-            inW = (int) input.size(4);
+            inW = input.size(4);
             intArgs = new int[] {1}; // 1 is channels first
         } else {
-            inD = (int) input.size(1);
+            inD = input.size(1);
-            inH = (int) input.size(2);
+            inH = input.size(2);
-            inW = (int) input.size(3);
+            inW = input.size(3);
-            inChannels = (int) input.size(4);
+            inChannels = input.size(4);
             intArgs = new int[] {0}; // 0 is channels last
         }
 
@@ -134,9 +133,8 @@ public class Upsampling3D extends AbstractLayer<org.deeplearning4j.nn.conf.layer
         }
 
         boolean ncdhw = layerConf().getDataFormat() == org.deeplearning4j.nn.conf.layers.Convolution3D.DataFormat.NCDHW;
-        // FIXME: int cast
+        long miniBatch = input.size(0);
-        int miniBatch = (int) input.size(0);
+        long inChannels, inD, inH, inW;
-        int inChannels, inD, inH, inW;
         int[] intArgs;
         int[] size = getSize();
         if(ncdhw){

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/layers/feedforward/embedding/EmbeddingLayer.java

@@ -20,6 +20,7 @@ import lombok.extern.slf4j.Slf4j;
 import lombok.val;
 import org.nd4j.linalg.api.buffer.DataType;
 import org.nd4j.linalg.api.ops.custom.ScatterUpdate;
+import org.nd4j.linalg.exception.ND4JArraySizeException;
 import org.nd4j.linalg.primitives.Pair;
 import org.deeplearning4j.exception.DL4JInvalidInputException;
 import org.deeplearning4j.nn.conf.NeuralNetConfiguration;
@@ -64,8 +65,7 @@ public class EmbeddingLayer extends BaseLayer<org.deeplearning4j.nn.conf.layers.
         INDArray weightGradients = gradientViews.get(DefaultParamInitializer.WEIGHT_KEY);
         weightGradients.assign(0);
 
-        // FIXME: int cast
+        long[] indexes = new long[(int) input.length()];
-        int[] indexes = new int[(int) input.length()];
         for (int i = 0; i < indexes.length; i++) {
             indexes[i] = input.getInt(i, 0);
         }
@@ -99,7 +99,8 @@ public class EmbeddingLayer extends BaseLayer<org.deeplearning4j.nn.conf.layers.
 
         val nIn = layerConf().getNIn();
 
-        // FIXME: int cast
+        if (input.length() > Integer.MAX_VALUE)
+            throw new ND4JArraySizeException();
         int[] indexes = new int[(int) input.length()];
         for (int i = 0; i < indexes.length; i++) {
             indexes[i] = input.getInt(i, 0);

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/layers/mkldnn/MKLDNNBatchNormHelper.java

@@ -56,7 +56,7 @@ public class MKLDNNBatchNormHelper implements BatchNormalizationHelper {
     }
 
     @Override
-    public Pair<Gradient, INDArray> backpropGradient(INDArray input, INDArray epsilon, int[] shape, INDArray gamma,
+    public Pair<Gradient, INDArray> backpropGradient(INDArray input, INDArray epsilon, long[] shape, INDArray gamma,
                                                      INDArray dGammaView, INDArray dBetaView, double eps, LayerWorkspaceMgr workspaceMgr) {
         //2019-02-14: Backprop disabled pending fixes. https://github.com/deeplearning4j/deeplearning4j/issues/7166
         //Also no MKL-DNN implemented for backprop anyway
@@ -82,7 +82,7 @@ public class MKLDNNBatchNormHelper implements BatchNormalizationHelper {
     }
 
     @Override
-    public INDArray preOutput(INDArray x, boolean training, int[] shape, INDArray gamma, INDArray beta, INDArray mean, INDArray var,
+    public INDArray preOutput(INDArray x, boolean training, long[] shape, INDArray gamma, INDArray beta, INDArray mean, INDArray var,
                               double decay, double eps, LayerWorkspaceMgr workspaceMgr) {
         if(x.dataType() != DataType.FLOAT)
             return null;    //MKL-DNN only supports float

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/layers/normalization/BatchNormalization.java

@@ -152,10 +152,9 @@ public class BatchNormalization extends BaseLayer<org.deeplearning4j.nn.conf.lay
                 eps = epsilon;
             }
 
-            // FIXME: int cast
             Pair<Gradient,INDArray> ret = null;
             try {
-                ret = helper.backpropGradient(in, eps, ArrayUtil.toInts(shape), gamma, dGammaView, dBetaView,
+                ret = helper.backpropGradient(in, eps, shape, gamma, dGammaView, dBetaView,
                         layerConf.getEps(), workspaceMgr);
             } catch (ND4JOpProfilerException e){
                 throw e;    //NaN panic etc for debugging
@@ -438,7 +437,6 @@ public class BatchNormalization extends BaseLayer<org.deeplearning4j.nn.conf.lay
             //Note that cudnn does not support dense (2d) batch norm case as of v7.1
             double decay = layerConf.getDecay();
 
-            // FIXME: int cast
             INDArray ret = null;
             try {
                 if(globalVarView == null){
@@ -448,7 +446,7 @@ public class BatchNormalization extends BaseLayer<org.deeplearning4j.nn.conf.lay
                     globalVarView.muli(globalVarView);
                 }
 
-                ret = helper.preOutput(in, training == TrainingMode.TRAIN, ArrayUtil.toInts(shape), gamma, beta, globalMeanView,
+                ret = helper.preOutput(in, training == TrainingMode.TRAIN, shape, gamma, beta, globalMeanView,
                         globalVarView, decay, layerConf.getEps(), workspaceMgr);
             } catch (ND4JOpProfilerException e){
                 throw e;    //NaN panic etc for debugging

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/layers/normalization/BatchNormalizationHelper.java

@@ -31,10 +31,10 @@ import org.deeplearning4j.nn.workspace.LayerWorkspaceMgr;
 public interface BatchNormalizationHelper extends LayerHelper {
     boolean checkSupported(double eps, boolean fixedGammaBeta);
 
-    Pair<Gradient, INDArray> backpropGradient(INDArray input, INDArray epsilon, int[] shape, INDArray gamma,
+    Pair<Gradient, INDArray> backpropGradient(INDArray input, INDArray epsilon, long[] shape, INDArray gamma,
                     INDArray dGammaView, INDArray dBetaView, double eps, LayerWorkspaceMgr workspaceMgr);
 
-    INDArray preOutput(INDArray x, boolean training, int[] shape, INDArray gamma, INDArray beta, INDArray mean,
+    INDArray preOutput(INDArray x, boolean training, long[] shape, INDArray gamma, INDArray beta, INDArray mean,
                     INDArray var, double decay, double eps, LayerWorkspaceMgr workspaceMgr);
 
     INDArray getMeanCache(DataType dataType);

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/layers/objdetect/Yolo2OutputLayer.java

@@ -114,10 +114,9 @@ public class Yolo2OutputLayer extends AbstractLayer<org.deeplearning4j.nn.conf.l
         double lambdaCoord = layerConf().getLambdaCoord();
         double lambdaNoObj = layerConf().getLambdaNoObj();
 
-        // FIXME: int cast
+        long mb = input.size(0);
-        int mb = (int) input.size(0);
+        long h = input.size(2);
-        int h = (int) input.size(2);
+        long w = input.size(3);
-        int w = (int) input.size(3);
         int b = (int) layerConf().getBoundingBoxes().size(0);
         int c = (int) labels.size(1)-4;
 
@@ -243,12 +242,12 @@ public class Yolo2OutputLayer extends AbstractLayer<org.deeplearning4j.nn.conf.l
 
         //Class prediction loss
         INDArray classPredictionsPreSoftmax2d = inputClassesPreSoftmax.permute(0,1,3,4,2) //[minibatch, b, c, h, w] To [mb, b, h, w, c]
-                .dup('c').reshape('c', new int[]{mb*b*h*w, c});
+                .dup('c').reshape('c', new long[]{mb*b*h*w, c});
         INDArray classLabelsBroadcast = Nd4j.createUninitialized(input.dataType(), new long[]{mb, b, c, h, w}, 'c');
         for(int i=0; i<b; i++ ){
             classLabelsBroadcast.get(all(), point(i), all(), all(), all()).assign(classLabels); //[mb, c, h, w] to [mb, b, c, h, w]
         }
-        INDArray classLabels2d = classLabelsBroadcast.permute(0,1,3,4,2).dup('c').reshape('c', new int[]{mb*b*h*w, c});
+        INDArray classLabels2d = classLabelsBroadcast.permute(0,1,3,4,2).dup('c').reshape('c', new long[]{mb*b*h*w, c});
 
         //Calculate the loss:
         ILossFunction lossConfidence = new LossL2();
@@ -297,7 +296,7 @@ public class Yolo2OutputLayer extends AbstractLayer<org.deeplearning4j.nn.conf.l
         // ----- Gradient Calculation (specifically: return dL/dIn -----
 
         INDArray epsOut = workspaceMgr.createUninitialized(ArrayType.ACTIVATION_GRAD, input.dataType(), input.shape(), 'c');
-        INDArray epsOut5 = Shape.newShapeNoCopy(epsOut, new int[]{mb, b, 5+c, h, w}, false);
+        INDArray epsOut5 = Shape.newShapeNoCopy(epsOut, new long[]{mb, b, 5+c, h, w}, false);
         INDArray epsClassPredictions = epsOut5.get(all(), all(), interval(5, 5+c), all(), all());    //Shape: [mb, b, 5+c, h, w]
         INDArray epsXY = epsOut5.get(all(), all(), interval(0,2), all(), all());
         INDArray epsWH = epsOut5.get(all(), all(), interval(2,4), all(), all());
@@ -426,16 +425,16 @@ public class Yolo2OutputLayer extends AbstractLayer<org.deeplearning4j.nn.conf.l
      * @return IOU and gradients
      */
     private static IOURet calculateIOULabelPredicted(INDArray labelTL, INDArray labelBR, INDArray predictedWH, INDArray predictedXYinGridBox, INDArray objectPresentMask, INDArray objectPresentMaskBool){
-        // FIXME: int cast
+
-        int mb = (int) labelTL.size(0);
+        long mb = labelTL.size(0);
-        int h = (int) labelTL.size(2);
+        long h = labelTL.size(2);
-        int w = (int) labelTL.size(3);
+        long w = labelTL.size(3);
-        int b = (int) predictedWH.size(1);
+        long b = predictedWH.size(1);
 
         INDArray labelWH = labelBR.sub(labelTL);                //4d [mb, 2, H, W], label W/H in terms of number of grid boxes
 
-        int gridH = (int) labelTL.size(2);
+        long gridH = labelTL.size(2);
-        int gridW = (int) labelTL.size(3);
+        long gridW = labelTL.size(3);
         //Add grid positions to the predicted XY values (to get predicted XY in terms of grid cell units in image,
         // from (0 to 1 in grid cell) format)
         INDArray linspaceX = Nd4j.linspace(0, gridW-1, gridW, predictedWH.dataType());

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/layers/objdetect/YoloUtils.java

@@ -45,12 +45,11 @@ public class YoloUtils {
     }
 
     public static INDArray activate(@NonNull INDArray boundingBoxPriors, @NonNull INDArray input, LayerWorkspaceMgr layerWorkspaceMgr){
-        // FIXME: int cast
+        long mb = input.size(0);
-        int mb = (int) input.size(0);
+        long h = input.size(2);
-        int h = (int) input.size(2);
+        long w = input.size(3);
-        int w = (int) input.size(3);
+        long b = boundingBoxPriors.size(0);
-        int b = (int) boundingBoxPriors.size(0);
+        long c = input.size(1)/b-5;  //input.size(1) == b * (5 + C) -> C = (input.size(1)/b) - 5
-        int c = (int) (input.size(1)/b)-5;  //input.size(1) == b * (5 + C) -> C = (input.size(1)/b) - 5
 
         INDArray output = layerWorkspaceMgr.create(ArrayType.ACTIVATIONS, input.dataType(), input.shape(), 'c');
         INDArray output5 = output.reshape('c', mb, b, 5+c, h, w);
@@ -77,7 +76,7 @@ public class YoloUtils {
         //TODO OPTIMIZE?
         INDArray inputClassesPreSoftmax = input5.get(all(), all(), interval(5, 5+c), all(), all());   //Shape: [minibatch, C, H, W]
         INDArray classPredictionsPreSoftmax2d = inputClassesPreSoftmax.permute(0,1,3,4,2) //[minibatch, b, c, h, w] To [mb, b, h, w, c]
-                .dup('c').reshape('c', new int[]{mb*b*h*w, c});
+                .dup('c').reshape('c', new long[]{mb*b*h*w, c});
         Transforms.softmax(classPredictionsPreSoftmax2d, false);
         INDArray postSoftmax5d = classPredictionsPreSoftmax2d.reshape('c', mb, b, h, w, c ).permute(0, 1, 4, 2, 3);
 
@@ -173,13 +172,12 @@ public class YoloUtils {
             throw new IllegalStateException("Invalid confidence threshold: must be in range [0,1]. Got: " + confThreshold);
         }
 
-        // FIXME: int cast
         //Activations format: [mb, 5b+c, h, w]
-        int mb = (int) networkOutput.size(0);
+        long mb = networkOutput.size(0);
-        int h = (int) networkOutput.size(2);
+        long h = networkOutput.size(2);
-        int w = (int) networkOutput.size(3);
+        long w = networkOutput.size(3);
-        int b = (int) boundingBoxPriors.size(0);
+        long b = boundingBoxPriors.size(0);
-        int c = (int) (networkOutput.size(1)/b)-5;  //input.size(1) == b * (5 + C) -> C = (input.size(1)/b) - 5
+        long c = (networkOutput.size(1)/b)-5;  //input.size(1) == b * (5 + C) -> C = (input.size(1)/b) - 5
 
         //Reshape from [minibatch, B*(5+C), H, W] to [minibatch, B, 5+C, H, W] to [minibatch, B, 5, H, W]
         INDArray output5 = networkOutput.dup('c').reshape(mb, b, 5+c, h, w);

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/layers/recurrent/LSTMHelpers.java

@@ -40,6 +40,7 @@ import org.nd4j.linalg.api.ndarray.INDArray;
 import org.nd4j.linalg.api.ops.impl.transforms.pairwise.arithmetic.MulOp;
 import org.nd4j.linalg.api.ops.impl.transforms.same.TimesOneMinus;
 import org.nd4j.linalg.api.shape.Shape;
+import org.nd4j.linalg.exception.ND4JArraySizeException;
 import org.nd4j.linalg.exception.ND4JOpProfilerException;
 import org.nd4j.linalg.factory.Nd4j;
 import org.nd4j.linalg.indexing.NDArrayIndex;
@@ -113,7 +114,9 @@ public class LSTMHelpers {
 
         input = input.castTo(inputWeights.dataType());  //No-op if already correct dtype
 
-        // FIXME
+        if ((!is2dInput && (input.size(2) > Integer.MAX_VALUE)) ||
+            recurrentWeights.size(0) > Integer.MAX_VALUE || input.size(0) > Integer.MAX_VALUE)
+            throw new ND4JArraySizeException();
         int timeSeriesLength = (int) (is2dInput ? 1 : input.size(2));
         int hiddenLayerSize = (int) recurrentWeights.size(0);
         int miniBatchSize = (int) input.size(0);
@@ -550,7 +553,8 @@ public class LSTMHelpers {
         for (long iTimeIndex = timeSeriesLength - 1; iTimeIndex >= endIdx; iTimeIndex--) {
             try(MemoryWorkspace ws = workspaceMgr.notifyScopeEntered(ArrayType.RNN_BP_LOOP_WORKING_MEM)) {
 
-                // FIXME: int cast
+                if (iTimeIndex > Integer.MAX_VALUE)
+                    throw new ND4JArraySizeException();
                 int time = (int) iTimeIndex;
                 int inext = 1;
 
@@ -574,8 +578,6 @@ public class LSTMHelpers {
                         (iTimeIndex == 0 ? fwdPass.prevAct : fwdPass.fwdPassOutputAsArrays[(int) (time - inext)]);
                 INDArray currMemCellState = fwdPass.memCellState[(int) time];
 
-
-                // FIXME: int cast
                 //LSTM unit output errors (dL/d(a_out)); not to be confused with \delta=dL/d(z_out)
                 INDArray epsilonSlice = (is2dInput ? epsilon : epsilon.tensorAlongDimension((int) time, 1, 0)); //(w^{L+1}*(delta^{(L+1)t})^T)^T or equiv.
 

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/layers/recurrent/RnnLossLayer.java

@@ -89,8 +89,7 @@ public class RnnLossLayer extends BaseLayer<org.deeplearning4j.nn.conf.layers.Rn
         ILossFunction lossFunction = layerConf().getLossFn();
         INDArray delta2d = lossFunction.computeGradient(labels2d, input2d.dup(input2d.ordering()), layerConf().getActivationFn(), maskReshaped);
 
-        // FIXME: int cast
+        INDArray delta3d = TimeSeriesUtils.reshape2dTo3d(delta2d, input.size(0), workspaceMgr, ArrayType.ACTIVATION_GRAD);
-        INDArray delta3d = TimeSeriesUtils.reshape2dTo3d(delta2d, (int) input.size(0), workspaceMgr, ArrayType.ACTIVATION_GRAD);
 
         // grab the empty gradient
         Gradient gradient = new DefaultGradient();
@@ -119,7 +118,6 @@ public class RnnLossLayer extends BaseLayer<org.deeplearning4j.nn.conf.layers.Rn
 
     @Override
     public int numLabels() {
-        // FIXME: int cast
         return (int) labels.size(1);
     }
 
@@ -167,7 +165,7 @@ public class RnnLossLayer extends BaseLayer<org.deeplearning4j.nn.conf.layers.Rn
 
         INDArray as2d = TimeSeriesUtils.reshape3dTo2d(input);
         INDArray out2d = layerConf().getActivationFn().getActivation(workspaceMgr.dup(ArrayType.ACTIVATIONS, as2d, as2d.ordering()), training);
-        return workspaceMgr.leverageTo(ArrayType.ACTIVATIONS, TimeSeriesUtils.reshape2dTo3d(out2d, (int)input.size(0), workspaceMgr, ArrayType.ACTIVATIONS));
+        return workspaceMgr.leverageTo(ArrayType.ACTIVATIONS, TimeSeriesUtils.reshape2dTo3d(out2d, input.size(0), workspaceMgr, ArrayType.ACTIVATIONS));
     }
 
     @Override
@@ -254,7 +252,6 @@ public class RnnLossLayer extends BaseLayer<org.deeplearning4j.nn.conf.layers.Rn
         //scoreArray: shape [minibatch*timeSeriesLength, 1]
         //Reshape it to [minibatch, timeSeriesLength] then sum over time step
 
-        // FIXME: int cast
         INDArray scoreArrayTs = TimeSeriesUtils.reshapeVectorToTimeSeriesMask(scoreArray, (int)input.size(0));
         INDArray summedScores = scoreArrayTs.sum(1);
 

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/layers/recurrent/RnnOutputLayer.java

@@ -70,8 +70,7 @@ public class RnnOutputLayer extends BaseOutputLayer<org.deeplearning4j.nn.conf.l
         this.input = inputTemp;
         INDArray epsilon2d = gradAndEpsilonNext.getSecond();
 
-        // FIXME: int cast
+        INDArray epsilon3d = TimeSeriesUtils.reshape2dTo3d(epsilon2d, input.size(0), workspaceMgr, ArrayType.ACTIVATION_GRAD);
-        INDArray epsilon3d = TimeSeriesUtils.reshape2dTo3d(epsilon2d, (int) input.size(0), workspaceMgr, ArrayType.ACTIVATION_GRAD);
 
         weightNoiseParams.clear();
 
@@ -145,8 +144,7 @@ public class RnnOutputLayer extends BaseOutputLayer<org.deeplearning4j.nn.conf.l
             }
         }
 
-        // FIXME: int cast
+        return TimeSeriesUtils.reshape2dTo3d(act2d, input.size(0), workspaceMgr, ArrayType.ACTIVATIONS);
-        return TimeSeriesUtils.reshape2dTo3d(act2d, (int) input.size(0), workspaceMgr, ArrayType.ACTIVATIONS);
     }
 
     @Override
@@ -205,7 +203,6 @@ public class RnnOutputLayer extends BaseOutputLayer<org.deeplearning4j.nn.conf.l
         //scoreArray: shape [minibatch*timeSeriesLength, 1]
         //Reshape it to [minibatch, timeSeriesLength] then sum over time step
 
-        // FIXME: int cast
         INDArray scoreArrayTs = TimeSeriesUtils.reshapeVectorToTimeSeriesMask(scoreArray, (int)input.size(0));
         INDArray summedScores = scoreArrayTs.sum(true, 1);
 

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/layers/variational/VariationalAutoencoder.java

@@ -41,6 +41,7 @@ import org.nd4j.linalg.api.blas.Level1;
 import org.nd4j.linalg.api.buffer.DataType;
 import org.nd4j.linalg.api.memory.MemoryWorkspace;
 import org.nd4j.linalg.api.ndarray.INDArray;
+import org.nd4j.linalg.exception.ND4JArraySizeException;
 import org.nd4j.linalg.factory.Nd4j;
 import org.nd4j.linalg.learning.regularization.Regularization;
 import org.nd4j.linalg.lossfunctions.ILossFunction;
@@ -552,7 +553,8 @@ public class VariationalAutoencoder implements Layer {
 
     @Override
     public int batchSize() {
-        // FIXME: int cast
+        if (input.size(0) > Integer.MAX_VALUE)
+            throw new ND4JArraySizeException();
         return (int) input.size(0);
     }
 
@@ -862,7 +864,8 @@ public class VariationalAutoencoder implements Layer {
 
     @Override
     public int getInputMiniBatchSize() {
-        // FIXME: int cast
+        if (input.size(0) > Integer.MAX_VALUE)
+            throw new ND4JArraySizeException();
         return (int) input.size(0);
     }
 

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/multilayer/MultiLayerNetwork.java

@@ -75,6 +75,7 @@ import org.nd4j.linalg.dataset.DataSet;
 import org.nd4j.linalg.dataset.api.MultiDataSet;
 import org.nd4j.linalg.dataset.api.iterator.DataSetIterator;
 import org.nd4j.linalg.dataset.api.iterator.MultiDataSetIterator;
+import org.nd4j.linalg.exception.ND4JArraySizeException;
 import org.nd4j.linalg.factory.Nd4j;
 import org.nd4j.linalg.heartbeat.Heartbeat;
 import org.nd4j.linalg.heartbeat.reports.Environment;
@@ -425,7 +426,8 @@ public class MultiLayerNetwork implements Serializable, Classifier, Layer, Neura
         try(MemoryWorkspace ws = workspaceMgr.notifyScopeEntered(ArrayType.FF_WORKING_MEM)) {
             if (layerWiseConfigurations.getInputPreProcess(layerIdx) != null) {
 
-                // FIXME: int cast
+                if (input.size(0) > Integer.MAX_VALUE)
+                    throw new ND4JArraySizeException();
                 outputOfPrevLayer = layerWiseConfigurations.getInputPreProcess(layerIdx).preProcess(outputOfPrevLayer, (int) input.size(0),
                         LayerWorkspaceMgr.noWorkspaces(helperWorkspaces));
             }
@@ -439,7 +441,8 @@ public class MultiLayerNetwork implements Serializable, Classifier, Layer, Neura
         //In 99+% of cases, the input and labels dimension 0 size should be identical
         //The only real exceptions: space to batch, and batch to space layers
         //In those cases, we should base it on the labels size, as this impacts gradient calculation
-        // FIXME: int cast
+        if (input.size(0) > Integer.MAX_VALUE || labels.size(0) > Integer.MAX_VALUE)
+            throw new ND4JArraySizeException();
         return labels == null ? (int) input.size(0) : (int)labels.size(0);
     }
 
@@ -2074,7 +2077,8 @@ public class MultiLayerNetwork implements Serializable, Classifier, Layer, Neura
             if (endTimeIdx > timeSeriesLength)
                 endTimeIdx = timeSeriesLength;
 
-            // FIXME: int cast
+            if (startTimeIdx > Integer.MAX_VALUE || endTimeIdx > Integer.MAX_VALUE)
+                throw new ND4JArraySizeException();
             INDArray[] subsets = getSubsetsForTbptt((int) startTimeIdx, (int) endTimeIdx, input, labels,
                     featuresMaskArray, labelsMaskArray);
 
@@ -2211,7 +2215,9 @@ public class MultiLayerNetwork implements Serializable, Classifier, Layer, Neura
     public int[] predict(INDArray d) {
         INDArray output = output(d, Layer.TrainingMode.TEST);
 
-        // FIXME: int cast
+        if (d.size(0) > Integer.MAX_VALUE)
+            throw new ND4JArraySizeException();
+
         int[] ret = new int[(int) d.size(0)];
         if (d.isRowVectorOrScalar())
             ret[0] = Nd4j.getBlasWrapper().iamax(output);
@@ -2335,7 +2341,8 @@ public class MultiLayerNetwork implements Serializable, Classifier, Layer, Neura
         org.deeplearning4j.nn.conf.layers.OutputLayer layerConf =
                 (org.deeplearning4j.nn.conf.layers.OutputLayer) getOutputLayer().conf().getLayer();
 
-        // FIXME: int cast
+        if (layerConf.getNOut() > Integer.MAX_VALUE)
+            throw new ND4JArraySizeException();
         fit(examples, FeatureUtil.toOutcomeMatrix(labels, (int) layerConf.getNOut()));
     }
 
@@ -2584,7 +2591,8 @@ public class MultiLayerNetwork implements Serializable, Classifier, Layer, Neura
         INDArray inputToOutputLayer = outputOfLayerDetached(training, FwdPassType.STANDARD,layers.length-2, data.getFeatures(),
                 data.getFeaturesMaskArray(), data.getLabelsMaskArray(), null);
 
-        // FIXME: int cast
+        if (data.getFeatures().size(0) > Integer.MAX_VALUE)
+            throw new ND4JArraySizeException();
         IOutputLayer ol = (IOutputLayer) getOutputLayer();
         if (getLayerWiseConfigurations().getInputPreProcess(layers.length - 1) != null) {
             inputToOutputLayer = getLayerWiseConfigurations().getInputPreProcess(layers.length - 1)
@@ -2647,7 +2655,8 @@ public class MultiLayerNetwork implements Serializable, Classifier, Layer, Neura
             IOutputLayer ol = (IOutputLayer) getOutputLayer();
             if(layerWiseConfigurations.getInputPreProcess(layers.length-1) != null){
 
-                // FIXME: int cast
+                if (data.getFeatures().size(0) > Integer.MAX_VALUE)
+                    throw new ND4JArraySizeException();
                 inputLast = layerWiseConfigurations.getInputPreProcess(layers.length-1).preProcess(inputLast,
                         (int) data.getFeatures().size(0), mgr);
             }
@@ -2811,7 +2820,8 @@ public class MultiLayerNetwork implements Serializable, Classifier, Layer, Neura
                 throw new IllegalArgumentException(
                         "Invalid input: length 0 (shape: " + Arrays.toString(input.shape()) + ")");
 
-            // FIXME: int cast
+            if (input.size(0) > Integer.MAX_VALUE)
+                throw new ND4JArraySizeException();
             setInputMiniBatchSize((int) input.size(0));
         }
     }
@@ -3086,7 +3096,8 @@ public class MultiLayerNetwork implements Serializable, Classifier, Layer, Neura
         if(!conf().isMiniBatch())
             return 1;
 
-        // FIXME: int cast
+        if (input.size(0) > Integer.MAX_VALUE)
+            throw new ND4JArraySizeException();
         return (int) input.size(0);
     }
 
@@ -3256,7 +3267,8 @@ public class MultiLayerNetwork implements Serializable, Classifier, Layer, Neura
     public void setLayerMaskArrays(INDArray featuresMaskArray, INDArray labelsMaskArray) {
         if (featuresMaskArray != null) {
 
-            // FIXME: int cast
+            if (featuresMaskArray.size(0) > Integer.MAX_VALUE)
+                throw new ND4JArraySizeException();
             //New approach: use feedForwardMaskArray method
             feedForwardMaskArray(featuresMaskArray, MaskState.Active, (int) featuresMaskArray.size(0));
 
@@ -3438,7 +3450,8 @@ public class MultiLayerNetwork implements Serializable, Classifier, Layer, Neura
                     val startTimeIdx = i * fwdLen;
                     val endTimeIdx = Math.min(startTimeIdx + fwdLen, tsLength);
 
-                    // FIXME: int cast
+                    if (endTimeIdx > Integer.MAX_VALUE)
+                        throw new ND4JArraySizeException();
                     INDArray[] subsets = getSubsetsForTbptt(startTimeIdx, (int) endTimeIdx, features, labels, fMask, lMask);
 
                     setLayerMaskArrays(subsets[2], subsets[3]);
@@ -3943,7 +3956,8 @@ public class MultiLayerNetwork implements Serializable, Classifier, Layer, Neura
         }
         FeedForwardLayer ffl = (FeedForwardLayer) conf;
 
-        // FIXME: int cast
+        if (ffl.getNOut() > Integer.MAX_VALUE)
+            throw new ND4JArraySizeException();
         return (int) ffl.getNOut();
     }
 
@@ -3969,7 +3983,8 @@ public class MultiLayerNetwork implements Serializable, Classifier, Layer, Neura
         }
         FeedForwardLayer ffl = (FeedForwardLayer) conf;
 
-        // FIXME: int cast
+        if (ffl.getNIn() > Integer.MAX_VALUE)
+            throw new ND4JArraySizeException();
         return (int) ffl.getNIn();
     }
 

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/params/VariationalAutoencoderParamInitializer.java

@@ -22,6 +22,7 @@ import org.deeplearning4j.nn.conf.layers.Layer;
 import org.deeplearning4j.nn.conf.layers.variational.VariationalAutoencoder;
 import org.deeplearning4j.nn.weights.IWeightInit;
 import org.nd4j.linalg.api.ndarray.INDArray;
+import org.nd4j.linalg.exception.ND4JArraySizeException;
 import org.nd4j.linalg.indexing.NDArrayIndex;
 
 import java.util.ArrayList;
@@ -108,7 +109,8 @@ public class VariationalAutoencoderParamInitializer extends DefaultParamInitiali
         }
 
         //Between last decoder layer and parameters for p(x|z):
-        // FIXME: int cast
+        if (nIn > Integer.MAX_VALUE)
+            throw new ND4JArraySizeException();
         val nDistributionParams = layer.getOutputDistribution().distributionInputSize((int) nIn);
         val lastDecLayerSize = decoderLayerSizes[decoderLayerSizes.length - 1];
         paramCount += (lastDecLayerSize + 1) * nDistributionParams;
@@ -294,7 +296,8 @@ public class VariationalAutoencoderParamInitializer extends DefaultParamInitiali
         }
 
         //Finally, p(x|z):
-        // FIXME: int cast
+        if (nIn > Integer.MAX_VALUE)
+            throw new ND4JArraySizeException();
         int nDistributionParams = layer.getOutputDistribution().distributionInputSize((int) nIn);
         int pxzWeightCount = decoderLayerSizes[decoderLayerSizes.length - 1] * nDistributionParams;
         INDArray pxzWeightView =
@@ -402,7 +405,8 @@ public class VariationalAutoencoderParamInitializer extends DefaultParamInitiali
         }
 
         //Finally, p(x|z):
-        // FIXME: int cast
+        if (nIn > Integer.MAX_VALUE)
+            throw new ND4JArraySizeException();
         int nDistributionParams = layer.getOutputDistribution().distributionInputSize((int) nIn);
         int pxzWeightCount = decoderLayerSizes[decoderLayerSizes.length - 1] * nDistributionParams;
         INDArray pxzWeightView =

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/updater/BaseMultiLayerUpdater.java

@@ -30,6 +30,7 @@ import org.nd4j.linalg.api.ops.CustomOp;
 import org.nd4j.linalg.api.ops.DynamicCustomOp;
 
 import org.nd4j.linalg.api.ops.impl.reduce.floating.Norm2;
+import org.nd4j.linalg.exception.ND4JArraySizeException;
 import org.nd4j.linalg.factory.Nd4j;
 import org.nd4j.linalg.indexing.NDArrayIndex;
 import org.deeplearning4j.nn.workspace.ArrayType;
@@ -111,7 +112,8 @@ public abstract class BaseMultiLayerUpdater<T extends Model> implements Updater
                     if (currentBlock == null || !UpdaterUtils.updaterConfigurationsEquals(lastLayer, lastVariable,
                                     layers[i], var)) {
 
-                        // FIXME: int cast
+                        if (paramsViewSoFar + paramSizeThisVariable > Integer.MAX_VALUE || paramsViewSoFar + paramSizeThisVariable > Integer.MAX_VALUE)
+                            throw new ND4JArraySizeException();
                         //Create a new block
                         List<UpdaterBlock.ParamState> list = new ArrayList<>();
                         list.add(new UpdaterBlock.ParamState(layers[i], var, paramsViewSoFar,
@@ -122,9 +124,11 @@ public abstract class BaseMultiLayerUpdater<T extends Model> implements Updater
 
                         updaterBlocks.add(currentBlock);
                     } else {
-                        // FIXME: int cast
+                        long newOffset = currentBlock.getParamOffsetEnd() + paramSizeThisVariable;
+                        if (newOffset > Integer.MAX_VALUE)
+                            throw new ND4JArraySizeException();
                         //Add to existing updater block
-                        currentBlock.setParamOffsetEnd((int) (currentBlock.getParamOffsetEnd() + paramSizeThisVariable));
+                        currentBlock.setParamOffsetEnd((int) newOffset);
                         currentBlock.setUpdaterViewOffsetEnd(
                                         currentBlock.getUpdaterViewOffsetEnd() + updaterStateSizeThisVariable);
                         currentBlock.getLayersAndVariablesInBlock()

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/util/Convolution1DUtils.java

@@ -29,6 +29,7 @@ import org.nd4j.base.Preconditions;
 import org.nd4j.linalg.api.ndarray.INDArray;
 import org.nd4j.linalg.api.ops.impl.broadcast.BroadcastCopyOp;
 import org.nd4j.linalg.api.shape.Shape;
+import org.nd4j.linalg.exception.ND4JArraySizeException;
 import org.nd4j.linalg.factory.Nd4j;
 
 import java.util.Arrays;
@@ -62,10 +63,9 @@ public class Convolution1DUtils {
      * @param dilation        Kernel dilation
      * @return Output size (width)
      */
-    public static int getOutputSize(int inH, int kernel, int strides, int padding,
+    public static long getOutputSize(long inH, int kernel, int strides, int padding,
                                     ConvolutionMode convolutionMode, int dilation) {
-        // FIXME: int cast
+        long eKernel = effectiveKernelSize(kernel, dilation);
-        int eKernel = effectiveKernelSize(kernel, dilation);
         if (convolutionMode == ConvolutionMode.Same) {
             return (int) Math.ceil(inH / ((double) strides));
         }
@@ -85,7 +85,8 @@ public class Convolution1DUtils {
      */
     public static int getOutputSize(INDArray inputData, int kernel, int strides, int padding,
                                     ConvolutionMode convolutionMode, int dilation) {
-        // FIXME: int cast
+        if (inputData.size(2) > Integer.MAX_VALUE)
+            throw new ND4JArraySizeException();
         int inH = (int) inputData.size(2);
         int eKernel = effectiveKernelSize(kernel, dilation);
         boolean atrous = (eKernel == kernel);

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/util/Convolution3DUtils.java

@@ -61,15 +61,14 @@ public class Convolution3DUtils {
                                         ConvolutionMode convolutionMode, int[] dilation, boolean isNCDHW) {
 
         // NCDHW vs. NDHWC
-        int inD = (int) (isNCDHW ? inputData.size(2) : inputData.size(1));
+        long inD = (isNCDHW ? inputData.size(2) : inputData.size(1));
-        int inH = (int) (isNCDHW ? inputData.size(3) : inputData.size(2));
+        long inH = (isNCDHW ? inputData.size(3) : inputData.size(2));
-        int inW = (int) (isNCDHW ? inputData.size(4) : inputData.size(3));
+        long inW = (isNCDHW ? inputData.size(4) : inputData.size(3));
 
         int[] eKernel = effectiveKernelSize(kernel, dilation);
         boolean atrous = (eKernel == kernel);
 
-        // FIXME: int cast
+        val inShape = new long[]{inD, inH, inW};
-        val inShape = new int[]{inD, inH, inW};
         validateShapes(ArrayUtil.toInts(inputData.shape()), eKernel, strides, padding, convolutionMode, dilation, inShape, atrous);
 
         if (convolutionMode == ConvolutionMode.Same) {
@@ -80,16 +79,16 @@ public class Convolution3DUtils {
             return new int[]{outD, outH, outW};
         }
 
-        int outD = (inD - eKernel[0] + 2 * padding[0]) / strides[0] + 1;
+        int outD = ((int)inD - eKernel[0] + 2 * padding[0]) / strides[0] + 1;
-        int outH = (inH - eKernel[1] + 2 * padding[1]) / strides[1] + 1;
+        int outH = ((int)inH - eKernel[1] + 2 * padding[1]) / strides[1] + 1;
-        int outW = (inW - eKernel[2] + 2 * padding[2]) / strides[2] + 1;
+        int outW = ((int)inW - eKernel[2] + 2 * padding[2]) / strides[2] + 1;
 
         return new int[]{outD, outH, outW};
     }
 
 
     private static void validateShapes(int[] inputDataShape, int[] eKernel, int[] strides, int[] padding,
-                                      ConvolutionMode convolutionMode, int[] dilation, int[] inShape,
+                                      ConvolutionMode convolutionMode, int[] dilation, long[] inShape,
                                       boolean atrous) {
 
         String[] dims = new String[]{"depth", "height", "width"};

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/util/ConvolutionUtils.java

@@ -36,6 +36,8 @@ import org.nd4j.linalg.api.ops.impl.broadcast.BroadcastCopyOp;
 import org.nd4j.linalg.api.ops.impl.layers.convolution.MaxPooling2D;
 import org.nd4j.linalg.api.ops.impl.layers.convolution.config.Pooling2DConfig;
 import org.nd4j.linalg.api.shape.Shape;
+import org.nd4j.linalg.exception.ND4JArraySizeException;
+import org.nd4j.linalg.factory.NDArrayFactory;
 import org.nd4j.linalg.factory.Nd4j;
 
 import java.util.Arrays;
@@ -73,7 +75,8 @@ public class ConvolutionUtils {
     public static int[] getDeconvolutionOutputSize(INDArray inputData, int[] kernel, int[] strides, int[] padding,
                                                    ConvolutionMode convolutionMode, int[] dilation) {
 
-        // FIXME: int cast
+        if (inputData.size(2) > Integer.MAX_VALUE || inputData.size(3) > Integer.MAX_VALUE)
+            throw new ND4JArraySizeException();
         int hIn = (int) inputData.size(2);
         int wIn = (int) inputData.size(3);
         int[] eKernel = effectiveKernelSize(kernel, dilation);
@@ -104,7 +107,8 @@ public class ConvolutionUtils {
      */
     public static int[] getOutputSize(INDArray inputData, int[] kernel, int[] strides, int[] padding,
                                       ConvolutionMode convolutionMode, int[] dilation) {
-        // FIXME: int cast
+        if (inputData.size(2) > Integer.MAX_VALUE || inputData.size(3) > Integer.MAX_VALUE)
+            throw new ND4JArraySizeException();
         int inH = (int) inputData.size(2);
         int inW = (int) inputData.size(3);
 
@@ -499,7 +503,7 @@ public class ConvolutionUtils {
         }
     }
 
-    public static INDArray reshape2dTo4d(INDArray in2d, int[] toShape, LayerWorkspaceMgr workspaceMgr, ArrayType type){
+    public static INDArray reshape2dTo4d(INDArray in2d, long[] toShape, LayerWorkspaceMgr workspaceMgr, ArrayType type){
         if(in2d.rank() != 2)
             throw new IllegalArgumentException("Invalid input: expect NDArray with rank 2");
         if (toShape.length != 4)
@@ -513,7 +517,7 @@ public class ConvolutionUtils {
         return workspaceMgr.leverageTo(type, out.permute(0, 3, 1, 2));
     }
 
-    public static INDArray reshape2dTo5d(Convolution3D.DataFormat format, INDArray in2d, int n, int d, int h, int w, int ch, LayerWorkspaceMgr workspaceMgr, ArrayType type){
+    public static INDArray reshape2dTo5d(Convolution3D.DataFormat format, INDArray in2d, long n, long d, long h, long w, long ch, LayerWorkspaceMgr workspaceMgr, ArrayType type){
         if(in2d.rank() != 2)
             throw new IllegalArgumentException("Invalid input: expect NDArray with rank 2");
 
@@ -580,14 +584,21 @@ public class ConvolutionUtils {
         int inW;
         int inDepth;
 
-        // FIXME: int cast
         if (inputType instanceof InputType.InputTypeConvolutional) {
             InputType.InputTypeConvolutional conv = (InputType.InputTypeConvolutional) inputType;
+            if (conv.getHeight() > Integer.MAX_VALUE || conv.getWidth() > Integer.MAX_VALUE ||
+                conv.getChannels() > Integer.MAX_VALUE){
+                throw new ND4JArraySizeException();
+            }
             inH = (int) conv.getHeight();
             inW = (int) conv.getWidth();
             inDepth = (int) conv.getChannels();
         } else if (inputType instanceof InputType.InputTypeConvolutionalFlat) {
             InputType.InputTypeConvolutionalFlat conv = (InputType.InputTypeConvolutionalFlat) inputType;
+            if (conv.getHeight() > Integer.MAX_VALUE || conv.getWidth() > Integer.MAX_VALUE ||
+                conv.getDepth() > Integer.MAX_VALUE) {
+                throw new ND4JArraySizeException();
+            }
             inH = (int) conv.getHeight();
             inW = (int) conv.getWidth();
             inDepth = (int) conv.getDepth();

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/util/TimeSeriesUtils.java

@@ -20,6 +20,7 @@ import lombok.val;
 import org.nd4j.base.Preconditions;
 import org.nd4j.linalg.api.ndarray.INDArray;
 import org.nd4j.linalg.api.shape.Shape;
+import org.nd4j.linalg.exception.ND4JArraySizeException;
 import org.nd4j.linalg.factory.Nd4j;
 import org.nd4j.linalg.indexing.BooleanIndexing;
 import org.nd4j.linalg.indexing.INDArrayIndex;
@@ -193,7 +194,7 @@ public class TimeSeriesUtils {
     }
 
 
-    public static INDArray reshape2dTo3d(INDArray in, int miniBatchSize, LayerWorkspaceMgr workspaceMgr, ArrayType arrayType) {
+    public static INDArray reshape2dTo3d(INDArray in, long miniBatchSize, LayerWorkspaceMgr workspaceMgr, ArrayType arrayType) {
         if (in.rank() != 2)
             throw new IllegalArgumentException("Invalid input: expect NDArray with rank 2");
         //Based on: RnnToFeedForwardPreProcessor
@@ -220,7 +221,6 @@ public class TimeSeriesUtils {
             in = in.dup('f');
         }
 
-        // FIXME: int cast
         int[] idxs = new int[(int) in.size(2)];
         int j=0;
         for( int i=idxs.length-1; i>=0; i--){
@@ -248,7 +248,8 @@ public class TimeSeriesUtils {
             in = workspaceMgr.dup(arrayType, in, 'f');
         }
 
-        // FIXME: int cast
+        if (in.size(2) > Integer.MAX_VALUE)
+            throw new ND4JArraySizeException();
         int[] idxs = new int[(int) in.size(2)];
         int j=0;
         for( int i=idxs.length-1; i>=0; i--){
@@ -291,7 +292,8 @@ public class TimeSeriesUtils {
                     + " with shape " + Arrays.toString(mask.shape()));
         }
 
-        // FIXME: int cast
+        if (mask.size(1) > Integer.MAX_VALUE)
+            throw new ND4JArraySizeException();
         int[] idxs = new int[(int) mask.size(1)];
         int j=0;
         for( int i=idxs.length-1; i>=0; i--){
@@ -319,7 +321,8 @@ public class TimeSeriesUtils {
                     + " with shape " + Arrays.toString(mask.shape()));
         }
 
-        // FIXME: int cast
+        if (mask.size(1) > Integer.MAX_VALUE)
+            throw new ND4JArraySizeException();
         int[] idxs = new int[(int) mask.size(1)];
         int j=0;
         for( int i=idxs.length-1; i>=0; i--){
@@ -358,9 +361,8 @@ public class TimeSeriesUtils {
         INDArray out;
         if (mask == null) {
 
-            // FIXME: int cast
             //No mask array -> extract same (last) column for all
-            int lastTS = (int) pullFrom.size(2) - 1;
+            long lastTS = pullFrom.size(2) - 1;
             out = pullFrom.get(NDArrayIndex.all(), NDArrayIndex.all(), NDArrayIndex.point(lastTS));
             fwdPassTimeSteps = null; //Null -> last time step for all examples
         } else {
@@ -396,9 +398,8 @@ public class TimeSeriesUtils {
         INDArray out;
         if (mask == null) {
 
-            // FIXME: int cast
             //No mask array -> extract same (last) column for all
-            int lastTS = (int) pullFrom.size(2) - 1;
+            long lastTS = pullFrom.size(2) - 1;
             out = pullFrom.get(NDArrayIndex.all(), NDArrayIndex.all(), NDArrayIndex.point(lastTS));
             fwdPassTimeSteps = null; //Null -> last time step for all examples
         } else {

deeplearning4j/deeplearning4j-nn/src/test/java/org/deeplearning4j/TestUtils.java

@@ -116,7 +116,6 @@ public class TestUtils {
     public static INDArray randomOneHot(long examples, long nOut, Random rng){
         INDArray arr = Nd4j.create(examples, nOut);
         for( int i=0; i<examples; i++ ){
-            // FIXME: int cast
             arr.putScalar(i, rng.nextInt((int) nOut), 1.0);
         }
         return arr;

deeplearning4j/deeplearning4j-scaleout/spark/dl4j-spark-nlp/src/main/java/org/deeplearning4j/spark/models/embeddings/word2vec/FirstIterationFunction.java

@@ -214,7 +214,6 @@ public class FirstIterationFunction implements
                 else {
                     nextRandom.set(Math.abs(nextRandom.get() * 25214903917L + 11));
 
-                    // FIXME: int cast
                     int idx = Math.abs((int) (nextRandom.get() >> 16) % (int) negativeHolder.getTable().length());
 
                     target = negativeHolder.getTable().getInt(idx);

deeplearning4j/deeplearning4j-scaleout/spark/dl4j-spark-nlp/src/main/java/org/deeplearning4j/spark/models/embeddings/word2vec/SecondIterationFunction.java

@@ -222,7 +222,6 @@ public class SecondIterationFunction implements FlatMapFunction<Iterator<Tuple2<
                 else {
                     nextRandom.set(Math.abs(nextRandom.get() * 25214903917L + 11));
 
-                    // FIXME: int cast
                     int idx = (int) Math.abs((int) (nextRandom.get() >> 16) % negativeHolder.getTable().length());
 
                     target = negativeHolder.getTable().getInt(idx);

deeplearning4j/deeplearning4j-scaleout/spark/dl4j-spark-nlp/src/main/java/org/deeplearning4j/spark/models/embeddings/word2vec/SentenceBatch.java

@@ -162,7 +162,6 @@ public class SentenceBatch implements Function<Word2VecFuncCall, Word2VecChange>
                     label = 1;
                 } else {
                     nextRandom.set(nextRandom.get() * 25214903917L + 11);
-                    // FIXME: int cast
                     target = table.getInt((int) (nextRandom.get() >> 16) % (int) table.length());
                     if (target == 0)
                         target = (int) nextRandom.get() % (numWords - 1) + 1;

deeplearning4j/deeplearning4j-scaleout/spark/dl4j-spark-nlp/src/main/java/org/deeplearning4j/spark/models/embeddings/word2vec/Word2VecPerformer.java

@@ -187,7 +187,6 @@ public class Word2VecPerformer implements VoidFunction<Pair<List<VocabWord>, Ato
                 } else {
                     nextRandom.set(nextRandom.get() * 25214903917L + 11);
 
-                    // FIXME: int cast
                     target = table.getInt((int) (nextRandom.get() >> 16) % (int) table.length());
                     if (target == 0)
                         target = (int) nextRandom.get() % (numWords - 1) + 1;

deeplearning4j/deeplearning4j-scaleout/spark/dl4j-spark-nlp/src/main/java/org/deeplearning4j/spark/models/embeddings/word2vec/Word2VecPerformerVoid.java

@@ -337,7 +337,6 @@ public class Word2VecPerformerVoid implements VoidFunction<Pair<List<VocabWord>,
                     label = 1;
                 } else {
                     nextRandom.set(nextRandom.get() * 25214903917L + 11);
-                    // FIXME: int cast
                     target = table.getInt((int) (nextRandom.get() >> 16) % (int) table.length());
                     if (target == 0)
                         target = (int) nextRandom.get() % (numWords - 1) + 1;

deeplearning4j/deeplearning4j-scaleout/spark/dl4j-spark/src/main/java/org/deeplearning4j/spark/api/stats/StatsCalculationHelper.java

@@ -39,7 +39,7 @@ public class StatsCalculationHelper {
     private long initialModelAfter;
     private long lastDataSetBefore;
     private long lastProcessBefore;
-    private int totalExampleCount;
+    private long totalExampleCount;
     private List<EventStats> dataSetGetTimes = new ArrayList<>();
     private List<EventStats> processMiniBatchTimes = new ArrayList<>();
 
@@ -65,7 +65,7 @@ public class StatsCalculationHelper {
         lastDataSetBefore = timeSource.currentTimeMillis();
     }
 
-    public void logNextDataSetAfter(int numExamples) {
+    public void logNextDataSetAfter(long numExamples) {
         long now = timeSource.currentTimeMillis();
         long duration = now - lastDataSetBefore;
         dataSetGetTimes.add(new BaseEventStats(lastDataSetBefore, duration));

deeplearning4j/deeplearning4j-scaleout/spark/dl4j-spark/src/main/java/org/deeplearning4j/spark/api/worker/ExecuteWorkerMultiDataSetFlatMap.java

@@ -84,9 +84,8 @@ public class ExecuteWorkerMultiDataSetFlatMap<R extends TrainingResult> implemen
                     s.logNextDataSetBefore();
                 MultiDataSet next = batchedIterator.next();
 
-                // FIXME: int cast
                 if (stats)
-                    s.logNextDataSetAfter((int) next.getFeatures(0).size(0));
+                    s.logNextDataSetAfter(next.getFeatures(0).size(0));
 
                 if (stats) {
                     s.logProcessMinibatchBefore();

deeplearning4j/deeplearning4j-scaleout/spark/dl4j-spark/src/main/java/org/deeplearning4j/spark/impl/graph/scoring/GraphFeedForwardWithKeyFunction.java

@@ -71,7 +71,7 @@ public class GraphFeedForwardWithKeyFunction<K> implements PairFlatMapFunction<I
 
         List<INDArray[]> featuresList = new ArrayList<>(batchSize);
         List<K> keyList = new ArrayList<>(batchSize);
-        List<Integer> origSizeList = new ArrayList<>();
+        List<Long> origSizeList = new ArrayList<>();
 
         long[][] firstShapes = null;
         boolean sizesDiffer = false;
@@ -96,8 +96,7 @@ public class GraphFeedForwardWithKeyFunction<K> implements PairFlatMapFunction<I
             featuresList.add(t2._2());
             keyList.add(t2._1());
 
-            // FIXME: int cast
+            origSizeList.add(t2._2()[0].size(0));
-            origSizeList.add((int) t2._2()[0].size(0));
             tupleCount++;
         }
 
@@ -156,7 +155,7 @@ public class GraphFeedForwardWithKeyFunction<K> implements PairFlatMapFunction<I
 
             examplesInBatch = 0;
             for (int i = firstIdx; i < nextIdx; i++) {
-                int numExamples = origSizeList.get(i);
+                long numExamples = origSizeList.get(i);
                 INDArray[] outSubset = new INDArray[out.length];
                 for (int j = 0; j < out.length; j++) {
                     outSubset[j] = getSubset(examplesInBatch, examplesInBatch + numExamples, out[j]);
@@ -174,7 +173,7 @@ public class GraphFeedForwardWithKeyFunction<K> implements PairFlatMapFunction<I
         return output.iterator();
     }
 
-    private INDArray getSubset(int exampleStart, int exampleEnd, INDArray from) {
+    private INDArray getSubset(long exampleStart, long exampleEnd, INDArray from) {
         switch (from.rank()) {
             case 2:
                 return from.get(NDArrayIndex.interval(exampleStart, exampleEnd), NDArrayIndex.all());

deeplearning4j/deeplearning4j-scaleout/spark/dl4j-spark/src/main/java/org/deeplearning4j/spark/impl/graph/scoring/ScoreFlatMapFunctionCGDataSet.java

@@ -36,7 +36,7 @@ import java.util.Iterator;
 import java.util.List;
 
 /** Function used to score a DataSet using a ComputationGraph */
-public class ScoreFlatMapFunctionCGDataSet implements FlatMapFunction<Iterator<DataSet>, Tuple2<Integer, Double>> {
+public class ScoreFlatMapFunctionCGDataSet implements FlatMapFunction<Iterator<DataSet>, Tuple2<Long, Double>> {
     private static final Logger log = LoggerFactory.getLogger(ScoreFlatMapFunctionCGDataSet.class);
     private String json;
     private Broadcast<INDArray> params;
@@ -50,9 +50,9 @@ public class ScoreFlatMapFunctionCGDataSet implements FlatMapFunction<Iterator<D
     }
 
     @Override
-    public Iterator<Tuple2<Integer, Double>> call(Iterator<DataSet> dataSetIterator) throws Exception {
+    public Iterator<Tuple2<Long, Double>> call(Iterator<DataSet> dataSetIterator) throws Exception {
         if (!dataSetIterator.hasNext()) {
-            return Collections.singletonList(new Tuple2<>(0, 0.0)).iterator();
+            return Collections.singletonList(new Tuple2<>(0L, 0.0)).iterator();
         }
 
         DataSetIterator iter = new IteratorDataSetIterator(dataSetIterator, minibatchSize); //Does batching where appropriate
@@ -65,13 +65,12 @@ public class ScoreFlatMapFunctionCGDataSet implements FlatMapFunction<Iterator<D
                             "Network did not have same number of parameters as the broadcast set parameters");
         network.setParams(val);
 
-        List<Tuple2<Integer, Double>> out = new ArrayList<>();
+        List<Tuple2<Long, Double>> out = new ArrayList<>();
         while (iter.hasNext()) {
             DataSet ds = iter.next();
             double score = network.score(ds, false);
 
-            // FIXME: int cast
+            long numExamples = ds.getFeatures().size(0);
-            int numExamples = (int) ds.getFeatures().size(0);
             out.add(new Tuple2<>(numExamples, score * numExamples));
         }
 

deeplearning4j/deeplearning4j-scaleout/spark/dl4j-spark/src/main/java/org/deeplearning4j/spark/impl/graph/scoring/ScoreFlatMapFunctionCGMultiDataSet.java

@@ -35,7 +35,7 @@ import java.util.Iterator;
 import java.util.List;
 
 /** Function used to score a MultiDataSet using a given ComputationGraph */
-public class ScoreFlatMapFunctionCGMultiDataSet implements FlatMapFunction<Iterator<MultiDataSet>, Tuple2<Integer, Double>> {
+public class ScoreFlatMapFunctionCGMultiDataSet implements FlatMapFunction<Iterator<MultiDataSet>, Tuple2<Long, Double>> {
 
     private static final Logger log = LoggerFactory.getLogger(ScoreFlatMapFunctionCGMultiDataSet.class);
     private String json;
@@ -50,9 +50,9 @@ public class ScoreFlatMapFunctionCGMultiDataSet implements FlatMapFunction<Itera
     }
 
     @Override
-    public Iterator<Tuple2<Integer, Double>> call(Iterator<MultiDataSet> dataSetIterator) throws Exception {
+    public Iterator<Tuple2<Long, Double>> call(Iterator<MultiDataSet> dataSetIterator) throws Exception {
         if (!dataSetIterator.hasNext()) {
-            return Collections.singletonList(new Tuple2<>(0, 0.0)).iterator();
+            return Collections.singletonList(new Tuple2<>(0L, 0.0)).iterator();
         }
 
         MultiDataSetIterator iter = new IteratorMultiDataSetIterator(dataSetIterator, minibatchSize); //Does batching where appropriate
@@ -66,13 +66,12 @@ public class ScoreFlatMapFunctionCGMultiDataSet implements FlatMapFunction<Itera
                             "Network did not have same number of parameters as the broadcast set parameters");
         network.setParams(val);
 
-        List<Tuple2<Integer, Double>> out = new ArrayList<>();
+        List<Tuple2<Long, Double>> out = new ArrayList<>();
         while (iter.hasNext()) {
             MultiDataSet ds = iter.next();
             double score = network.score(ds, false);
 
-            // FIXME: int cast
+            long numExamples = ds.getFeatures(0).size(0);
-            int numExamples = (int) ds.getFeatures(0).size(0);
             out.add(new Tuple2<>(numExamples, score * numExamples));
         }
 

deeplearning4j/deeplearning4j-scaleout/spark/dl4j-spark/src/main/java/org/deeplearning4j/spark/impl/multilayer/scoring/FeedForwardWithKeyFunction.java

@@ -105,7 +105,6 @@ public class FeedForwardWithKeyFunction<K>
             fMaskList.add(t2._2()._2());
             keyList.add(t2._1());
 
-            // FIXME: int cast
             origSizeList.add((int) t2._2()._1().size(0));
             tupleCount++;
         }

deeplearning4j/deeplearning4j-scaleout/spark/dl4j-spark/src/main/java/org/deeplearning4j/spark/impl/multilayer/scoring/ScoreFlatMapFunction.java

@@ -64,7 +64,6 @@ public class ScoreFlatMapFunction implements FlatMapFunction<Iterator<DataSet>,
             DataSet ds = iter.next();
             double score = network.score(ds, false);
 
-            // FIXME: int cast
             val numExamples = (int) ds.getFeatures().size(0);
             out.add(new Tuple2<>(numExamples, score * numExamples));
         }

deeplearning4j/deeplearning4j-scaleout/spark/dl4j-spark/src/main/java/org/deeplearning4j/spark/impl/paramavg/ParameterAveragingTrainingWorker.java

@@ -247,10 +247,8 @@ public class ParameterAveragingTrainingWorker extends BaseTrainingWorker<Paramet
                 trainingHook.postUpdate(dataSet, graph);
             }
         }
-
-        // FIXME: int cast
         if (configuration.isCollectTrainingStats())
-            stats.logFitEnd((int) dataSet.getFeatures(0).size(0));
+            stats.logFitEnd(dataSet.getFeatures(0).size(0));
 
         Nd4j.getExecutioner().commit();
 

deeplearning4j/deeplearning4j-scaleout/spark/dl4j-spark/src/main/java/org/deeplearning4j/spark/impl/paramavg/stats/ParameterAveragingTrainingWorkerStats.java

@@ -195,7 +195,7 @@ public class ParameterAveragingTrainingWorkerStats implements SparkTrainingStats
             lastFitStartTime = timeSource.currentTimeMillis();
         }
 
-        public void logFitEnd(int numExamples) {
+        public void logFitEnd(long numExamples) {
             long now = timeSource.currentTimeMillis();
             fitTimes.add(new ExampleCountEventStats(lastFitStartTime, now - lastFitStartTime, numExamples));
         }

deeplearning4j/deeplearning4j-scaleout/spark/dl4j-spark/src/main/java/org/deeplearning4j/spark/iterator/BaseDataSetIterator.java

@@ -19,6 +19,7 @@ package org.deeplearning4j.spark.iterator;
 import org.nd4j.linalg.dataset.DataSet;
 import org.nd4j.linalg.dataset.api.DataSetPreProcessor;
 import org.nd4j.linalg.dataset.api.iterator.DataSetIterator;
+import org.nd4j.linalg.exception.ND4JArraySizeException;
 
 import java.util.Collection;
 import java.util.Iterator;
@@ -31,8 +32,8 @@ public abstract class BaseDataSetIterator<T> implements DataSetIterator {
     protected Collection<T> dataSetStreams;
     protected DataSetPreProcessor preprocessor;
     protected Iterator<T> iter;
-    protected int totalOutcomes = -1;
+    protected long totalOutcomes = -1;
-    protected int inputColumns = -1;
+    protected long inputColumns = -1;
     protected int batch = -1;
     protected DataSet preloadedDataSet;
     protected int cursor = 0;
@@ -112,7 +113,9 @@ public abstract class BaseDataSetIterator<T> implements DataSetIterator {
     private void preloadDataSet() {
         preloadedDataSet = load(iter.next());
 
-        // FIXME: int cast
+        if (preloadedDataSet.getLabels().size(1) > Integer.MAX_VALUE ||
+            preloadedDataSet.getFeatures().size(1) > Integer.MAX_VALUE)
+            throw new ND4JArraySizeException();
         totalOutcomes = (int) preloadedDataSet.getLabels().size(1);
         inputColumns = (int) preloadedDataSet.getFeatures().size(1);
     }

deeplearning4j/deeplearning4j-scaleout/spark/dl4j-spark/src/main/java/org/deeplearning4j/spark/iterator/PathSparkDataSetIterator.java

@@ -67,7 +67,6 @@ public class PathSparkDataSetIterator extends BaseDataSetIterator<String> {
             ds = load(iter.next());
         }
 
-        // FIXME: int cast
         totalOutcomes = ds.getLabels() == null ? 0 : (int) ds.getLabels().size(1); //May be null for layerwise pretraining
         inputColumns = (int) ds.getFeatures().size(1);
         batch = ds.numExamples();

deeplearning4j/deeplearning4j-scaleout/spark/dl4j-spark/src/main/java/org/deeplearning4j/spark/iterator/PortableDataStreamDataSetIterator.java

@@ -18,6 +18,7 @@ package org.deeplearning4j.spark.iterator;
 
 import org.apache.spark.input.PortableDataStream;
 import org.nd4j.linalg.dataset.DataSet;
+import org.nd4j.linalg.exception.ND4JArraySizeException;
 
 import java.io.IOException;
 import java.io.InputStream;
@@ -53,7 +54,9 @@ public class PortableDataStreamDataSetIterator extends BaseDataSetIterator<Porta
             ds = load(iter.next());
         }
 
-        // FIXME: int cast
+        if (ds.getLabels().size(1) > Integer.MAX_VALUE ||
+            ds.getFeatures().size(1) > Integer.MAX_VALUE)
+            throw new ND4JArraySizeException();
         totalOutcomes = (int) ds.getLabels().size(1);
         inputColumns = (int) ds.getFeatures().size(1);
         batch = ds.numExamples();

deeplearning4j/deeplearning4j-scaleout/spark/dl4j-spark/src/main/java/org/deeplearning4j/spark/stats/ExampleCountEventStats.java

@@ -26,9 +26,9 @@ import lombok.Getter;
 public class ExampleCountEventStats extends BaseEventStats {
 
     @Getter
-    private final int totalExampleCount;
+    private final long totalExampleCount;
 
-    public ExampleCountEventStats(long startTime, long durationMs, int totalExampleCount) {
+    public ExampleCountEventStats(long startTime, long durationMs, long totalExampleCount) {
         super(startTime, durationMs);
         this.totalExampleCount = totalExampleCount;
     }

deeplearning4j/deeplearning4j-scaleout/spark/dl4j-spark/src/main/java/org/deeplearning4j/spark/util/MLLibUtil.java

@@ -31,6 +31,7 @@ import org.datavec.api.split.InputStreamInputSplit;
 import org.datavec.api.writable.Writable;
 import org.nd4j.linalg.api.ndarray.INDArray;
 import org.nd4j.linalg.dataset.DataSet;
+import org.nd4j.linalg.exception.ND4JArraySizeException;
 import org.nd4j.linalg.factory.Nd4j;
 import org.nd4j.linalg.util.FeatureUtil;
 import scala.Tuple2;
@@ -122,7 +123,8 @@ public class MLLibUtil {
         if (!arr.isVector()) {
             throw new IllegalArgumentException("passed in array must be a vector");
         }
-        // FIXME: int cast
+        if (arr.length() > Integer.MAX_VALUE)
+            throw new ND4JArraySizeException();
         double[] ret = new double[(int) arr.length()];
         for (int i = 0; i < arr.length(); i++) {
             ret[i] = arr.getDouble(i);

deeplearning4j/deeplearning4j-ui-parent/deeplearning4j-ui/src/main/java/org/deeplearning4j/ui/weights/ConvolutionalIterationListener.java

@@ -33,6 +33,7 @@ import org.deeplearning4j.ui.api.UIServer;
 import org.deeplearning4j.ui.storage.mapdb.MapDBStatsStorage;
 import org.deeplearning4j.util.UIDProvider;
 import org.nd4j.linalg.api.ndarray.INDArray;
+import org.nd4j.linalg.exception.ND4JArraySizeException;
 import org.nd4j.linalg.io.ClassPathResource;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -205,7 +206,8 @@ public class ConvolutionalIterationListener extends BaseTrainingListener {
                     if(layers[i].type() == Layer.Type.CONVOLUTIONAL){
                         INDArray output = activations.get(i+1); //Offset by 1 - activations list includes input
 
-                        // FIXME: int cast
+                        if (output.shape()[0] - 1 > Integer.MAX_VALUE)
+                            throw new ND4JArraySizeException();
                         int sampleDim = output.shape()[0] == 1 ? 0 : rnd.nextInt((int) output.shape()[0] - 1) + 1;
                         if (cnt == 0) {
                             INDArray inputs = layers[i].input();
@@ -426,7 +428,8 @@ public class ConvolutionalIterationListener extends BaseTrainingListener {
 
         val height = (numRows * (tShape[1] + border + padding_col)) + padding_col + zoomPadding + zoomWidth;
 
-        // FIXME: int cast
+        if (height > Integer.MAX_VALUE)
+            throw new ND4JArraySizeException();
         BufferedImage outputImage = new BufferedImage(maxWidth, (int) height, BufferedImage.TYPE_BYTE_GRAY);
         Graphics2D graphics2D = outputImage.createGraphics();
 
@@ -571,7 +574,8 @@ public class ConvolutionalIterationListener extends BaseTrainingListener {
             */
 
             graphics2D.setPaint(borderColor);
-            // FIXME: int cast
+            if (tad2D.shape()[0] > Integer.MAX_VALUE || tad2D.shape()[1] > Integer.MAX_VALUE)
+                throw new ND4JArraySizeException();
             graphics2D.drawRect(columnOffset, rowOffset, (int) tad2D.shape()[0], (int) tad2D.shape()[1]);
 
 

deeplearning4j/deeplearning4j-zoo/src/main/java/org/deeplearning4j/zoo/util/BaseLabels.java

@@ -90,9 +90,8 @@ public abstract class BaseLabels implements Labels {
         Preconditions.checkState(predictions.size(1) == labels.size(), "Invalid input array:" +
                 " expected array with size(1) equal to numLabels (%s), got array with shape %s", labels.size(), predictions.shape());
 
-        // FIXME: int cast
+        long rows = predictions.size(0);
-        int rows = (int) predictions.size(0);
+        long cols = predictions.size(1);
-        int cols = (int) predictions.size(1);
         if (predictions.isColumnVectorOrScalar()) {
             predictions = predictions.ravel();
             rows = (int) predictions.size(0);

deeplearning4j/dl4j-integration-tests/src/test/java/org/deeplearning4j/integration/TestUtils.java

@@ -116,7 +116,6 @@ public class TestUtils {
     public static INDArray randomOneHot(long examples, long nOut, Random rng){
         INDArray arr = Nd4j.create(examples, nOut);
         for( int i=0; i<examples; i++ ){
-            // FIXME: int cast
             arr.putScalar(i, rng.nextInt((int) nOut), 1.0);
         }
         return arr;