如何创建Jax数组并在后续将字典存入其中？

直接把原生字典存入Jax数组不可行——因为Jax数组是同构数据结构，要求所有元素的类型、形状严格一致，而字典属于异构的键值对结构，无法直接适配Jax数组的存储规则。不过可以通过以下几种方式实现类似需求：

1. 转换为Jax结构化数组 #

如果字典的键值对可以映射为固定的字段和对应类型，可将字典转为Jax结构化数组，这种方式能保留键的语义，同时符合Jax数组的要求。

示例代码：

import jax.numpy as jnp

# 定义结构化数据类型：键对应字段名，值对应数据类型
struct_dtype = [('feature1', jnp.float32), ('feature2', jnp.int32), ('feature3', jnp.bool_)]

# 待存入的字典
data_dict = {'feature1': 3.14, 'feature2': 42, 'feature3': True}

# 将字典值按dtype顺序转为元组，再生成结构化数组
jax_struct_array = jnp.array(tuple(data_dict.values()), dtype=struct_dtype)

# 访问方式：通过字段名或索引
print(jax_struct_array['feature1'])  # 输出: 3.14

2. 将字典值堆叠为Jax数组的维度 #

如果字典中所有值都是形状一致的Jax数组，可以把这些值堆叠成新的维度，同时保留键到维度索引的映射关系。

示例代码：

import jax.numpy as jnp

# 字典值均为同形状的Jax数组
array_dict = {'x': jnp.array([1, 2, 3]), 'y': jnp.array([4, 5, 6]), 'z': jnp.array([7, 8, 9])}

# 将所有值沿新维度堆叠成一个Jax数组
jax_array = jnp.stack(list(array_dict.values()), axis=0)

# 建立键到索引的映射，方便后续通过键访问对应数据
key_index_map = {key: idx for idx, key in enumerate(array_dict.keys())}

# 示例：通过键获取对应数据
print(jax_array[key_index_map['y']])  # 输出: [4 5 6]

3. 用Jax树结构处理异构字典 #

如果需要保留字典的异构特性（比如值的类型、形状各不相同），可以使用Jax的jax.tree_util工具，它能把字典当作“Jax树”处理，支持JIT编译、自动微分等Jax核心功能，虽然不是传统意义上的数组，但能满足对异构数据的Jax操作需求。

示例代码：

import jax
import jax.numpy as jnp

# 异构字典：包含标量、不同形状的数组
hetero_dict = {'scalar': 2.5, 'vec': jnp.array([1, 2]), 'mat': jnp.array([[1,2],[3,4]])}

# 将字典转为Jax可处理的树结构（自动把非数组类型转为Jax数组）
jax_tree = jax.tree_util.tree_map(lambda x: jnp.asarray(x), hetero_dict)

# 对整个树结构进行JIT编译的操作
@jax.jit
def process_tree(tree):
    return {k: v * 2 for k, v in tree.items()}

result = process_tree(jax_tree)
print(result['mat'])  # 输出: [[2 4]
                      #        [6 8]]

内容的提问来源于stack exchange，提问作者imk