[rust] Provides better error message for unsupported ops (#3424)

deepjavalibrary · Aug 18, 2024 · 7932bdb · 7932bdb
1 parent cb898a6
commit 7932bdb
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 5 deletions.
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
@@ -171,6 +171,12 @@ public RsNDArray toType(DataType dataType, boolean copy) {
             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);

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
@@ -20,7 +20,6 @@
 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 void testToDataType() {
 
     @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
+            }
         }
     }