我有一个(N, T, d)形状的数组x.我有两个函数f和g，它们都采用形状数组(some_dimension, d)并返回形状数组(some_dimension, ).
我想把所有的x都算f.这很简单:f(x.reshape(-1, d)).
然后我只想在第二个维度的第一个切片上计算g，也就是g(x[:, 0, :])，然后将它减go 所有维度的fON的值.代码中举例说明了这一点
MWe--效率低下的方式
import numpy as np
# Reproducibility
seed = 1234
rng = np.random.default_rng(seed=seed)
# Generate x
N = 100
T = 10
d = 2
x = rng.normal(loc=0.0, scale=1.0, size=(N, T, d))
# In practice the functions are not this simple
def f(x):
return x[:, 0] + x[:, 1]
def g(x):
return x[:, 0]**2 - x[:, 1]**2
# Compute f on all the (flattened) array
fx = f(x.reshape(-1, d)).reshape(N, T)
# Compute g only on the first slice of second dimension. Here are two ways of doing so
gx = np.tile(g(x[:, 0])[:, None], reps=(1, T))
gx = np.repeat(g(x[:, 0]), axis=0, repeats=T).reshape(N, T)
# Finally compute what I really want to compute
diff = fx - gx
有没有更有效的方法？我觉得一定要用广播，但我想不通.