jax.flatten_util 模块

函数列表

ravel_pytree(pytree)

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