JAX的运行速度比NumPy慢吗

Question

我最近开始学习JAX。我用NumPy编写了一个简短的代码片段，并用JAX编写了相应的代码片段。我原以为JAX会更快，但当我评测代码时，NumPy代码比JAX代码快得多。我想知道这是否普遍正确，或者我的实现中是否存在问题。

NumPy代码：

import numpy as np
from time import time as tm

def gp(x):
return np.maximum(np.zeros(x.shape), x)

# -- inputs
n_q1 = 25
n_q2 = 5
x = np.random.rand(1, n_q1)  # todo: changes w/ time
y = np.random.rand(1, n_q2)  # todo: changes w/ time
# -- activations
n_p1 = 3
n_p2 = 2
v_p1 = np.random.rand(1, n_p1)
v_p2 = np.random.rand(1, n_p2)
a_p1 = 0.5
a_p2 = 0.5
# -- weights
W_q1p1 = np.random.rand(n_q1, n_p1)
W_p2q2 = np.random.rand(n_p2, n_q2)
W_p1p1 = np.random.rand(n_p1, n_p1)
W_p1p2 = np.random.rand(n_p1, n_p2)
W_p2p1 = np.random.rand(n_p2, n_p1)
# -- computation
t1=tm()
for t in range(2000):
z_p1 = np.matmul(v_p1, W_p1p1) + np.matmul(v_p2, W_p2p1) + np.matmul(x, W_q1p1)
v_p1_new = a_p1 * v_p1 + (1 - a_p1) * gp(z_p1)
z_p2 = np.matmul(v_p1, W_p1p2)
v_p2_new = a_p2 * v_p2 + (1 - a_p2) * gp(z_p2)
v_p1, v_p2 = v_p1_new, v_p2_new
print(tm()-t1)

这产生：0.02118515968322754

JAX代码：

from jax import random, nn, numpy as jnp
from time import time as tm

def gp(x):
return nn.relu(x)

# -- inputs
n_q1 = 25
n_q2 = 5
key = random.PRNGKey(0)
x = random.normal(key, (1, n_q1))
y = random.normal(key, (1, n_q2))  # todo: check if I need to advance "key" manually
# -- activations
n_p1 = 3
n_p2 = 2
v_p1 = random.normal(key, (1, n_p1))
v_p2 = random.normal(key, (1, n_p2))
a_p1 = 0.5
a_p2 = 0.5
# -- weights
W_q1p1 = random.normal(key, (n_q1, n_p1))
W_p2q2 = random.normal(key, (n_p2, n_q2))
W_p1p1 = random.normal(key, (n_p1, n_p1))
W_p1p2 = random.normal(key, (n_p1, n_p2))
W_p2p1 = random.normal(key, (n_p2, n_p1))
# -- computation
t1=tm()
for t in range(2000):
z_p1 = jnp.matmul(v_p1, W_p1p1) + jnp.matmul(v_p2, W_p2p1) + jnp.matmul(x, W_q1p1)
v_p1_new = a_p1 * v_p1 + (1 - a_p1) * gp(z_p1)
z_p2 = jnp.matmul(v_p1, W_p1p2)
v_p2_new = a_p2 * v_p2 + (1 - a_p2) * gp(z_p2)
v_p1, v_p2 = v_p1_new, v_p2_new

print(tm()-t1)

这产生：2.5548229217529297

Answer 1

JAX文档在其常见问题解答中有一个有用的部分：https://jax.readthedocs.io/en/latest/faq.html#is-jax比numpy 快

TL；DR：很复杂。对于CPU上的单个矩阵操作，JAX通常比NumPy慢，但JAX中JIT编译的操作序列通常比NumPy快，并且一旦转到GPU/TPU，JAX一般会比NumPy快得多。