[pytorch] Implements gammaln operator for PyTorch

Fixes #3259
deepjavalibrary · Jun 18, 2024 · 4b9c9d5 · 4b9c9d5
1 parent 04d6561
commit 4b9c9d5
Show file tree

Hide file tree

Showing 6 changed files with 26 additions and 6 deletions.
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
@@ -842,7 +842,7 @@ public PtNDArray exp() {
     /** {@inheritDoc} */
     @Override
     public NDArray gammaln() {
-        throw new UnsupportedOperationException("Not implemented yet.");
+        return JniUtils.gammaln(this);
     }
 
     /** {@inheritDoc} */

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
@@ -1112,6 +1112,11 @@ public static PtNDArray exp(PtNDArray ndArray) {
                 ndArray.getManager(), PyTorchLibrary.LIB.torchExp(ndArray.getHandle()));
     }
 
+    public static PtNDArray gammaln(PtNDArray ndArray) {
+        return new PtNDArray(
+                ndArray.getManager(), PyTorchLibrary.LIB.torchLgamma(ndArray.getHandle()));
+    }
+
     public static PtNDArray log(PtNDArray ndArray) {
         return new PtNDArray(
                 ndArray.getManager(), PyTorchLibrary.LIB.torchLog(ndArray.getHandle()));

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
@@ -324,6 +324,8 @@ native long[] torchUnique(
 
     native long torchExp(long handle);
 
+    native long torchLgamma(long handle);
+
     native long torchLog(long handle);
 
     native long torchLog10(long handle);

diff --git a/...torch/pytorch-native/src/main/native/ai_djl_pytorch_jni_PyTorchLibrary_torch_pointwise.cc b/...torch/pytorch-native/src/main/native/ai_djl_pytorch_jni_PyTorchLibrary_torch_pointwise.cc
@@ -283,6 +283,14 @@ JNIEXPORT jlong JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_torchExp(JNIEnv*
   API_END_RETURN()
 }
 
+JNIEXPORT jlong JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_torchLgamma(JNIEnv* env, jobject jthis, jlong jhandle) {
+  API_BEGIN()
+  const auto* tensor_ptr = reinterpret_cast<torch::Tensor*>(jhandle);
+  const auto* result_ptr = new torch::Tensor(tensor_ptr->lgamma());
+  return reinterpret_cast<uintptr_t>(result_ptr);
+  API_END_RETURN()
+}
+
 JNIEXPORT jlong JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_torchLog(JNIEnv* env, jobject jthis, jlong jhandle) {
   API_BEGIN()
   const auto* tensor_ptr = reinterpret_cast<torch::Tensor*>(jhandle);

diff --git a/examples/src/test/java/ai/djl/examples/training/TrainTimeSeriesTest.java b/examples/src/test/java/ai/djl/examples/training/TrainTimeSeriesTest.java
@@ -13,7 +13,7 @@
 
 package ai.djl.examples.training;
 
-import ai.djl.testing.TestRequirements;
+import ai.djl.engine.Engine;
 import ai.djl.training.TrainingResult;
 import ai.djl.translate.TranslateException;
 
@@ -26,10 +26,15 @@ public class TrainTimeSeriesTest {
 
     @Test
     public void testTrainTimeSeries() throws TranslateException, IOException {
-        TestRequirements.linux();
+        String[] args;
+        Engine engine = Engine.getEngine("PyTorch");
+        if (engine.getGpuCount() > 0) {
+            // TODO: PyTorch -- cuDNN error: CUDNN_STATUS_VERSION_MISMATCH
+            args = new String[] {"-g", "1", "-e", "5", "-b", "32", "--engine", "MXNet"};
+        } else {
+            args = new String[] {"-g", "1", "-e", "5", "-b", "32"};
+        }
 
-        // TODO: PyTorch -- PtNDArray.gammaln not implemented
-        String[] args = new String[] {"-g", "1", "-e", "5", "-b", "32", "--engine", "MXNet"};
         TrainingResult result = TrainTimeSeries.runExample(args);
         Assert.assertNotNull(result);
         float loss = result.getTrainLoss();

diff --git a/extensions/timeseries/src/main/java/ai/djl/timeseries/distribution/DistributionLoss.java b/extensions/timeseries/src/main/java/ai/djl/timeseries/distribution/DistributionLoss.java
@@ -57,7 +57,7 @@ public NDArray evaluate(NDList labels, NDList predictions) {
             NDArray lossWeights = predictions.get("loss_weights");
             NDArray weightedValue =
                     NDArrays.where(lossWeights.neq(0), loss.mul(lossWeights), loss.zerosLike());
-            NDArray sumWeights = lossWeights.sum().maximum(1.);
+            NDArray sumWeights = lossWeights.sum().maximum(1f);
             loss = weightedValue.sum().div(sumWeights);
         }
         return loss;