Improve Block usability

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} */

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} */

api/src/main/java/ai/djl/modality/nlp/EncoderDecoder.java

@@ -97,19 +97,18 @@ public NDList forward(
      * @param manager the NDManager to initialize the parameters
      * @param dataType the datatype of the parameters
      * @param inputShapes the shapes of the inputs to the block
-     * @return the shapes of the outputs of the block
      */
     @Override
-    public Shape[] initialize(NDManager manager, DataType dataType, Shape... inputShapes) {
+    public void initialize(NDManager manager, DataType dataType, Shape... inputShapes) {
         beforeInitialize(inputShapes);
         encoder.initialize(manager, dataType, inputShapes[0]);
-        return decoder.initialize(manager, dataType, inputShapes[1]);
+        decoder.initialize(manager, dataType, inputShapes[1]);
     }
 
     /** {@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} */

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);
     }
 }

api/src/main/java/ai/djl/nn/AbstractBlock.java

@@ -30,7 +30,7 @@
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Locale;
-import java.util.function.Function;
+import java.util.function.Predicate;
 
 /**
  * {@code AbstractBlock} is an abstract implementation of {@link Block}.
@@ -43,12 +43,11 @@
  * <ul>
  *   <li>Define a version for serializing parameter and metadata and pass it to the parent
  *       constructor
- *   <li>Use {@link AbstractBlock#addParameter(Parameter, Shape)} or {@link
- *       AbstractBlock#addParameter(Parameter, Function)} to add parameters to your block in the
+ *   <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
@@ -61,9 +60,9 @@
  * </ul>
  *
  * <p>If you use {@link AbstractBlock#addParameter(Parameter)} to add parameters, you have to take
- * care of parameter initialization yourself. In this case, you need to override {@link
- * AbstractBlock#getParameterShape(String, Shape[])} to determine the shape of your parameters. If
- * you use the other variants of {@code addParameter} this is done for you.
+ * care of parameter initialization yourself. In this case, you need to setShape to your parameters
+ * if you know the shape of Parameter or you can implement prepare to setShape when you see the
+ * input shape.
  */
 // Using LinkedHashMap instead of Map is intentional: we want to make sure that consumers
 // of this API know the children and parameters are always iterated over in insertion order.
@@ -99,14 +98,6 @@ public abstract class AbstractBlock implements Block {
      */
     protected LinkedHashMap<String, Parameter> parameters = new LinkedHashMap<>();
 
-    /**
-     * Callbacks to determine the shape of a parameter. Values may be null in which case extending
-     * classes need to override {@link Block#getParameterShape(String, Shape[])} and implement
-     * parameter shape resolution manually.
-     */
-    protected LinkedHashMap<String, Function<Shape[], Shape>> parameterShapeCallbacks =
-            new LinkedHashMap<>();
-
     /**
      * Builds an empty block with the given version for parameter serialization.
      *
@@ -195,73 +186,20 @@ protected final <B extends Block> B addChildBlock(String name, B block) {
         return block;
     }
 
-    /**
-     * Adds a parameter to this block. If parameters are added with this method, subclasses need to
-     * override {@link Block#getParameterShape(String, Shape[])} and return the shapes of parameters
-     * themselves.
-     *
-     * @param parameter the parameter to add, not null
-     * @param <P> the specific parameter subclass
-     * @return the parameter passed as arguments to make it easier to create and assign paramters in
-     *     one line
-     */
-    protected final <P extends Parameter> P addParameter(P parameter) {
-        return addParameter(parameter, (Function<Shape[], Shape>) null);
-    }
-
     /**
      * Adds a parameter to this block. If parameters are added with this method, intialization of
      * the parameter works out of the box
      *
-     * @param parameter the parameter to add, not null
-     * @param shape the shape of the parameter
      * @param <P> the specific parameter subclass
-     * @return the parameter passed as arguments to make it easier to create and assign paramters in
-     *     one line
-     */
-    protected final <P extends Parameter> P addParameter(P parameter, Shape shape) {
-        return addParameter(parameter, (inputShapes) -> shape);
-    }
-
-    /**
-     * Adds a parameter to this block. If parameters are added with this method, intialization of
-     * the parameter works out of the box
-     *
      * @param parameter the parameter to add, not null
-     * @param shapeCallback the method to call once the input shape of this block is known to
-     *     determine the shape of the given parameter
-     * @param <P> the specific parameter subclass
      * @return the parameter passed as arguments to make it easier to create and assign parameters
      *     in one line
      */
-    protected final <P extends Parameter> P addParameter(
-            P parameter, Function<Shape[], Shape> shapeCallback) {
+    protected final <P extends Parameter> P addParameter(P parameter) {
         parameters.put(parameter.getName(), parameter);
-        parameterShapeCallbacks.put(parameter.getName(), shapeCallback);
         return parameter;
     }
 
-    /** {@inheritDoc} */
-    @Override
-    public Shape getParameterShape(String name, Shape[] inputShapes) {
-        Function<Shape[], Shape> callback = parameterShapeCallbacks.get(name);
-        if (callback == null) {
-            Parameter parameter = parameters.get(name);
-            if (parameter == null) {
-                throw new IllegalArgumentException(
-                        "No parameter named " + name + " found in this block.");
-            } else {
-                throw new IllegalStateException(
-                        "No shape initializer for parameter "
-                                + name
-                                + "found. "
-                                + "Either pass an initializer for the shape when adding the "
-                                + "parameter or override getParameterShape in the subclass.");
-            }
-        }
-        return callback.apply(inputShapes);
-    }
-
     /** {@inheritDoc} */
     @Override
     public BlockList getChildren() {
@@ -285,13 +223,9 @@ public PairList<String, Shape> describeInput() {
 
     /** {@inheritDoc} */
     @Override
-    public void setInitializer(Initializer initializer) {
-        for (Parameter parameter : parameters.values()) {
-            parameter.setInitializer(initializer, false);
-        }
-        for (Block child : children.values()) {
-            child.setInitializer(initializer);
-        }
+    public void setInitializer(Initializer initializer, Parameter.Type params) {
+        Predicate<Parameter> predicate = parameter -> parameter.getType().equals(params);
+        setInitializer(initializer, predicate);
     }
 
     /** {@inheritDoc} */
@@ -301,18 +235,50 @@ public void setInitializer(Initializer initializer, String paramName) {
         if (parameter == null) {
             throw new IllegalArgumentException("Could not find parameter " + paramName);
         }
-        parameter.setInitializer(initializer, true);
+        parameter.setInitializer(initializer);
     }
 
     /** {@inheritDoc} */
     @Override
-    public Shape[] initialize(NDManager manager, DataType dataType, Shape... inputShapes) {
+    public void setInitializer(Initializer initializer, Predicate<Parameter> predicate) {
+        List<Parameter> params = getParameters().values();
+        for (Parameter param : params) {
+            if (predicate.test(param)) {
+                param.setInitializer(initializer);
+            }
+        }
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public void initialize(NDManager manager, DataType dataType, Shape... inputShapes) {
         beforeInitialize(inputShapes);
+        // if parameters are initialized, skip it
+        if (!isInitialized()) {
+            // setShape for all params
+            prepare(inputShapes);
+        }
         for (Parameter parameter : parameters.values()) {
-            parameter.initialize(manager, dataType, inputShapes);
+            parameter.initialize(manager, dataType);
         }
         initializeChildBlocks(manager, dataType, inputShapes);
-        return getOutputShapes(manager, inputShapes);
+    }
+
+    /**
+     * Performs any action necessary before initialization. For example, keep the input information
+     * or verify the layout.
+     *
+     * @param inputShapes the expected shapes of the input
+     */
+    protected void beforeInitialize(Shape... inputShapes) {
+        if (inputNames.isEmpty()) {
+            // automatically assign input names
+            inputNames = new ArrayList<>();
+            for (int i = 0; i < inputShapes.length; ++i) {
+                inputNames.add("data" + i);
+            }
+        }
+        this.inputShapes = inputShapes;
     }
 
     /**
@@ -355,20 +321,11 @@ public ParameterList getDirectParameters() {
     }
 
     /**
-     * Performs any action necessary before initialization.
+     * Sets the shape of {@link Parameter}s.
      *
-     * @param inputShapes the expected shapes of the input
+     * @param inputShapes the shapes of inputs
      */
-    protected void beforeInitialize(Shape[] inputShapes) {
-        if (inputNames.isEmpty()) {
-            // automatically assign input names
-            inputNames = new ArrayList<>();
-            for (int i = 0; i < inputShapes.length; ++i) {
-                inputNames.add("data" + i);
-            }
-        }
-        this.inputShapes = inputShapes;
-    }
+    protected void prepare(Shape[] inputShapes) {}
 
     /** {@inheritDoc} */
     @Override
@@ -494,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");

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!");
     }
 }

api/src/main/java/ai/djl/nn/Block.java

@@ -25,6 +25,7 @@
 import java.io.DataInputStream;
 import java.io.DataOutputStream;
 import java.io.IOException;
+import java.util.function.Predicate;
 
 /**
  * A {@code Block} is a composable function that forms a neural network.
@@ -158,11 +159,12 @@ default NDList forward(
     }
 
     /**
-     * Sets an {@link Initializer} to the block.
+     * Sets an {@link Initializer} to all the parameters that match parameter type in the block.
      *
      * @param initializer the initializer to set
+     * @param type the Parameter Type we want to setInitializer
      */
-    void setInitializer(Initializer initializer);
+    void setInitializer(Initializer initializer, Parameter.Type type);
 
     /**
      * Sets an {@link Initializer} to the specified direct parameter of the block, overriding the
@@ -173,15 +175,22 @@ default NDList forward(
      */
     void setInitializer(Initializer initializer, String paramName);
 
+    /**
+     * Sets an {@link Initializer} to all the parameters that match Predicate in the block.
+     *
+     * @param initializer the initializer to be set
+     * @param predicate predicate function to indicate parameters you want to set
+     */
+    void setInitializer(Initializer initializer, Predicate<Parameter> predicate);
+
     /**
      * Initializes the parameters of the block. This method must be called before calling `forward`.
      *
      * @param manager the NDManager to initialize the parameters
      * @param dataType the datatype of the parameters
      * @param inputShapes the shapes of the inputs to the block
-     * @return the shapes of the outputs of the block
      */
-    Shape[] initialize(NDManager manager, DataType dataType, Shape... inputShapes);
+    void initialize(NDManager manager, DataType dataType, Shape... inputShapes);
 
     /**
      * Returns a boolean whether the block is initialized.
@@ -232,25 +241,13 @@ default NDList forward(
      */
     ParameterList getParameters();
 
-    /**
-     * Returns the shape of the specified direct parameter of this block given the shapes of the
-     * input to the block.
-     *
-     * @param name the name of the parameter
-     * @param inputShapes the shapes of the input to the block
-     * @return the shape of the parameter specified
-     * @throws IllegalArgumentException if the parameter name specified is invalid
-     */
-    Shape getParameterShape(String name, Shape[] inputShapes);
-
     /**
      * 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.

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()];