[pytorch] Adds IValue Dict(str, IValue) support (#1765)

Change-Id: I277ff2f07383d6389c6865ba0e3e77aef3190d0a
deepjavalibrary · Jul 5, 2022 · a1c7448 · a1c7448
1 parent 0eaf1b5
commit a1c7448
Show file tree

Hide file tree

Showing 4 changed files with 56 additions and 0 deletions.
diff --git a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/IValue.java b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/IValue.java
@@ -287,6 +287,24 @@ public static IValue stringMapFrom(Map<String, PtNDArray> map) {
         return new IValue(PyTorchLibrary.LIB.iValueFromStringMap(keys, handles));
     }
 
+    /**
+     * Creates a new {@code IValue} of type {@code Map[String, IValue]}.
+     *
+     * @param map the Map[String, IValue] value
+     * @return a new {@code IValue} of type {@code Map[String, IValue]}
+     */
+    public static IValue stringIValueMapFrom(Map<String, IValue> map) {
+        String[] keys = new String[map.size()];
+        long[] handles = new long[map.size()];
+        int i = 0;
+        for (Map.Entry<String, IValue> entry : map.entrySet()) {
+            keys[i] = entry.getKey();
+            handles[i] = entry.getValue().getHandle();
+            ++i;
+        }
+        return new IValue(PyTorchLibrary.LIB.iValueFromStringIValueMap(keys, handles));
+    }
+
     /**
      * Returns the {@code boolean} value of this IValue.
      *

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
@@ -524,6 +524,8 @@ native long moduleLoad(
 
     native long iValueFromStringMap(String[] keys, long[] tensorHandles);
 
+    native long iValueFromStringIValueMap(String[] keys, long[] tensorHandles);
+
     native long iValueToTensor(long iValueHandle);
 
     native boolean iValueToBool(long iValueHandle);

diff --git a/engines/pytorch/pytorch-engine/src/test/java/ai/djl/pytorch/integration/IValueTest.java b/engines/pytorch/pytorch-engine/src/test/java/ai/djl/pytorch/integration/IValueTest.java
@@ -123,6 +123,18 @@ public void testIValue() {
                 Assert.assertEquals(list.get("data1"), array1);
             }
 
+            // (Dict(str, Tensor[])
+            Map<String, IValue> iValueMap = new ConcurrentHashMap<>();
+            try (IValue v1 = IValue.listFrom(array1);
+                    IValue v2 = IValue.listFrom(array2)) {
+                iValueMap.put("data1", v1);
+                iValueMap.put("data2", v2);
+                try (IValue ivalue = IValue.stringIValueMapFrom(iValueMap)) {
+                    Assert.assertTrue(ivalue.isMap());
+                    Assert.assertEquals(ivalue.getType(), "Dict(str, Tensor[])");
+                }
+            }
+
             try (IValue iv1 = IValue.from(1);
                     IValue iv2 = IValue.from(2);
                     IValue ivalue = IValue.listFrom(iv1, iv2)) {

diff --git a/engines/pytorch/pytorch-native/src/main/native/ai_djl_pytorch_jni_PyTorchLibrary_ivalue.cc b/engines/pytorch/pytorch-native/src/main/native/ai_djl_pytorch_jni_PyTorchLibrary_ivalue.cc
@@ -176,6 +176,30 @@ JNIEXPORT jlong JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_iValueFromStringM
   API_END_RETURN()
 }
 
+JNIEXPORT jlong JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_iValueFromStringIValueMap(
+    JNIEnv* env, jobject jthis, jobjectArray jkeys, jlongArray jvalues) {
+  API_BEGIN()
+  auto len = static_cast<size_t>(env->GetArrayLength(jvalues));
+  jlong* jptrs = env->GetLongArrayElements(jvalues, JNI_FALSE);
+  if (len == 0) {
+    const auto* ivalue_ptr = new torch::IValue{c10::impl::GenericDict(c10::StringType::get(), c10::TensorType::get())};
+    return reinterpret_cast<uintptr_t>(ivalue_ptr);
+  }
+
+  auto* firstEntryValue = reinterpret_cast<torch::IValue*>(jptrs[0]);
+  c10::impl::GenericDict dict(c10::StringType::get(), c10::unshapedType(firstEntryValue->type()));
+  for (size_t i = 0; i < len; ++i) {
+    auto jname = (jstring) env->GetObjectArrayElement(jkeys, i);
+    std::string name = djl::utils::jni::GetStringFromJString(env, jname);
+    dict.insert(name, *reinterpret_cast<torch::IValue*>(jptrs[i]));
+  }
+  env->ReleaseLongArrayElements(jvalues, jptrs, JNI_ABORT);
+  env->DeleteLocalRef(jkeys);
+  const auto* ivalue_ptr = new torch::IValue{dict};
+  return reinterpret_cast<uintptr_t>(ivalue_ptr);
+  API_END_RETURN()
+}
+
 JNIEXPORT jlong JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_iValueToTensor(
     JNIEnv* env, jobject jthis, jlong jhandle) {
   API_BEGIN()