small bugfix

Author: jjcmoon

src/circuit.cpp

@@ -504,7 +504,8 @@ m.doc() = "Layerize arithmetic circuits";
 nb::class_<NodePtr>(m, "NodePtr")
 .def("__repr__", &NodePtr::to_string)
 .def(nb::self == nb::self)
-.def("__hash__", &NodePtr::as_int);
+.def("__hash__", &NodePtr::as_int)
+.def("get_ix", [](NodePtr a) {return a.get()->ix;});
 
 nb::class_<Circuit>(m, "Circuit", "Circuits are the main class added by KLay, and require no arguments to construct.\n\n:code:`circuit = klay.Circuit()` ")
 .def(nb::init<>())

src/klay/backends/torch_backend.py

@@ -129,7 +129,7 @@ def forward(self, x):
     def get_edge_weights(self):
         exp_weights, _ = self._safe_exp(self.weights)
         norm = self._scatter_reduce(exp_weights, "sum")
-        return exp_weights / norm
+        return exp_weights / norm[self.csr]
 
 
 class ProbabilisticLogSumLayer(ProbabilisticKnowledgeLayer):
@@ -139,7 +139,7 @@ def forward(self, x, epsilon=10e-16):
 
     def get_edge_weights(self, epsilon):
         norm = self._logsumexp_scatter_reduce(self.weights, epsilon)
-        return self.weights - norm
+        return self.weights - norm[self.csr]
 
 
 def get_semiring(name: str, probabilistic: bool):