Repository: deepjavalibrary/djl
Pull request #3423 (2024-08-17 18:37:17)

[pytorch] Adds pad operator
## Description ##

Brief description of what this PR is about

- If this change is a backward incompatible change, why must this change be made?
- Interesting edge cases to note here



Patch:
diff --git a/api/src/main/java/ai/djl/ndarray/NDArray.java b/api/src/main/java/ai/djl/ndarray/NDArray.java
index d9227a3c3..3d93c9878 100644
--- a/api/src/main/java/ai/djl/ndarray/NDArray.java
+++ b/api/src/main/java/ai/djl/ndarray/NDArray.java
@@ -3518,6 +3518,24 @@ public interface NDArray extends NDResource, BytesSupplier {
         return ifft2(sizes, new long[] {-2, -1});
     }
 
+    /**
+     * Pads this {@code NDArray} with the given {@link Shape}.
+     *
+     * <p>Examples
+     *
+     * <pre>
+     * NDArray array = manager.zeros(3, 3, 4, 2);
+     * array.pad(new Shape(1, 1), 0); # pad last dim by 1 on each side
+     * array.getShape() => (3, 3, 4, 4)
+     * </pre>
+     *
+     * @param padding the padding shape, must be even
+     * @return a padded {@code NDArray}
+     * @throws IllegalArgumentException thrown if the given {@code padding} does not match the size
+     *     of the current shape
+     */
+    NDArray pad(Shape padding, double value);
+
     /**
      * Reshapes this {@code NDArray} to the given {@link Shape}.
      *
diff --git a/api/src/main/java/ai/djl/ndarray/NDArrayAdapter.java b/api/src/main/java/ai/djl/ndarray/NDArrayAdapter.java
index 27c362c88..04ea056b4 100644
--- a/api/src/main/java/ai/djl/ndarray/NDArrayAdapter.java
+++ b/api/src/main/java/ai/djl/ndarray/NDArrayAdapter.java
@@ -936,6 +936,12 @@ public abstract class NDArrayAdapter implements NDArray {
         throw new UnsupportedOperationException("Not implemented yet.");
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public NDArray pad(Shape padding, double value) {
+        throw new UnsupportedOperationException("Not implemented yet.");
+    }
+
     /** {@inheritDoc} */
     @Override
     public NDArray ifft2(long[] sizes, long[] axes) {
diff --git a/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxNDArray.java b/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxNDArray.java
index f6d2c0aa0..85f940d54 100644
--- a/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxNDArray.java
+++ b/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxNDArray.java
@@ -1194,6 +1194,12 @@ public class MxNDArray extends NativeResource<Pointer> implements LazyNDArray {
         throw new UnsupportedOperationException("Not implemented yet.");
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public NDArray pad(Shape padding, double value) {
+        throw new UnsupportedOperationException("Not implemented yet.");
+    }
+
     /** {@inheritDoc} */
     @Override
     public NDArray reshape(Shape shape) {
diff --git a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtNDArray.java b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtNDArray.java
index 1edd61226..04c15f79b 100644
--- a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtNDArray.java
+++ b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtNDArray.java
@@ -1144,6 +1144,12 @@ public class PtNDArray extends NativeResource<Long> implements NDArray {
         return JniUtils.ifft2(this, sizes, axes);
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public NDArray pad(Shape padding, double value) {
+        return JniUtils.pad(this, padding.getShape(), value);
+    }
+
     /** {@inheritDoc} */
     @Override
     public PtNDArray reshape(Shape shape) {
diff --git a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/JniUtils.java b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/JniUtils.java
index b80c096cb..8f1ac296c 100644
--- a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/JniUtils.java
+++ b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/JniUtils.java
@@ -616,6 +616,12 @@ public final class JniUtils {
                 ndArray.getManager(), PyTorchLibrary.LIB.tensorClone(ndArray.getHandle()));
     }
 
+    public static PtNDArray pad(PtNDArray ndArray, long[] shape, double value) {
+        return new PtNDArray(
+                ndArray.getManager(),
+                PyTorchLibrary.LIB.torchPad(ndArray.getHandle(), shape, value));
+    }
+
     public static PtNDArray reshape(PtNDArray ndArray, long[] shape) {
         return new PtNDArray(
                 ndArray.getManager(), PyTorchLibrary.LIB.torchReshape(ndArray.getHandle(), shape));
diff --git a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/PyTorchLibrary.java b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/PyTorchLibrary.java
index 5c1fd290f..dfe7e0d88 100644
--- a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/PyTorchLibrary.java
+++ b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/PyTorchLibrary.java
@@ -165,6 +165,8 @@ final class PyTorchLibrary {
 
     native long torchLogicalNot(long handle);
 
+    native long torchPad(long handle, long[] shape, double value);
+
     native long torchReshape(long handle, long[] shape);
 
     native long torchSoftmax(long handle, long dim, int dType);
diff --git a/engines/pytorch/pytorch-native/src/main/native/ai_djl_pytorch_jni_PyTorchLibrary_nn_functional.cc b/engines/pytorch/pytorch-native/src/main/native/ai_djl_pytorch_jni_PyTorchLibrary_nn_functional.cc
index 3f1bcfd58..033cc5576 100644
--- a/engines/pytorch/pytorch-native/src/main/native/ai_djl_pytorch_jni_PyTorchLibrary_nn_functional.cc
+++ b/engines/pytorch/pytorch-native/src/main/native/ai_djl_pytorch_jni_PyTorchLibrary_nn_functional.cc
@@ -19,6 +19,16 @@
 
 // The file is the implementation for PyTorch neural network functional ops
 
+JNIEXPORT jlong JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_torchPad(
+    JNIEnv* env, jobject jthis, jlong jhandle, jlongArray jshape, jdouble jvalue) {
+  API_BEGIN()
+  const auto* tensor_ptr = reinterpret_cast<torch::Tensor*>(jhandle);
+  const auto shape_vec = djl::utils::jni::GetVecFromJLongArray(env, jshape);
+  const auto* result_ptr = new torch::Tensor(pad(*tensor_ptr, c10::ArrayRef(shape_vec), "constant", jvalue));
+  return reinterpret_cast<uintptr_t>(result_ptr);
+  API_END_RETURN()
+}
+
 JNIEXPORT jlong JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_torchSoftmax(
     JNIEnv* env, jobject jthis, jlong jhandle, jlong jdim, jint jdtype) {
   API_BEGIN()
diff --git a/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/TfNDArray.java b/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/TfNDArray.java
index 88df43498..4845ffd09 100644
--- a/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/TfNDArray.java
+++ b/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/TfNDArray.java
@@ -1215,6 +1215,12 @@ public class TfNDArray extends NativeResource<TFE_TensorHandle> implements NDArr
         throw new UnsupportedOperationException("Not implemented yet.");
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public NDArray pad(Shape padding, double value) {
+        throw new UnsupportedOperationException("Not implemented yet.");
+    }
+
     /** {@inheritDoc} */
     @Override
     public NDArray ifft2(long[] sizes, long[] axes) {
diff --git a/extensions/tokenizers/src/main/java/ai/djl/engine/rust/RsNDArray.java b/extensions/tokenizers/src/main/java/ai/djl/engine/rust/RsNDArray.java
index 97e1148bf..9495d5b54 100644
--- a/extensions/tokenizers/src/main/java/ai/djl/engine/rust/RsNDArray.java
+++ b/extensions/tokenizers/src/main/java/ai/djl/engine/rust/RsNDArray.java
@@ -1125,6 +1125,12 @@ public class RsNDArray extends NativeResource<Long> implements NDArray {
         throw new UnsupportedOperationException("Not implemented");
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public NDArray pad(Shape padding, double value) {
+        throw new UnsupportedOperationException("Not implemented");
+    }
+
     /** {@inheritDoc} */
     @Override
     public NDArray ifft2(long[] sizes, long[] axes) {


Test patch:
diff --git a/integration/src/main/java/ai/djl/integration/tests/ndarray/NDArrayShapesManipulationOpTest.java b/integration/src/main/java/ai/djl/integration/tests/ndarray/NDArrayShapesManipulationOpTest.java
index 87cfd4435..44858406a 100644
--- a/integration/src/main/java/ai/djl/integration/tests/ndarray/NDArrayShapesManipulationOpTest.java
+++ b/integration/src/main/java/ai/djl/integration/tests/ndarray/NDArrayShapesManipulationOpTest.java
@@ -86,6 +86,22 @@ public class NDArrayShapesManipulationOpTest {
         }
     }
 
+    @Test
+    public void testPad() {
+        try (NDManager manager = NDManager.newBaseManager(TestUtils.getEngine())) {
+            NDArray array = manager.zeros(new Shape(3, 3, 4, 2));
+            NDArray result = array.pad(new Shape(1, 1), 1d);
+            Assert.assertEquals(result.getShape(), new Shape(3, 3, 4, 4));
+            Assert.assertEquals(result.getFloat(0, 0, 0, 0), 1f);
+
+            result = array.pad(new Shape(1, 1, 2, 2), 1d);
+            Assert.assertEquals(result.getShape(), new Shape(3, 3, 8, 4));
+
+            result = array.pad(new Shape(0, 1, 2, 1, 3, 3), 1d);
+            Assert.assertEquals(result.getShape(), new Shape(3, 9, 7, 3));
+        }
+    }
+
     @Test
     public void testReshape() {
         try (NDManager manager = NDManager.newBaseManager(TestUtils.getEngine())) {