Support Retrieving Named Attributes

src/wrappers/jit.rs

@@ -2,7 +2,7 @@
 use super::utils::{path_to_cstring, ptr_to_string};
 use super::{device::Device, kind::Kind};
 use crate::{nn::Path, TchError, Tensor};
-use libc::{c_int, c_void};
+use libc::{c_char, c_int, c_void};
 use std::borrow::Borrow;
 use std::convert::TryFrom;
 use torch_sys::*;
@@ -403,6 +403,14 @@ impl IValue {
         }
         Ok(v)
     }
+
+    pub(super) extern "C" fn add_callback(data: *mut c_void, name: *const c_char, c_ivalue: *mut CIValue) {
+        let v = unsafe { &mut *(data as *mut Vec<(String, Result<IValue, TchError>)>) };
+        let name = unsafe { std::ffi::CStr::from_ptr(name).to_str().unwrap() };
+        let name = name.replace('|', ".");
+        let value = Self::from_c(c_ivalue);
+        v.push((name, value));
+    }
 }
 
 /// A jit PyTorch module.
@@ -593,6 +601,17 @@ impl CModule {
         Ok(v)
     }
 
+    /// Loads the named attributes on a module
+    pub fn named_attributes(&self) -> Result<Vec<(String, IValue)>, TchError> {
+        let mut v: Vec<(String, Result<IValue, TchError>)> = vec![];
+        unsafe_torch_err!(atm_named_attributes(
+            self.c_module,
+            &mut v as *mut _ as *mut c_void,
+            IValue::add_callback
+        ));
+        v.into_iter().map(|(k, v)| Ok((k, v?))).collect()
+    }
+
     /// Create a new module by tracing the application of the specified function on
     /// the given inputs.
     pub fn create_by_tracing<F>(

tests/jit_tests.rs

@@ -76,6 +76,23 @@ fn jit4() {
     assert_eq!(v, 14.0);
     let named_parameters = mod_.named_parameters().unwrap();
     assert_eq!(named_parameters, vec![]);
+
+    // Even "empty" models have some attributes.
+    let named_attributes = mod_.named_attributes().unwrap();
+    assert_eq!(named_attributes, vec![
+        (String::from("training"), IValue::Bool(true)),
+        (String::from("_is_full_backward_hook"), IValue::None),
+    ]);
+}
+
+#[test]
+fn jit_named_attributes() {
+    // Check that models with user-defined attributes are correctly loaded.
+    let mod_ = tch::CModule::load("tests/foo9.pt").unwrap();
+    let named_attributes = mod_.named_attributes().unwrap();
+    assert!(named_attributes.len() > 2);
+    assert!(named_attributes.iter().any(|(name, _)| name == "embedding_length"));
+    assert!(named_attributes.iter().any(|(name, _)| name == "max_distance"));
 }
 
 #[test]

torch-sys/libtch/torch_api.cpp

@@ -1317,6 +1317,15 @@ void atm_named_parameters(module m, void *data, void (*f)(void *, char *, tensor
   )
 }
 
+void atm_named_attributes(module m, void *data, void (*f)(void *, char *, ivalue)) {
+  PROTECT(
+    for (const auto &p : m->named_attributes()) {
+      auto v = p.value;
+      f(data, (char*)p.name.c_str(), new torch::IValue(v));
+    }
+  )
+}
+
 ivalue ati_tensor(tensor t) {
   PROTECT(
     return new torch::jit::IValue(*t);

torch-sys/libtch/torch_api.h

@@ -226,6 +226,7 @@ void atm_set_profiling_mode(int);
 void atm_fuser_cuda_set_enabled(bool);
 bool atm_fuser_cuda_is_enabled();
 void atm_named_parameters(module, void *data, void (*f)(void *, char *, tensor));
+void atm_named_attributes(module, void *data, void (*f)(void *, char *, ivalue));
 
 // This function has to be followed by a call to atm_end_tracing.
 module atm_create_for_tracing(char *modl_name, tensor *inputs, int ninputs);

torch-sys/src/lib.rs

@@ -319,6 +319,11 @@ extern "C" {
         data: *mut c_void,
         f: extern "C" fn(*mut c_void, name: *const c_char, t: *mut C_tensor),
     );
+    pub fn atm_named_attributes(
+        m: *mut CModule_,
+        data: *mut c_void,
+        f: extern "C" fn(*mut c_void, name: *const c_char, t: *mut CIValue),
+    );
     pub fn atm_create_for_tracing(
         modl_name: *const c_char,
         inputs: *const *mut C_tensor,