jax.flatten_util 模块

函数列表

ravel_pytree(pytree)

将一个数组的 pytree 展平为 1D 数组。