Fix the usage of the repeat function for embedding

api/src/main/java/ai/djl/nn/core/ConstantEmbedding.java

@@ -54,7 +54,8 @@ protected NDList forwardInternal(
         NDArray base = manager.create(embedding.getShape());
         embedding.copyTo(base);
         Shape shape = inputs.get(0).getShape().addAll(embedding.getShape());
-        return new NDList(base.repeat(shape));
+        return new NDList(
+                base.reshape(1, embedding.size()).repeat(0, inputs.get(0).size()).reshape(shape));
     }
 
     /** {@inheritDoc} */
@@ -104,8 +105,8 @@ public long embed(Object item) {
     public NDArray embed(NDManager manager, Object[] items) {
         NDArray base = manager.create(embedding.getShape());
         embedding.copyTo(base);
-        Shape shape = new Shape(items.length).addAll(embedding.getShape());
-        return base.repeat(shape);
+        return base.repeat(0, items.length)
+                .reshape(new Shape(items.length).addAll(embedding.getShape()));
     }
 
     /** {@inheritDoc} */

integration/src/main/java/ai/djl/integration/tests/nn/BlockCoreTest.java

@@ -34,6 +34,7 @@
 import ai.djl.nn.convolutional.Conv2d;
 import ai.djl.nn.convolutional.Conv2dTranspose;
 import ai.djl.nn.convolutional.Conv3d;
+import ai.djl.nn.core.ConstantEmbedding;
 import ai.djl.nn.core.Linear;
 import ai.djl.nn.core.LinearCollection;
 import ai.djl.nn.core.Multiplication;
@@ -586,6 +587,40 @@ public void testEmbedding() throws IOException, MalformedModelException {
         }
     }
 
+    @Test
+    public void testConstantEmbedding() {
+        TrainingConfig config =
+                new DefaultTrainingConfig(Loss.l2Loss())
+                        .optInitializer(Initializer.ONES, Parameter.Type.WEIGHT);
+
+        try (Model model = Model.newInstance("model", TestUtils.getEngine())) {
+            NDArray embedding = model.getNDManager().create(new float[] {1, 2}, new Shape(2));
+            ConstantEmbedding block = new ConstantEmbedding(embedding);
+            model.setBlock(block);
+
+            try (Trainer trainer = model.newTrainer(config)) {
+                Shape inputShape = new Shape(2);
+                trainer.initialize(inputShape);
+
+                NDManager manager = trainer.getManager();
+
+                Assert.assertEquals(
+                        trainer.forward(new NDList(manager.create(0.1f))).singletonOrThrow(),
+                        manager.create(new float[] {1, 2}, new Shape(2)));
+
+                Assert.assertEquals(
+                        trainer.forward(new NDList(manager.create(new float[] {1, 1})))
+                                .singletonOrThrow(),
+                        manager.create(new float[] {1, 2, 1, 2}, new Shape(2, 2)));
+
+                Assert.assertEquals(block.embed("x"), 0);
+                Assert.assertEquals(
+                        block.embed(manager, new String[] {"x"}),
+                        manager.create(new float[] {1, 2}, new Shape(1, 2)));
+            }
+        }
+    }
+
     @Test
     public void testConv1d() throws IOException, MalformedModelException {
         TrainingConfig config =