finish jit-able speciate function

next time i'll create a new branch

examples/jax_playground.py

@@ -3,56 +3,45 @@ import jax.numpy as jnp
 from jax import jit, vmap
 from time_utils import using_cprofile
 from time import time
-
+#
 @jit
 def fx(x, y):
     return x + y
-
-
+#
+#
+# # @jit
+# def fy(z):
+#     z1, z2 = z, z + 1
+#     vmap_fx = vmap(fx)
+#     return vmap_fx(z1, z2)
+#
 # @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 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))
 
 
 
-@using_cprofile
+
+# @using_cprofile
 def main():
-    z = jnp.zeros((100000, ))
+    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))
 
-    num = 100
 
-    nums = jnp.arange(num) * 10
-
-    f = jit(vmap(test_while, in_axes=(0, None))).lower(nums, z).compile()
-    def test_time(*args):
-        return f(*args)
-
-    print(test_time(nums, z))
-
-    #
-    #
-    # for i in range(10):
-    #     num = 10 ** i
-    #     st = time()
-    #     res = test_time(num, z)
-    #     print(res)
-    #     t = time() - st
-    #     print(f"num: {num}, time: {t}")
 
 if __name__ == '__main__':
     main()