jax.flatten_util 模块

函数列表

ravel_pytree(pytree)

将数组的 pytree 展平（扁平化）为 1D 数组。