用 jax 计算逐行（或逐轴）点积的最佳方法是什么？答案

【问题标题】：What's the best way to compute row-wise (or axis-wise) dot products with jax?用 jax 计算逐行（或逐轴）点积的最佳方法是什么？
【发布时间】：2020-04-20 02:54:50
【问题描述】：

我有两个形状为 (N, M) 的数值数组。我想计算一个逐行的点积。 IE。产生一个形状为 (N,) 的数组，使得第 n 行是每个数组的第 n 行的点积。

我知道 numpy 的 inner1d 方法。使用 jax 执行此操作的最佳方法是什么？ jax 有jax.numpy.inner，但这有别的作用。

【问题讨论】：

标签： numpy linear-algebra jax

【解决方案1】：

您可以在几行 jax 代码中定义自己的 jit 编译版本的 inner1d：

import jax
@jax.jit
def inner1d(X, Y):
  return (X * Y).sum(-1)

测试一下：

import jax.numpy as jnp
import numpy as np
from numpy.core import umath_tests


X = np.random.rand(5, 10)
Y = np.random.rand(5, 10)

print(umath_tests.inner1d(X, Y))
print(inner1d(jnp.array(X), jnp.array(Y)))
# [2.23219571 2.1013316  2.70353783 2.14094973 2.62582531]
# [2.2321959 2.1013315 2.703538  2.1409497 2.6258256]

【讨论】：

【解决方案2】：

你可以试试jax.numpy.einsum。这里使用 numpy einsum 实现

import numpy as np
from numpy.core.umath_tests import inner1d

arr1 = np.random.randint(0,10,[5,5])
arr2 = np.random.randint(0,10,[5,5])

arr = np.inner1d(arr1,arr2)
arr
array([ 87, 200, 229,  81,  53])
np.einsum('...i,...i->...',arr1,arr2)
array([ 87, 200, 229,  81,  53])

【讨论】：