Added __call__ functions.

Author: vbharadwaj-bk

openequivariance/openequivariance/impl_jax/TensorProduct.py

@@ -69,6 +69,12 @@ def __init__(self, config: TPProblem):
     def forward(self, X: jax.ndarray, Y: jax.ndarray, W: jax.ndarray) -> jax.ndarray:
         return forward(X, Y, W, self.L3_dim, self.config.irrep_dtype, self.attrs)
 
+    def __call__(self,
+            X: jax.numpy.ndarray, 
+            Y: jax.numpy.ndarray, 
+            W: jax.numpy.ndarray) -> jax.numpy.ndarray:
+        return self.forward(X, Y, W)
+
 
 def jax_to_torch(x):
     import numpy as np

openequivariance/openequivariance/impl_jax/TensorProductConv.py

@@ -55,7 +55,7 @@ def __init__(self, config: TPProblem, deterministic: bool = False, kahan: bool =
         super().__init__(
             config,
             dp, extlib.postprocess_kernel,
-            idx_dtype=np.int32, # Note: this is distinct from PyTorch 
+            idx_dtype=np.int32, # N.B. this is distinct from the PyTorch version 
             torch_op=False,
             deterministic=deterministic,
             kahan=kahan
@@ -99,6 +99,15 @@ def forward(
             self.L3_dim, 
             self.config.irrep_dtype, 
             self.attrs)
+    
+    def __call__(self,
+            X: jax.numpy.ndarray, 
+            Y: jax.numpy.ndarray, 
+            W: jax.numpy.ndarray, 
+            rows: jax.numpy.ndarray, 
+            cols: jax.numpy.ndarray, 
+            sender_perm: Optional[jax.numpy.ndarray] = None) -> jax.numpy.ndarray:
+        return self.forward(X, Y, W, rows, cols, sender_perm)
 
 if __name__=="__main__":
     X_ir, Y_ir, Z_ir = Irreps("1x2e"), Irreps("1x3e"), Irreps("1x2e")