generally complete, but not work well. Debug

examples/jax_playground.py

@@ -5,33 +5,10 @@ from jax import random
 from jax import vmap, jit
 
 
-def plus1(x):
-    return x + 1
+seed = jax.random.PRNGKey(42)
+seed, *subkeys = random.split(seed, 3)
 
 
-def minus1(x):
-    return x - 1
-
-
-def func(rand_key, x):
-    r = jax.random.uniform(rand_key, shape=())
-    return jax.lax.cond(r > 0.5, plus1, minus1, x)
-
-
-def func2(rand_key):
-    r = jax.random.uniform(rand_key, ())
-    if r < 0.3:
-        return 1
-    elif r < 0.5:
-        return 2
-    else:
-        return 3
-
-
-
-key = random.PRNGKey(0)
-print(func(key, 0))
-
-batch_func = vmap(jit(func))
-keys = random.split(key, 100)
-print(batch_func(keys, jnp.zeros(100)))
+c = random.split(seed, 1)
+print(seed, subkeys)
+print(c)

examples/time_utils.py

@@ -0,0 +1,28 @@
+import cProfile
+from io import StringIO
+import pstats
+
+
+def using_cprofile(func, root_abs_path=None, replace_pattern=None, save_path=None):
+    def inner(*args, **kwargs):
+        pr = cProfile.Profile()
+        pr.enable()
+        ret = func(*args, **kwargs)
+        pr.disable()
+        profile_stats = StringIO()
+        stats = pstats.Stats(pr, stream=profile_stats)
+        if root_abs_path is not None:
+            stats.sort_stats('cumulative').print_stats(root_abs_path)
+        else:
+            stats.sort_stats('cumulative').print_stats()
+        output = profile_stats.getvalue()
+        if replace_pattern is not None:
+            output = output.replace(replace_pattern, "")
+        if save_path is None:
+            print(output)
+        else:
+            with open(save_path, "w") as f:
+                f.write(output)
+        return ret
+
+    return inner

examples/xor.py

@@ -1,10 +1,12 @@
 from typing import Callable, List
+from functools import partial
 
 import jax
 import numpy as np
 
 from utils import Configer
 from algorithms.neat import Pipeline
+from time_utils import using_cprofile
 
 xor_inputs = np.array([[0, 0], [0, 1], [1, 0], [1, 1]])
 xor_outputs = np.array([[0], [1], [1], [0]])
@@ -17,22 +19,24 @@ def evaluate(forward_func: Callable) -> List[float]:
     """
     outs = forward_func(xor_inputs)
     outs = jax.device_get(outs)
-    fitnesses = np.mean((outs - xor_outputs) ** 2, axis=(1, 2))
+    fitnesses = -np.mean((outs - xor_outputs) ** 2, axis=(1, 2))
+    # print(fitnesses)
     return fitnesses.tolist()  # returns a list
 
 
+# @using_cprofile
+@partial(using_cprofile, root_abs_path='/mnt/e/neat-jax/', replace_pattern="/mnt/e/neat-jax/")
 def main():
     config = Configer.load_config()
     pipeline = Pipeline(config)
-    forward_func = pipeline.ask(batch=True)
-    fitnesses = evaluate(forward_func)
-    pipeline.tell(fitnesses)
+    pipeline.auto_run(evaluate)
 
-
-    # for i in range(100):
+    # for _ in range(100):
+    #     s = time.time()
     #     forward_func = pipeline.ask(batch=True)
     #     fitnesses = evaluate(forward_func)
     #     pipeline.tell(fitnesses)
+    #     print(time.time() - s)
 
 
 if __name__ == '__main__':