Specialized ivalue dicts using string keys and tensor values. (#598)

* Specialized ivalue dicts using string keys and tensor values. * Clippy fixes.
LaurentMazare · Dec 27, 2022 · 51cc937 · 51cc937
1 parent 23db431
commit 51cc937
Show file tree

Hide file tree

Showing 12 changed files with 47 additions and 4 deletions.
diff --git a/tests/create_jit_models.py b/tests/create_jit_models.py
@@ -1,4 +1,5 @@
 import torch
+from typing import Dict, Tuple
 
 class Foo(torch.jit.ScriptModule):
     def __init__(self, v):
@@ -107,3 +108,12 @@ def make_input_object(self, foo, bar):
 
 foo_7 = TorchScriptExample()
 foo_7.save("foo7.pt")
+
+# https://github.com/LaurentMazare/tch-rs/issues/597
+class DictExample(torch.jit.ScriptModule):
+    @torch.jit.script_method
+    def generate(self, batch: Dict[str, torch.Tensor]) -> Tuple[torch.Tensor, torch.Tensor]:
+        return batch["foo"], batch["bar"]
+
+foo_8 = DictExample()
+foo_8.save("foo8.pt")
diff --git a/tests/foo.pt b/tests/foo.pt
diff --git a/tests/foo1.pt b/tests/foo1.pt
diff --git a/tests/foo2.pt b/tests/foo2.pt
diff --git a/tests/foo3.pt b/tests/foo3.pt
diff --git a/tests/foo4.pt b/tests/foo4.pt
diff --git a/tests/foo5.pt b/tests/foo5.pt
diff --git a/tests/foo6.pt b/tests/foo6.pt
diff --git a/tests/foo7.pt b/tests/foo7.pt
diff --git a/tests/foo8.pt b/tests/foo8.pt
diff --git a/tests/jit_tests.rs b/tests/jit_tests.rs
@@ -1,4 +1,4 @@
-use std::convert::TryFrom;
+use std::convert::{TryFrom, TryInto};
 use tch::{IValue, Kind, Tensor};
 
 #[test]
@@ -160,3 +160,20 @@ fn jit_double_free() {
     };
     assert_eq!(Vec::<f64>::from(&result), [5.0, 7.0, 9.0])
 }
+
+// https://github.com/LaurentMazare/tch-rs/issues/597
+#[test]
+fn specialized_dict() {
+    let mod_ = tch::CModule::load("tests/foo8.pt").unwrap();
+    let input = IValue::GenericDict(vec![
+        (IValue::String("bar".to_owned()), IValue::Tensor(Tensor::of_slice(&[1_f32, 7_f32]))),
+        (
+            IValue::String("foo".to_owned()),
+            IValue::Tensor(Tensor::of_slice(&[1_f32, 2_f32, 3_f32])),
+        ),
+    ]);
+    let result = mod_.method_is("generate", &[input]).unwrap();
+    let result: (Tensor, Tensor) = result.try_into().unwrap();
+    assert_eq!(Vec::<f64>::from(&result.0), [1.0, 2.0, 3.0]);
+    assert_eq!(Vec::<f64>::from(&result.1), [1.0, 7.0])
+}
diff --git a/torch-sys/libtch/torch_api.cpp b/torch-sys/libtch/torch_api.cpp
@@ -1212,11 +1212,27 @@ ivalue ati_generic_list(ivalue *is, int nvalues) {
   return nullptr;
 }
 
+using generic_dict = c10::Dict<torch::jit::IValue, torch::jit::IValue>;
+
 ivalue ati_generic_dict(ivalue *is, int nvalues) {
-  c10::Dict<torch::jit::IValue, torch::jit::IValue> dict(c10::AnyType::get(), c10::AnyType::get());
   PROTECT(
-    for (int i = 0; i < nvalues; ++i) dict.insert(*(is[2*i]), *(is[2*i+1]));
-    return new torch::jit::IValue(dict);
+    bool all_keys_are_str = true;
+    for (int i = 0; i < nvalues; ++i) {
+        if (!is[2*i]->isString()) all_keys_are_str = false;
+    }
+    bool all_values_are_tensor = true;
+    for (int i = 0; i < nvalues; ++i) {
+        if (!is[2*i+1]->isTensor()) all_values_are_tensor = false;
+    }
+    if (all_keys_are_str && all_values_are_tensor) {
+      generic_dict dict(c10::StringType::get(), c10::TensorType::get());
+      for (int i = 0; i < nvalues; ++i) dict.insert(is[2*i]->toString(), is[2*i+1]->toTensor());
+      return new torch::jit::IValue(dict);
+    } else {
+      generic_dict dict(c10::AnyType::get(), c10::AnyType::get());
+      for (int i = 0; i < nvalues; ++i) dict.insert(*(is[2*i]), *(is[2*i+1]));
+      return new torch::jit::IValue(dict);
+    }
   )
   return nullptr;
 }