本文介绍如何通过张量重塑（reshape）与隐式广播（broadcasting）替代显式 tf.repeat，在保持计算正确性的同时显著降低内存开销，适用于大规模分块矩阵乘场景。
本文介绍如何通过张量重塑（reshape）与隐式广播（broadcasting）替代显式 `tf.repeat`，在保持计算正确性的同时显著降低内存开销，适用于大规模分块矩阵乘场景。
在深度学习与科学计算中，常遇到一类结构化矩阵乘需求：张量 a 形状为 (n//f, c, c)，需作用于 b（形状 (n, c, c)），其中 a 的每个 slice 实际对应 b 中连续 f 行的共享变换。朴素做法是用 tf.repeat(a, f, axis=0) 构造 (n, c, c) 的 a_prime 再执行 a_prime @ b，但该操作会复制 f 倍内存，当 n 和 f 较大时极易触发 OOM。
更优解法是规避显式复制，转而利用 TensorFlow 的广播机制与批量矩阵乘（tf.matmul）的维度对齐能力。核心思想是将 a 与 b 分别升维并错位对齐，使广播自动完成“每块 a[i] 作用于对应 f 行 b”的逻辑，再通过 tf.reduce_sum 聚合结果。
具体步骤如下：
-
重塑张量以启用广播：
- 将 a 从 (n//f, c, c) 变形为 (1, n//f, c, c)，引入 batch 维度 1；
- 将 b 从 (n, c, c) 变形为 (n, 1, c, c)，引入 broadcast 维度 1；
此时 a_reshaped 与 b_reshaped 在前两维上形成 (1, n//f) × (n, 1) 广播关系，等效于对每个 i ∈ [0, n//f) 和 j ∈ [i*f, (i+1)*f) 自动匹配 a[i] 与 b[j]。
执行批处理矩阵乘：
tf.matmul(a_reshaped, b_reshaped) 输出形状为 (n, n//f, c, c)，其中第 (j, i) 项即 a[i] @ b[j]（注意 j 与 i 的映射需满足 j ∈ [i*f, (i+1)*f)）。按组聚合结果：
由于 a[i] 应统一作用于 b 的第 i 组 f 行（即索引 i*f 到 (i+1)*f−1），我们需沿 axis=1（即 n//f 维）求和——但直接 reduce_sum(axis=1) 会错误地跨组累加。关键修正：实际应先将 b 按 f 分组 reshape，再与 a 对齐。更鲁棒的实现如下：
import tensorflow as tf
import numpy as np
n, c, f = 100, 5, 10
a = tf.constant(np.random.rand(n // f, c, c)) # shape: (10, 5, 5)
b = tf.constant(np.random.rand(n, c, c)) # shape: (100, 5, 5)
# Step 1: Reshape b into (n//f, f, c, c) — group rows by f
b_grouped = tf.reshape(b, (n // f, f, c, c)) # shape: (10, 10, 5, 5)
# Step 2: Expand a to (n//f, 1, c, c) for broadcasting over f-dim
a_expanded = tf.expand_dims(a, axis=1) # shape: (10, 1, 5, 5)
# Step 3: Broadcast-matmul → (n//f, f, c, c)
intermediate = tf.matmul(a_expanded, b_grouped) # a[i] @ b_grouped[i,j] for each j
# Step 4: Flatten first two dims → (n, c, c)
result = tf.reshape(intermediate, (n, c, c))
print("Result shape:", result.shape) # (100, 5, 5)
✅ 优势总结：
- 零内存冗余：全程无 repeat 或 tile，仅通过 reshape 和 expand_dims 改变视图；
- GPU友好：tf.matmul 充分利用硬件加速，远快于 Python 循环或逐行 tf.linalg.matvec；
- 可扩展性强：支持任意 f（无需整除校验），且易于集成到 tf.function 图中。
⚠️ 注意事项：
- 确保 n % f == 0，否则 reshape 会报错；若不满足，需先截断或补零；
- 该方法本质是分组广播矩阵乘，语义严格等价于 for i in range(n//f): result[i*f:(i+1)*f] = a[i] @ b[i*f:(i+1)*f]；
- 若 b 实际为 (n, c, 1)（如问题描述末尾所述），则 @ 应为 tf.linalg.matvec 或调整 b 的 c 维为 1，此时 result 形状为 (n, c, 1)，逻辑不变。
