2、JAX张量操作：创建张量、张量属性、数据类型、索引与切片、广播机制、重塑与转置
好，咱们进入正题。张量操作，说白了就是JAX这门语言的「基本功」。你想想看，不管是图像处理还是生成模型，底层全是张量在跑来跑去。我刚开始接触JAX时，觉得它跟NumPy长得太像了，差点掉以轻心。结果呢？踩了几个坑才明白——相似归相似，细节上JAX有自己的脾气。
2.1 创建张量：从哪儿来？
创建张量，最直接的方式就是用jnp.array()。我个人习惯先写NumPy再转JAX，因为调试起来更顺手。
import jax.numpy as jnp
import numpy as np
# 从Python列表创建
a = jnp.array([1, 2, 3, 4])
print(a) # [1 2 3 4]
# 从NumPy数组创建
b = jnp.array(np.random.randn(3, 4))
print(b.shape) # (3, 4)
除了array()，还有几个常用函数：
jnp.zeros((2, 3))— 全零张量jnp.ones((2, 3))— 全一张量jnp.eye(4)— 单位矩阵jnp.arange(10)— 等差数列jnp.linspace(0, 1, 5)— 等间隔序列
小提示：JAX默认创建的是32位浮点数，跟NumPy的64位不一样。我在做图像生成时，经常用jnp.float32来省显存。
2.2 张量属性：看看它长什么样
拿到一个张量，第一件事就是看它的形状、维度和数据类型。这三个属性我几乎每天都要用。
x = jnp.array([[1, 2, 3], [4, 5, 6]])
print(x.shape) # (2, 3)
print(x.ndim) # 2
print(x.dtype) # int32
print(x.size) # 6
嗯，这里要注意：shape返回的是一个元组，ndim就是元组的长度。我曾经在写模型时，把shape[0]当成了batch size，结果发现维度搞反了——调试了半小时才发现。
2.3 数据类型：别小看它
JAX支持的数据类型跟NumPy基本一致，但默认是32位。这一点跟PyTorch很像，跟NumPy不同。
| 类型 | 说明 | 常用场景 |
|---|---|---|
jnp.float32 | 32位浮点 | 模型参数、图像数据 |
jnp.float64 | 64位浮点 | 高精度计算（但JAX默认禁用） |
jnp.int32 | 32位整型 | 索引、标签 |
jnp.bool_ | 布尔型 | 掩码、条件判断 |
注意：JAX默认不支持float64！如果你需要高精度，得在启动时设置jax.config.update("jax_enable_x64", True)。我第一次用JAX做科学计算时，发现结果精度不对，查了半天才发现是这个原因。
类型转换也很简单：
x = jnp.array([1, 2, 3], dtype=jnp.float32)
y = x.astype(jnp.int32)
2.4 索引与切片：怎么取数据？
索引和切片，跟NumPy完全一样。说白了就是方括号里写坐标。
x = jnp.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
# 取单个元素
print(x[1, 2]) # 6
# 取一行
print(x[0, :]) # [1 2 3]
# 取一列
print(x[:, 1]) # [2 5 8]
# 花式索引
print(x[[0, 2], [0, 2]]) # [1 9]
我个人觉得，花式索引在图像处理中特别有用。比如你要从一批图片中随机选几张，用数组索引就行。
避坑指南：JAX的索引返回的是新数组，不是视图。这意味着你修改索引结果不会影响原数组。我曾经以为能像NumPy那样原地修改，结果发现JAX的数组是不可变的。
2.5 广播机制：自动扩展的魔法
广播机制，说白了就是让不同形状的张量能一起做运算。JAX的广播规则跟NumPy完全一致：从最后一个维度开始比较，要么相等，要么一个是1，要么缺失。
# 标量广播
a = jnp.array([1, 2, 3])
b = 2
print(a + b) # [3 4 5]
# 向量广播到矩阵
a = jnp.array([[1, 2, 3], [4, 5, 6]])
b = jnp.array([10, 20, 30])
print(a + b) # [[11 22 33], [14 25 36]]
你想想看，如果没有广播机制，你得手动把b复制成跟a一样的形状，多麻烦。我在做图像归一化时，经常用广播来减去均值、除以标准差。
注意：广播虽然方便，但容易出错。我曾经把形状为(3, 1)和(1, 3)的两个数组做加法，结果得到了(3, 3)的矩阵——完全不是我想要的结果。所以，做广播前最好用reshape()显式调整形状。
2.6 重塑与转置：改变形状的艺术
重塑和转置，是张量操作中最常用的两个功能。重塑改变形状但不改变数据顺序，转置交换维度。
x = jnp.array([[1, 2, 3], [4, 5, 6]])
# 重塑
y = x.reshape(3, 2)
print(y)
# [[1 2]
# [3 4]
# [5 6]]
# 转置
z = x.T
print(z)
# [[1 4]
# [2 5]
# [3 6]]
嗯，这里有个细节：reshape是按行优先（C风格）填充的。如果你想要按列优先，可以用order='F'参数。
还有一个常用操作是flatten()，把多维张量拉成一维：
x = jnp.array([[1, 2], [3, 4]])
y = x.flatten()
print(y) # [1 2 3 4]
小技巧：在卷积神经网络中，经常要把特征图从(batch, height, width, channels)转成(batch, channels, height, width)。用transpose()可以轻松搞定：x.transpose(0, 3, 1, 2)。
知识体系总览
下面这张图，是我自己总结的张量操作知识体系。你可以把它当成一张地图，随时回来查漏补缺。
好了，这一章的内容就这些。张量操作是JAX的基石，你练得越熟，后面写模型就越顺手。别急，慢慢来，每个操作都亲手敲一遍代码，比光看强十倍。
公众号:蓝海资料掘金营，微信deep3321