jax的vmap用法详解

根据矩阵的行列，按照参数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]]