根据矩阵的行列,按照参数in_axes进行矩阵的行列操作,然后根据out_axes的输出轴进行优先输出,一般可以理解为按照行输出优先输出,或者按照列输出优先输出。具体可以看列子。
import jax.numpy as jnp
import numpy as np
import jax
a = np.array(([1,3],[23,-5]))
print(a)
#[[ 1 3]
# [23 -5]]
b = np.array(([11,7],[19,13]))
print(b)
#[[11 7]
#[19 13]]
print(jnp.add(a,b))
#正常的两个矩阵element-wise的相加
#[[1+11, 3+7]]
#[[23+19, -5+13]]
#----------------
#[[12 10]
#[42 8]]
print(jax.vmap(jnp.add, in_axes=(0,0), out_axes=0)(a,b))
#矩阵a的行 + 矩阵b的行,然后根据out_axes=0输出,0表示行输出
#[[1+11, 3+7]]
#[[23+19, -5+13]]
#----------------
#[[12 10]
#[42 8]]
print(jax.vmap(jnp.add, in_axes=(0,0), out_axes=1)(a,b))
#矩阵a的行 + 矩阵b的行,然后根据out_axes=1输出,1表示列输出
#[[1+11, 3+7]]
#[[23+19, -5+13]]
#---------类似转置,以列先输出
#[[1+11, 23+19]]
#[[3+7, -5+13]]
#----------------
#[[12 42]
#[10 8]]
print(jax.vmap(jnp.add, in_axes=(0,1), out_axes=0)(a,b))
#矩阵a的行 + 矩阵b的列,然后根据out_axes=1输出,1表示列输出
#[[1+11, 3+19]]
#[[23+7, -5+13]]
#--------以行先输出,可以认为直接输出不进行任何操作
#[[1+11, 3+19]]
#[[23+7, -5+13]]
#----------------
#[[12 22]
#[30 8]]
print(jax.vmap(jnp.add, in_axes=(1,0), out_axes=-2)(a,b))
#矩阵a的列 + 矩阵b的行,然后根据out_axes=-2输出,-2相当于0,行输出
#[[1+11, 23+7]]
#[[3+19, -5+13]]
#--------以行先输出,可以认为直接输出不进行任何操作
#[[1+11, 23+7]]
#[[3+19, -5+13]]
#----------------
#[[12 30]
#[22 8]]