如何有效地将矢量解码为numpy运算

Question

我正在尝试生成一种编码，该编码将向量映射到numpy数组之间的操作，并能有效地做到这一点。

例如，如果我有两个操作np.add&np.foo，我将np.add编码为0，将np.foo编码为1。从向量np.array([0,1,2](到np.add(1,2(的最快方法是什么

本质上是一种更有效的方法：

operation_list = [np.add, np.foo, np.witchcraft]
# In practice this array will be generated
# automatically
arr_of_ops = np.array([[0,1,2],  #np.add(1,2)
[1,1,2],  #np.foo(1,2)
[0,2,2],  #np.add(2,2)
[3,3,2]])  #np.witchcraft(3,2)
#This is the function I want to implement in an #efficient way as possible
def evaluate_encoding(arr_of_ops):
results = []
for row in arr_of_ops:
if row[0] == 0:
results.append(np.add(row[1],row[2]))
elif row[0] == 1:
results.append(np.foo(row[1],row[2]))
elif row[0] == 2:
results.append(np.witchcraft(row[1],row[2]))
return np.array(results)

Answer 1

一种方法是使用pandas，这是相当有效的，因为它也是矢量化的。

您可以通过操作将数据转换为数据帧并进行查询

import numpy as np
import pandas as pd
arr_of_ops = np.array([
[0, 1, 2],
[1, 1, 2],
[0, 2, 2],
[1, 3, 2]
])
df = pd.DataFrame(data=arr_of_ops, columns=("operation", "a", "b"))
indices_sum = df.operation.eq(0)
indices_diff = df.operation.eq(1)
df["results"] = None
df.results[indices_sum] = df.a[indices_sum] + df.b[indices_sum]
df.results[indices_diff] = df.a[indices_diff] - df.b[indices_diff]
print(df)
#    operation  a  b results
# 0          0  1  2     3.0
# 1          1  1  2    -1.0
# 2          0  2  2     4.0
# 3          1  3  2     1.0

Answer 2

1. 使用JAX，您就可以摆脱cpu/gpu/tpu的选择，XLA为您选择的硬件提供后端
2. 使用dict并丢失ifelse(以jax.numpy.op:op_code的形式(
3. 丢失循环，如果需要，可以使用jax.pmap在多台机器上映射数据上的操作