Repository: deepjavalibrary/djl
Pull request #3424 (2024-08-17 19:00:32)

[rust] Provides better error message for unsupported ops
## 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..8e6e4f3c7 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
@@ -171,6 +171,12 @@ public class RsNDArray extends NativeResource<Long> implements NDArray {
             long newHandle = RustLibrary.toBoolean(getHandle());
             return toArray(newHandle, dataType, false, true);
         }
+        if (this.dataType == DataType.INT64
+                && dataType == DataType.FLOAT16
+                && getDevice().isGpu()) {
+            // TODO:
+            throw new UnsupportedOperationException("FP16 to I64 is not supported on GPU.");
+        }
         int dType = manager.toRustDataType(dataType);
         long newHandle = RustLibrary.toDataType(getHandle(), dType);
         return toArray(newHandle, dataType, false, true);
@@ -1125,6 +1131,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/extensions/tokenizers/src/test/java/ai/djl/engine/rust/NDArrayTests.java b/extensions/tokenizers/src/test/java/ai/djl/engine/rust/NDArrayTests.java
index 5ef89ee95..2bc0f79b0 100644
--- a/extensions/tokenizers/src/test/java/ai/djl/engine/rust/NDArrayTests.java
+++ b/extensions/tokenizers/src/test/java/ai/djl/engine/rust/NDArrayTests.java
@@ -20,7 +20,6 @@ import ai.djl.ndarray.types.DataType;
 import ai.djl.ndarray.types.Shape;
 import ai.djl.nn.Activation;
 import ai.djl.testing.Assertions;
-import ai.djl.testing.TestRequirements;
 
 import org.testng.Assert;
 import org.testng.annotations.Test;
@@ -119,14 +118,17 @@ public class NDArrayTests {
 
     @Test
     public void testI64toF16() {
-        TestRequirements.notGpu("Rust");
         try (NDManager manager = NDManager.newBaseManager("Rust")) {
             NDArray array = manager.create(2);
             Assert.assertEquals(array.getDataType(), DataType.INT32);
             NDArray int64 = array.toType(DataType.INT64, false);
-            NDArray f16 = int64.toType(DataType.FLOAT16, false);
-            NDArray bool = f16.toType(DataType.BOOLEAN, false);
-            Assert.assertTrue(bool.getBoolean());
+            try {
+                NDArray f16 = int64.toType(DataType.FLOAT16, false);
+                NDArray bool = f16.toType(DataType.BOOLEAN, false);
+                Assert.assertTrue(bool.getBoolean());
+            } catch (UnsupportedOperationException e) {
+                // ignore on GPU
+            }
         }
     }
 
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())) {