Current Progress: After final design presentation

examples/jax_playground.py

@@ -1,59 +1,21 @@
+from functools import partial
+
 import jax
-import jax.numpy as jnp
-from jax import jit, vmap
-from time import time
-import numpy as np
+from jax import jit
 
-
-@jit
-def jax_mutate(seed, x):
-    noise = jax.random.normal(seed, x.shape) * 0.1
-    return x + noise
-
-
-def numpy_mutate(x):
-    noise = np.random.normal(size=x.shape) * 0.1
-    return x + noise
-
-
-def jax_mutate_population(seed, pop_x):
-    seeds = jax.random.split(seed, len(pop_x))
-    func = vmap(jax_mutate, in_axes=(0, 0))
-    return func(seeds, pop_x)
-
-
-def numpy_mutate_population(pop_x):
-    return np.stack([numpy_mutate(x) for x in pop_x])
-
-
-def numpy_mutate_population_vmap(pop_x):
-    noise = np.random.normal(size=pop_x.shape) * 0.1
-    return pop_x + noise
+from configs import Configer
+from neat.pipeline_ import Pipeline
 
 
 def main():
-    seed = jax.random.PRNGKey(0)
-    i = 10
-    while i < 200000:
-        pop_x = jnp.ones((i, 100, 100))
-        jax_pop_func = jit(jax_mutate_population).lower(seed, pop_x).compile()
+    config = Configer.load_config("xor.ini")
+    print(config)
+    pipeline = Pipeline(config)
 
-        tic = time()
-        res = jax.device_get(jax_pop_func(seed, pop_x))
-        jax_time = time() - tic
 
-        tic = time()
-        res = numpy_mutate_population(pop_x)
-        numpy_time = time() - tic
-
-        tic = time()
-        res = numpy_mutate_population_vmap(pop_x)
-        numpy_time_vmap = time() - tic
-
-        # print(f'POP_SIZE: {i} | JAX: {jax_time:.4f} | Numpy: {numpy_time:.4f} | Speedup: {numpy_time / jax_time:.4f}')
-        print(f'POP_SIZE: {i} | JAX: {jax_time:.4f} | Numpy: {numpy_time:.4f} | Numpy Vmap: {numpy_time_vmap:.4f}')
-
-        i = int(i * 1.3)
+@jit
+def f(x, jit_config):
+    return x + jit_config["bias_mutate_rate"]
 
 
 if __name__ == '__main__':