如何获取jaxlib.xla_extension.ArrayImpl对象中的数值？

从报错IndexError: Too many indices for array: 1 non-None/Ellipsis indices for dim 0可知，z1[0]是0维标量数组，而非1维数组，因此不能通过[0]二次索引。

提取标量JAX数组的Python数值可采用以下方法：

• 直接用类型转换：float(z1[0])（整数类型用int()）
• 调用.item()方法：z1[0].item()，自动匹配对应Python数值类型
• 转为NumPy数组后取值：np.array(z1[0]).item()（需先导入numpy）

示例代码：

import jax.numpy as jnp

z1 = jnp.array([0.71530414])
scalar_arr = z1[0]
print(float(scalar_arr))  # 输出 0.71530414
print(scalar_arr.item())  # 输出 0.71530414

若需将整个JAX数组转为Python列表：

• 0维标量数组：[scalar_arr.item()]
• 多维数组：调用.tolist()方法，示例：

arr = jnp.array([[1.0, 2.0], [3.0, 4.0]])
print(arr.tolist())  # 输出 [[1.0, 2.0], [3.0, 4.0]]

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