Current Progress: After final design presentation

examples/jax_playground.py

@@ -1,50 +1,57 @@
 import jax
 import jax.numpy as jnp
 from jax import jit, vmap
-from time_utils import using_cprofile
 from time import time
-#
 import numpy as np
+
+
 @jit
-def fx(x):
-    return jnp.arange(x, x + 10)
-#
-#
-# # @jit
-# def fy(z):
-#     z1, z2 = z, z + 1
-#     vmap_fx = vmap(fx)
-#     return vmap_fx(z1, z2)
-#
-# @jit
-# def test_while(num, init_val):
-#     def cond_fun(carry):
-#         i, cumsum = carry
-#         return i < num
-#
-#     def body_fun(carry):
-#         i, cumsum = carry
-#         cumsum += i
-#         return i + 1, cumsum
-#
-#     return jax.lax.while_loop(cond_fun, body_fun, (0, init_val))
+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
 
 
-# @using_cprofile
+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
+
 def main():
-    print(fx(1))
+    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()
 
-    # vmap_f = vmap(fx, in_axes=(None, 0))
-    # vmap_vmap_f = vmap(vmap_f, in_axes=(0, None))
-    # a = jnp.array([20,10,30])
-    # b = jnp.array([6, 5, 4])
-    # res = vmap_vmap_f(a, b)
-    # print(res)
-    # print(jnp.argmin(res, axis=1))
+        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)
 
 if __name__ == '__main__':
     main()