deepjavalibrary · stu1130 · Mar 1, 2021 · Mar 1, 2021
diff --git a/api/src/main/java/ai/djl/modality/nlp/Decoder.java b/api/src/main/java/ai/djl/modality/nlp/Decoder.java
@@ -64,8 +64,8 @@ public void initializeChildBlocks(NDManager manager, DataType dataType, Shape...
 
     /** {@inheritDoc} */
     @Override
-    public Shape[] getOutputShapes(NDManager manager, Shape[] inputShapes) {
-        return block.getOutputShapes(manager, inputShapes);
+    public Shape[] getOutputShapes(Shape[] inputShapes) {
+        return block.getOutputShapes(inputShapes);
     }
 
     /** {@inheritDoc} */

diff --git a/api/src/main/java/ai/djl/modality/nlp/Encoder.java b/api/src/main/java/ai/djl/modality/nlp/Encoder.java
@@ -79,8 +79,8 @@ public void initializeChildBlocks(NDManager manager, DataType dataType, Shape...
 
     /** {@inheritDoc} */
     @Override
-    public Shape[] getOutputShapes(NDManager manager, Shape[] inputShapes) {
-        return block.getOutputShapes(manager, inputShapes);
+    public Shape[] getOutputShapes(Shape[] inputShapes) {
+        return block.getOutputShapes(inputShapes);
     }
 
     /** {@inheritDoc} */

diff --git a/api/src/main/java/ai/djl/modality/nlp/EncoderDecoder.java b/api/src/main/java/ai/djl/modality/nlp/EncoderDecoder.java
@@ -107,8 +107,8 @@ public void initialize(NDManager manager, DataType dataType, Shape... inputShape
 
     /** {@inheritDoc} */
     @Override
-    public Shape[] getOutputShapes(NDManager manager, Shape[] inputShapes) {
-        return decoder.getOutputShapes(manager, new Shape[] {inputShapes[1]});
+    public Shape[] getOutputShapes(Shape[] inputShapes) {
+        return decoder.getOutputShapes(new Shape[] {inputShapes[1]});
     }
 
     /** {@inheritDoc} */

diff --git a/api/src/main/java/ai/djl/modality/nlp/embedding/TrainableTextEmbedding.java b/api/src/main/java/ai/djl/modality/nlp/embedding/TrainableTextEmbedding.java
@@ -88,7 +88,7 @@ public void initializeChildBlocks(NDManager manager, DataType dataType, Shape...
 
     /** {@inheritDoc} */
     @Override
-    public Shape[] getOutputShapes(NDManager manager, Shape[] inputShapes) {
-        return trainableWordEmbedding.getOutputShapes(manager, inputShapes);
+    public Shape[] getOutputShapes(Shape[] inputShapes) {
+        return trainableWordEmbedding.getOutputShapes(inputShapes);
     }
 }
diff --git a/api/src/main/java/ai/djl/nn/AbstractBlock.java b/api/src/main/java/ai/djl/nn/AbstractBlock.java
@@ -46,8 +46,8 @@
  *   <li>Use {@link AbstractBlock#addParameter(Parameter)} to add parameters to your block in the
  *       constructor if necessary.
  *   <li>Use {@link AbstractBlock#addChildBlock(String, Block)} to add child blocks if necessary.
- *   <li>Override {@link AbstractBlock#getOutputShapes(NDManager, Shape[])} to determine the shape
- *       of your custom block's output based on the input it will receive.
+ *   <li>Override {@link Block#getOutputShapes(Shape[])} to determine the shape of your custom
+ *       block's output based on the input it will receive.
  *   <li>Override {@link AbstractBlock#initializeChildBlocks(NDManager, DataType, Shape...)} if you
  *       added child blocks to initialize them based on the input shape your block will receive. You
  *       can skip this if your block does not contain child blocks
@@ -451,7 +451,7 @@ public String toString() {
             appendShape(sb, inputShapeDescription.values().toArray(new Shape[0]));
             sb.append(" -> ");
             Shape[] outputShapes =
-                    getOutputShapes(null, inputShapeDescription.values().toArray(new Shape[0]));
+                    getOutputShapes(inputShapeDescription.values().toArray(new Shape[0]));
             appendShape(sb, outputShapes);
         } else {
             sb.append("Uninitialized");

diff --git a/api/src/main/java/ai/djl/nn/AbstractSymbolBlock.java b/api/src/main/java/ai/djl/nn/AbstractSymbolBlock.java
@@ -12,7 +12,6 @@
  */
 package ai.djl.nn;
 
-import ai.djl.ndarray.NDManager;
 import ai.djl.ndarray.types.Shape;
 
 /** {@code AbstractSymbolBlock} is an abstract implementation of {@link SymbolBlock}. */
@@ -29,7 +28,7 @@ public AbstractSymbolBlock(byte version) {
 
     /** {@inheritDoc} */
     @Override
-    public Shape[] getOutputShapes(NDManager manager, Shape[] inputShapes) {
+    public Shape[] getOutputShapes(Shape[] inputShapes) {
         throw new UnsupportedOperationException("not implement!");
     }
 }
diff --git a/api/src/main/java/ai/djl/nn/Block.java b/api/src/main/java/ai/djl/nn/Block.java
@@ -244,11 +244,10 @@ default NDList forward(
     /**
      * Returns the expected output shapes of the block for the specified input shapes.
      *
-     * @param manager an NDManager
      * @param inputShapes the shapes of the inputs
      * @return the expected output shapes of the block
      */
-    Shape[] getOutputShapes(NDManager manager, Shape[] inputShapes);
+    Shape[] getOutputShapes(Shape[] inputShapes);
 
     /**
      * Writes the parameters of the block to the given outputStream.

diff --git a/api/src/main/java/ai/djl/nn/LambdaBlock.java b/api/src/main/java/ai/djl/nn/LambdaBlock.java
@@ -70,11 +70,11 @@ protected NDList forwardInternal(
 
     /** {@inheritDoc} */
     @Override
-    public Shape[] getOutputShapes(NDManager manager, Shape[] inputShapes) {
-        try (NDManager subManager = manager.newSubManager()) {
+    public Shape[] getOutputShapes(Shape[] inputShapes) {
+        try (NDManager manager = NDManager.newBaseManager()) {
             NDList input = new NDList(inputShapes.length);
             for (Shape shape : inputShapes) {
-                input.add(subManager.zeros(shape));
+                input.add(manager.zeros(shape));
             }
             NDList output = lambda.apply(input);
             Shape[] outputShapes = new Shape[output.size()];

diff --git a/api/src/main/java/ai/djl/nn/ParallelBlock.java b/api/src/main/java/ai/djl/nn/ParallelBlock.java
@@ -149,16 +149,16 @@ public void initializeChildBlocks(NDManager manager, DataType dataType, Shape...
 
     /** {@inheritDoc} */
     @Override
-    public Shape[] getOutputShapes(NDManager manager, Shape[] inputShapes) {
+    public Shape[] getOutputShapes(Shape[] inputShapes) {
         Preconditions.checkArgument(!children.isEmpty(), "The parallel block is empty");
 
-        try (NDManager subManager = manager.newSubManager()) {
+        try (NDManager manager = NDManager.newBaseManager()) {
             List<NDList> inputs = new ArrayList<>();
             for (Block block : children.values()) {
-                Shape[] shapes = block.getOutputShapes(manager, inputShapes);
+                Shape[] shapes = block.getOutputShapes(inputShapes);
                 NDList output = new NDList(shapes.length);
                 for (Shape shape : shapes) {
-                    output.add(subManager.create(shape));
+                    output.add(manager.create(shape));
                 }
                 inputs.add(output);
             }

diff --git a/api/src/main/java/ai/djl/nn/SequentialBlock.java b/api/src/main/java/ai/djl/nn/SequentialBlock.java
@@ -155,19 +155,19 @@ public void initializeChildBlocks(NDManager manager, DataType dataType, Shape...
         Shape[] shapes = inputShapes;
         for (Block child : getChildren().values()) {
             child.initialize(manager, dataType, shapes);
-            shapes = child.getOutputShapes(manager, shapes);
+            shapes = child.getOutputShapes(shapes);
         }
     }
 
     /** {@inheritDoc} */
     @Override
-    public Shape[] getOutputShapes(NDManager manager, Shape[] inputs) {
+    public Shape[] getOutputShapes(Shape[] inputs) {
         if (children.isEmpty()) {
             throw new IllegalArgumentException("The sequential block is empty");
         }
         Shape[] current = inputs;
         for (Block block : children.values()) {
-            current = block.getOutputShapes(manager, current);
+            current = block.getOutputShapes(current);
         }
         return current;
     }

diff --git a/api/src/main/java/ai/djl/nn/convolutional/Convolution.java b/api/src/main/java/ai/djl/nn/convolutional/Convolution.java
@@ -16,7 +16,6 @@
 import ai.djl.MalformedModelException;
 import ai.djl.ndarray.NDArray;
 import ai.djl.ndarray.NDList;
-import ai.djl.ndarray.NDManager;
 import ai.djl.ndarray.types.LayoutType;
 import ai.djl.ndarray.types.Shape;
 import ai.djl.nn.AbstractBlock;
@@ -169,7 +168,7 @@ protected void prepare(Shape[] inputs) {
 
     /** {@inheritDoc} */
     @Override
-    public Shape[] getOutputShapes(NDManager manager, Shape[] inputs) {
+    public Shape[] getOutputShapes(Shape[] inputs) {
         long[] shape = new long[numDimensions()];
         shape[0] = inputs[0].get(0);
         shape[1] = filters;

diff --git a/api/src/main/java/ai/djl/nn/convolutional/Deconvolution.java b/api/src/main/java/ai/djl/nn/convolutional/Deconvolution.java
@@ -16,7 +16,6 @@
 import ai.djl.MalformedModelException;
 import ai.djl.ndarray.NDArray;
 import ai.djl.ndarray.NDList;
-import ai.djl.ndarray.NDManager;
 import ai.djl.ndarray.types.LayoutType;
 import ai.djl.ndarray.types.Shape;
 import ai.djl.nn.AbstractBlock;
@@ -146,7 +145,7 @@ protected void prepare(Shape[] inputs) {
 
     /** {@inheritDoc} */
     @Override
-    public Shape[] getOutputShapes(NDManager manager, Shape[] inputs) {
+    public Shape[] getOutputShapes(Shape[] inputs) {
         long[] shape = new long[numDimensions()];
         shape[0] = inputs[0].get(0);
         shape[1] = filters;

diff --git a/api/src/main/java/ai/djl/nn/core/ConstantEmbedding.java b/api/src/main/java/ai/djl/nn/core/ConstantEmbedding.java
@@ -57,7 +57,7 @@ protected NDList forwardInternal(
 
     /** {@inheritDoc} */
     @Override
-    public Shape[] getOutputShapes(NDManager manager, Shape[] inputShapes) {
+    public Shape[] getOutputShapes(Shape[] inputShapes) {
         return new Shape[] {inputShapes[0].addAll(embedding.getShape())};
     }
 

diff --git a/api/src/main/java/ai/djl/nn/core/Embedding.java b/api/src/main/java/ai/djl/nn/core/Embedding.java
@@ -113,7 +113,7 @@ public void prepare(Shape[] inputShapes) {
 
     /** {@inheritDoc} */
     @Override
-    public Shape[] getOutputShapes(NDManager manager, Shape[] inputShapes) {
+    public Shape[] getOutputShapes(Shape[] inputShapes) {
         return new Shape[] {inputShapes[0].addAll(new Shape(embeddingSize))};
     }
 

diff --git a/api/src/main/java/ai/djl/nn/core/Linear.java b/api/src/main/java/ai/djl/nn/core/Linear.java
@@ -16,7 +16,6 @@
 import ai.djl.MalformedModelException;
 import ai.djl.ndarray.NDArray;
 import ai.djl.ndarray.NDList;
-import ai.djl.ndarray.NDManager;
 import ai.djl.ndarray.types.Shape;
 import ai.djl.nn.AbstractBlock;
 import ai.djl.nn.Block;
@@ -90,7 +89,7 @@ protected NDList forwardInternal(
 
     /** {@inheritDoc} */
     @Override
-    public Shape[] getOutputShapes(NDManager manager, Shape[] inputs) {
+    public Shape[] getOutputShapes(Shape[] inputs) {
         return new Shape[] {inputs[0].slice(0, inputs[0].dimension() - 1).add(units)};
     }
 

diff --git a/api/src/main/java/ai/djl/nn/core/Prelu.java b/api/src/main/java/ai/djl/nn/core/Prelu.java
@@ -15,7 +15,6 @@
 import ai.djl.MalformedModelException;
 import ai.djl.ndarray.NDArray;
 import ai.djl.ndarray.NDList;
-import ai.djl.ndarray.NDManager;
 import ai.djl.ndarray.internal.NDArrayEx;
 import ai.djl.ndarray.types.Shape;
 import ai.djl.nn.AbstractBlock;
@@ -66,7 +65,7 @@ protected NDList forwardInternal(
 
     /** {@inheritDoc} */
     @Override
-    public Shape[] getOutputShapes(NDManager manager, Shape[] inputs) {
+    public Shape[] getOutputShapes(Shape[] inputs) {
         return new Shape[] {inputs[0]};
     }
 

diff --git a/api/src/main/java/ai/djl/nn/norm/BatchNorm.java b/api/src/main/java/ai/djl/nn/norm/BatchNorm.java
@@ -16,7 +16,6 @@
 import ai.djl.MalformedModelException;
 import ai.djl.ndarray.NDArray;
 import ai.djl.ndarray.NDList;
-import ai.djl.ndarray.NDManager;
 import ai.djl.ndarray.internal.NDArrayEx;
 import ai.djl.ndarray.types.Shape;
 import ai.djl.nn.AbstractBlock;
@@ -145,7 +144,7 @@ protected NDList forwardInternal(
 
     /** {@inheritDoc} */
     @Override
-    public Shape[] getOutputShapes(NDManager manager, Shape[] inputShapes) {
+    public Shape[] getOutputShapes(Shape[] inputShapes) {
         return new Shape[] {inputShapes[0]};
     }
 

diff --git a/api/src/main/java/ai/djl/nn/norm/Dropout.java b/api/src/main/java/ai/djl/nn/norm/Dropout.java
@@ -15,7 +15,6 @@
 import ai.djl.MalformedModelException;
 import ai.djl.ndarray.NDArray;
 import ai.djl.ndarray.NDList;
-import ai.djl.ndarray.NDManager;
 import ai.djl.ndarray.internal.NDArrayEx;
 import ai.djl.ndarray.types.Shape;
 import ai.djl.nn.AbstractBlock;
@@ -76,7 +75,7 @@ protected NDList forwardInternal(
 
     /** {@inheritDoc} */
     @Override
-    public Shape[] getOutputShapes(NDManager manager, Shape[] inputShapes) {
+    public Shape[] getOutputShapes(Shape[] inputShapes) {
         return new Shape[] {inputShapes[0]};
     }
 

diff --git a/api/src/main/java/ai/djl/nn/recurrent/RecurrentBlock.java b/api/src/main/java/ai/djl/nn/recurrent/RecurrentBlock.java
@@ -13,7 +13,6 @@
 package ai.djl.nn.recurrent;
 
 import ai.djl.MalformedModelException;
-import ai.djl.ndarray.NDManager;
 import ai.djl.ndarray.types.LayoutType;
 import ai.djl.ndarray.types.Shape;
 import ai.djl.nn.AbstractBlock;
@@ -91,7 +90,7 @@ public RecurrentBlock(BaseBuilder<?> builder) {
 
     /** {@inheritDoc} */
     @Override
-    public Shape[] getOutputShapes(NDManager manager, Shape[] inputs) {
+    public Shape[] getOutputShapes(Shape[] inputs) {
         Shape inputShape = inputs[0];
         Shape outputShape =
                 new Shape(inputShape.get(0), inputShape.get(1), stateSize * getNumDirections());

diff --git a/api/src/main/java/ai/djl/nn/transformer/BertBlock.java b/api/src/main/java/ai/djl/nn/transformer/BertBlock.java
@@ -156,7 +156,7 @@ public int getTypeDictionarySize() {
 
     /** {@inheritDoc} */
     @Override
-    public Shape[] getOutputShapes(NDManager manager, Shape[] inputShapes) {
+    public Shape[] getOutputShapes(Shape[] inputShapes) {
         long batch = inputShapes[0].get(0);
         long seqLength = inputShapes[0].get(1);
         return new Shape[] {
@@ -172,7 +172,7 @@ public void initializeChildBlocks(NDManager manager, DataType dataType, Shape...
         Shape[] tokenShape = {inputShapes[0]};
         Shape[] typeShape = {inputShapes[1]};
         this.tokenEmbedding.initialize(manager, dataType, tokenShape);
-        Shape[] embeddingOutput = this.tokenEmbedding.getOutputShapes(manager, tokenShape);
+        Shape[] embeddingOutput = this.tokenEmbedding.getOutputShapes(tokenShape);
         this.typeEmbedding.initialize(manager, dataType, typeShape);
         this.embeddingNorm.initialize(manager, dataType, embeddingOutput);
         this.embeddingDropout.initialize(manager, dataType, embeddingOutput);

diff --git a/api/src/main/java/ai/djl/nn/transformer/BertMaskedLanguageModelBlock.java b/api/src/main/java/ai/djl/nn/transformer/BertMaskedLanguageModelBlock.java
@@ -142,7 +142,7 @@ protected NDList forwardInternal(
 
     /** {@inheritDoc} */
     @Override
-    public Shape[] getOutputShapes(final NDManager manager, final Shape[] inputShapes) {
+    public Shape[] getOutputShapes(final Shape[] inputShapes) {
         int batchSize = (int) inputShapes[0].get(0);
         int indexCount = (int) inputShapes[1].get(1);
         int dictionarySize = (int) inputShapes[2].get(0);

diff --git a/api/src/main/java/ai/djl/nn/transformer/BertNextSentenceBlock.java b/api/src/main/java/ai/djl/nn/transformer/BertNextSentenceBlock.java
@@ -53,7 +53,7 @@ protected NDList forwardInternal(
 
     /** {@inheritDoc} */
     @Override
-    public Shape[] getOutputShapes(NDManager manager, Shape[] inputShapes) {
+    public Shape[] getOutputShapes(Shape[] inputShapes) {
         return new Shape[] {new Shape(inputShapes[0].get(0), 2)};
     }
 }
diff --git a/api/src/main/java/ai/djl/nn/transformer/BertPretrainingBlock.java b/api/src/main/java/ai/djl/nn/transformer/BertPretrainingBlock.java
@@ -52,7 +52,7 @@ public BertPretrainingBlock(final BertBlock.Builder builder) {
     public void initializeChildBlocks(NDManager manager, DataType dataType, Shape... inputShapes) {
         inputNames = Arrays.asList("tokenIds", "typeIds", "sequenceMasks", "maskedIndices");
         bertBlock.initialize(manager, dataType, inputShapes);
-        Shape[] bertOutputShapes = bertBlock.getOutputShapes(manager, inputShapes);
+        Shape[] bertOutputShapes = bertBlock.getOutputShapes(inputShapes);
         Shape embeddedSequence = bertOutputShapes[0];
         Shape pooledOutput = bertOutputShapes[1];
         Shape maskedIndices = inputShapes[2];
@@ -98,7 +98,7 @@ protected NDList forwardInternal(
 
     /** {@inheritDoc} */
     @Override
-    public Shape[] getOutputShapes(NDManager manager, Shape[] inputShapes) {
+    public Shape[] getOutputShapes(Shape[] inputShapes) {
         long batchSize = inputShapes[0].get(0);
         long maskedIndexCount = inputShapes[3].get(1);
         return new Shape[] {

diff --git a/api/src/main/java/ai/djl/nn/transformer/IdEmbedding.java b/api/src/main/java/ai/djl/nn/transformer/IdEmbedding.java
@@ -55,7 +55,7 @@ private IdEmbedding(Builder builder) {
 
     /** {@inheritDoc} */
     @Override
-    public Shape[] getOutputShapes(NDManager manager, Shape[] inputShapes) {
+    public Shape[] getOutputShapes(Shape[] inputShapes) {
         return new Shape[] {inputShapes[0].addAll(new Shape(embeddingSize))};
     }
 

diff --git a/api/src/main/java/ai/djl/nn/transformer/PointwiseFeedForwardBlock.java b/api/src/main/java/ai/djl/nn/transformer/PointwiseFeedForwardBlock.java
@@ -59,9 +59,9 @@ public PointwiseFeedForwardBlock(
 
     /** {@inheritDoc} */
     @Override
-    public Shape[] getOutputShapes(NDManager manager, Shape[] inputShapes) {
+    public Shape[] getOutputShapes(Shape[] inputShapes) {
         for (Block child : children.values()) {
-            inputShapes = child.getOutputShapes(manager, inputShapes);
+            inputShapes = child.getOutputShapes(inputShapes);
         }
         return inputShapes;
     }
@@ -84,7 +84,7 @@ public void initializeChildBlocks(NDManager manager, DataType dataType, Shape...
         Shape lastShape = inputShape;
         for (Block child : children.values()) {
             child.initialize(manager, dataType, lastShape);
-            lastShape = getOutputShapes(manager, new Shape[] {lastShape})[0];
+            lastShape = getOutputShapes(new Shape[] {lastShape})[0];
         }
     }
 

diff --git a/api/src/main/java/ai/djl/nn/transformer/ScaledDotProductAttentionBlock.java b/api/src/main/java/ai/djl/nn/transformer/ScaledDotProductAttentionBlock.java
@@ -146,7 +146,7 @@ public Linear getResultProjection() {
 
     /** {@inheritDoc} */
     @Override
-    public Shape[] getOutputShapes(NDManager manager, Shape[] inputShapes) {
+    public Shape[] getOutputShapes(Shape[] inputShapes) {
         // Return shape is the shape of the query. For 2 or less inputs we have self-attention, i.e.
         // the shape of the output is the shape of the input
         if (inputShapes.length == 1 || inputShapes.length == 2) {

diff --git a/api/src/main/java/ai/djl/nn/transformer/TransformerEncoderBlock.java b/api/src/main/java/ai/djl/nn/transformer/TransformerEncoderBlock.java
@@ -82,7 +82,7 @@ public TransformerEncoderBlock(
 
     /** {@inheritDoc} */
     @Override
-    public Shape[] getOutputShapes(NDManager manager, Shape[] inputShapes) {
+    public Shape[] getOutputShapes(Shape[] inputShapes) {
         return inputShapes;
     }
 

diff --git a/api/src/test/java/ai/djl/nn/convolutional/ShapeUtils.java b/api/src/test/java/ai/djl/nn/convolutional/ShapeUtils.java
@@ -31,7 +31,7 @@ private ShapeUtils() {}
      * @return the corresponding output shape for the provided input
      */
     public static Shape outputShapeForBlock(NDManager manager, Block block, Shape inputShape) {
-        Shape[] outputs = block.getOutputShapes(manager, new Shape[] {inputShape});
+        Shape[] outputs = block.getOutputShapes(new Shape[] {inputShape});
         return outputs[0];
     }