Adding "_getter" functions in python for TorchBind warnings (#152)

* expose irrep dtype for tensor product

* add getters to fake registration

* add getters to fake registration for conv

---------

Co-authored-by: Austin Glover <austin_glover@berkeley.com>

Author: asglover

openequivariance/extension/libtorch_tp_jit.cpp

@@ -134,7 +134,7 @@ class __attribute__ ((visibility ("default"))) TorchJITProduct : public torch::C
     Map_t fwd_dict, bwd_dict, dbl_bwd_dict, kernel_dims;
     JITTPImpl<JITKernel> internal;
     KernelProp kernelProp;
-    int64_t L3_dim; 
+    int64_t L3_dim, irrep_dtype; 
 
     TorchJITProduct(string kernel_plaintext, Map_t fwd_dict_i, Map_t bwd_dict_i, Map_t dbl_bwd_dict_i, Map_t kernel_dims_i) :
         fwd_dict(fwd_dict_i.copy()),
@@ -148,7 +148,8 @@ class __attribute__ ((visibility ("default"))) TorchJITProduct : public torch::C
                 to_map(kernel_dims_i)
             ),
         kernelProp(kernel_dims, false),
-        L3_dim(kernelProp.L3_dim) 
+        L3_dim(kernelProp.L3_dim),
+        irrep_dtype(kernel_dims_i.at("irrep_dtype")) 
         { }
 
     tuple<  tuple<string, string>, 
@@ -647,6 +648,7 @@ TORCH_LIBRARY_FRAGMENT(libtorch_tp_jit, m) {
             return 0;
         })
         .def_readonly("L3_dim", &TorchJITProduct::L3_dim)
+        .def_readonly("irrep_dtype", &TorchJITProduct::irrep_dtype)
         .def("__eq__", [](const c10::IValue & self, const c10::IValue& other) -> bool {
             return self.is(other); 
         })

openequivariance/implementations/LoopUnrollTP.py

@@ -183,9 +183,12 @@ def exec_tensor_product_rawptr(*args, **kwargs):
             def backward_rawptr(*args, **kwargs):
                 pass
 
-            def get_L3_dim(self):
+            def L3_dim_getter(self):
                 return self.kernel_dims["L3_dim"]
 
+            def irrep_dtype_getter(self):
+                return self.kernel_dims["irrep_dtype"]
+
         @torch.library.register_fake("libtorch_tp_jit::jit_tp_forward")
         def fake_forward(jit, L1_in, L2_in, W):
             L3_dim = None

openequivariance/implementations/convolution/LoopUnrollConv.py

@@ -296,6 +296,12 @@ def backward_rawptrs(*args, **kwargs):
             def double_backward_rawptrs(*args, **kwargs):
                 pass
 
+            def L3_dim_getter(self):
+                return self.kernel_dims["L3_dim"]
+
+            def irrep_dtype_getter(self):
+                return self.kernel_dims["irrep_dtype"]
+
         @torch.library.register_fake("libtorch_tp_jit::jit_conv_forward")
         def fake_forward(
             jit, L1_in, L2_in, W, rows, cols, workspace_buffer, sender_perm