remove create_func....

utils/__init__.py

@@ -1,4 +1,35 @@
-from .activation import Activation
-from .aggregation import Aggregation
+from .activation import Activation, act
+from .aggregation import Aggregation, agg
 from .tools import *
-from .graph import *
+from .graph import *
+
+Activation.name2func = {
+    'sigmoid': Activation.sigmoid_act,
+    'tanh': Activation.tanh_act,
+    'sin': Activation.sin_act,
+    'gauss': Activation.gauss_act,
+    'relu': Activation.relu_act,
+    'elu': Activation.elu_act,
+    'lelu': Activation.lelu_act,
+    'selu': Activation.selu_act,
+    'softplus': Activation.softplus_act,
+    'identity': Activation.identity_act,
+    'clamped': Activation.clamped_act,
+    'inv': Activation.inv_act,
+    'log': Activation.log_act,
+    'exp': Activation.exp_act,
+    'abs': Activation.abs_act,
+    'hat': Activation.hat_act,
+    'square': Activation.square_act,
+    'cube': Activation.cube_act,
+}
+
+Aggregation.name2func = {
+    'sum': Aggregation.sum_agg,
+    'product': Aggregation.product_agg,
+    'max': Aggregation.max_agg,
+    'min': Aggregation.min_agg,
+    'maxabs': Aggregation.maxabs_agg,
+    'median': Aggregation.median_agg,
+    'mean': Aggregation.mean_agg,
+}

utils/activation.py

@@ -1,8 +1,8 @@
+import jax
 import jax.numpy as jnp
 
 
 class Activation:
-
     name2func = {}
 
     @staticmethod
@@ -89,23 +89,11 @@ class Activation:
         return z ** 3
 
 
-Activation.name2func = {
-    'sigmoid': Activation.sigmoid_act,
-    'tanh': Activation.tanh_act,
-    'sin': Activation.sin_act,
-    'gauss': Activation.gauss_act,
-    'relu': Activation.relu_act,
-    'elu': Activation.elu_act,
-    'lelu': Activation.lelu_act,
-    'selu': Activation.selu_act,
-    'softplus': Activation.softplus_act,
-    'identity': Activation.identity_act,
-    'clamped': Activation.clamped_act,
-    'inv': Activation.inv_act,
-    'log': Activation.log_act,
-    'exp': Activation.exp_act,
-    'abs': Activation.abs_act,
-    'hat': Activation.hat_act,
-    'square': Activation.square_act,
-    'cube': Activation.cube_act,
-}
+def act(idx, z, act_funcs):
+    """
+    calculate activation function for each node
+    """
+    idx = jnp.asarray(idx, dtype=jnp.int32)
+    # change idx from float to int
+    res = jax.lax.switch(idx, act_funcs, z)
+    return res

utils/aggregation.py

@@ -1,8 +1,8 @@
+import jax
 import jax.numpy as jnp
 
 
 class Aggregation:
-
     name2func = {}
 
     @staticmethod
@@ -52,12 +52,16 @@ class Aggregation:
         return mean_without_zeros
 
 
-Aggregation.name2func = {
-    'sum': Aggregation.sum_agg,
-    'product': Aggregation.product_agg,
-    'max': Aggregation.max_agg,
-    'min': Aggregation.min_agg,
-    'maxabs': Aggregation.maxabs_agg,
-    'median': Aggregation.median_agg,
-    'mean': Aggregation.mean_agg,
-}
+def agg(idx, z, agg_funcs):
+    """
+    calculate activation function for inputs of node
+    """
+    idx = jnp.asarray(idx, dtype=jnp.int32)
+
+    def all_nan():
+        return 0.
+
+    def not_all_nan():
+        return jax.lax.switch(idx, agg_funcs, z)
+
+    return jax.lax.cond(jnp.all(jnp.isnan(z)), all_nan, not_all_nan)