reduce the use of device_get in speciate

examples/jax_playground.py

@@ -8,56 +8,14 @@ from functools import partial
 from examples.time_utils import using_cprofile
 
 
-def func(x, y):
-    """
-    :param x: (100, )
-    :param y: (100,
-    :return:
-    """
-    return x * y
-
-
-def func2(x, y, s):
-    """
-    :param x: (100, )
-    :param y: (100,
-    :return:
-    """
-    if s == '123':
-        return 0
-    else:
-        return x * y
-
-
 @jit
-def func3(x, y):
-    return func2(x, y, '123')
+def func(x, y):
+    return x + y
 
 
-# @using_cprofile
-def main():
-    key = jax.random.PRNGKey(42)
+a, b, c = jnp.array([1]), jnp.array([2]), jnp.array([3])
+li = [a, b, c]
 
-    x1, y1 = jax.random.normal(key, shape=(1000,)), jax.random.normal(key, shape=(1000,))
+cpu_li = jax.device_get(li)
 
-    jit_lower_func = jit(func).lower(1, 2).compile()
-    print(type(jit_lower_func))
-    print(jit_lower_func.memory_analysis())
-
-    jit_compiled_func2 = jit(func2, static_argnames=['s']).lower(x1, y1, '123').compile()
-    print(jit_compiled_func2(x1, y1))
-
-    # print(jit_compiled_func2(x1, y1))
-
-    f = func3.lower(x1, y1).compile()
-
-    print(f(x1, y1))
-
-    # print(jit_lower_func(x1, y1))
-
-    # x2, y2 = jax.random.normal(key, shape=(200,)), jax.random.normal(key, shape=(200,))
-    # print(jit_lower_func(x2, y2))
-
-
-if __name__ == '__main__':
-    main()
+print(cpu_li)