jax树形结构展平

Pytree 是 Haiku 和 JAX 中一种用于表示具有树结构的数据的通用方式。这个术语来源于函数式编程中的树形数据结构。在 JAX 中，Pytree 提供了一种通用的方式来处理和操作树形结构的数据，这对于深度学习和自动微分非常有用。在 Pytree 中，数据结构可以是由嵌套的元组、列表、字典等组成的树形结构，其中每个节点都可以包含标量、数组或其他 Pytree。这种灵活性使得 Pytree 能够表示和处理各种复杂的数据结构。

import jax
import jax.numpy as jnp

# 定义一个 Pytree
pytree_example = {
    'a': [1, 2, 3],
    'b': (4, 5, 6),
    'c': {
        'x': 7,
        'y': [8, 9]
    }
}


### 1. jax.tree_util.tree_flatten

# jax.tree_util.tree_flatten 将树形结构展平为一维的列表，并输出展平后的数据和结构信息。
# 展平后的数据可以用于进行一维的操作，而结构信息可以用于还原原始的树形结构。
flattened_list, values_tree_def = jax.tree_util.tree_flatten(pytree_example)

# 输出展平后的列表和原始结构
print("原始数据为:")
print(pytree_example)
print("展平后的数据为:")
print(flattened_list)
print("数据结构为:")
print(values_tree_def)

### 2. jax.tree_util.tree_unflatten
# 使用 tree_unflatten 还原成原始张量
restored_tensor = jax.tree_util.tree_un