jax.flatten_util.ravel_pytree

jax.flatten_util.ravel_pytree(pytree)

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

参数:

pytree – 要展平的数组和标量的 pytree。

返回:

一个对，其中第一个元素是表示展平并连接的叶值的 1D 数组，其 dtype 由提升叶值的 dtype 确定，第二个元素是用于将相同长度的 1D 向量解展平回与输入 pytree 结构相同的 pytree 的可调用对象。如果输入 pytree 为空（即没有叶子），则按照惯例，在输出的第一个组件中返回 dtype 为 float32 的 1D 空数组。

有关 dtype 提升的详细信息，请参阅 https://jax.net.cn/en/latest/type_promotion.html。