张量谜题

• 作者：Sasha Rush - srush_nlp（与 Marcos Treviso 合作）

在学习像 PyTorch 或 NumPy 这样的张量编程语言时，人们很容易依赖标准库（或者更坦率地说，StackOverflow）来找到解决所有问题的“魔法函数”。但实际上，张量语言具有极强的表达能力，你可以通过基本原理和巧妙地利用广播机制来完成大多数任务。

这是一组包含 21 个张量谜题的集合。就像国际象棋谜题一样，这些谜题并不是为了模拟真实程序的复杂性，而是为了在一个简化的环境中进行练习。每个谜题都要求你不用任何“魔法”函数，重新实现 NumPy 标准库中的一个函数。

我建议在 Colab 中运行。点击这里并复制笔记本即可开始。

如果你感兴趣，还有一个关于这些谜题的 YouTube 演示视频：

!pip install -qqq torchtyping hypothesis pytest git+https://github.com/danoneata/chalk@srush-patch-1
!wget -q https://github.com/srush/Tensor-Puzzles/raw/main/lib.py

from lib import draw_examples, make_test, run_test
import torch
import numpy as np
from torchtyping import TensorType as TT
tensor = torch.tensor

规则

1. 这些谜题围绕 广播 展开。请务必了解这一规则。

2. 每个谜题都需要用一行代码（不超过 80 列）来解答。
3. 你可以使用 @、算术运算、比较运算、shape、任何索引操作（例如 a[:j], a[:, None], a[arange(10)]），以及之前谜题中定义的函数。
4. 你 不允许 使用其他任何东西。不能使用 view、sum、take、squeeze、tensor 等函数。
5. 你可以从以下两个函数开始：

def arange(i: int):
    "使用此函数代替 for 循环。"
    return torch.tensor(range(i))

draw_examples("arange", [{"" : arange(i)} for i in [5, 3, 9]])

# 广播的示例。
examples = [(arange(4), arange(5)[:, None]) ,
            (arange(3)[:, None], arange(2))]
draw_examples("broadcast", [{"a": a, "b":b, "ret": a + b} for a, b in examples])

def where(q, a, b):
    "使用此函数代替 if 语句。"
    return (q * a) + (~q) * b

# 在图示中，橙色表示正数/True，零表示假/False，蓝色表示负数。

examples = [(tensor([False]), tensor([10]), tensor([0])),
            (tensor([False, True]), tensor([1, 1]), tensor([-10, 0])),
            (tensor([False, True]), tensor([1]), tensor([-10, 0])),
            (tensor([[False, True], [True, False]]), tensor([1]), tensor([-10, 0])),
            (tensor([[False, True], [True, False]]), tensor([[0], [10]]), tensor([-10, 0])),
           ]
draw_examples("where", [{"q": q, "a":a, "b":b, "ret": where(q, a, b)} for q, a, b in examples])

谜题 1 - ones

计算 ones——全为 1 的向量。

def ones_spec(out):
    for i in range(len(out)):
        out[i] = 1
        
def ones(i: int) -> TT["i"]:
    raise NotImplementedError

test_ones = make_test("one", ones, ones_spec, add_sizes=["i"])

# run_test(test_ones)

谜题 2 - sum

计算 sum——向量的总和。

def sum_spec(a, out):
    out[0] = 0
    for i in range(len(a)):
        out[0] += a[i]
        
def sum(a: TT["i"]) -> TT[1]:
    raise NotImplementedError


test_sum = make_test("sum", sum, sum_spec)

# run_test(test_sum)

谜题 3 - outer

计算 outer——两个向量的外积。

def outer_spec(a, b, out):
    for i in range(len(out)):
        for j in range(len(out[0])):
            out[i][j] = a[i] * b[j]
            
def outer(a: TT["i"], b: TT["j"]) -> TT["i", "j"]:
    raise NotImplementedError
    
test_outer = make_test("outer", outer, outer_spec)

# run_test(test_outer)

谜题 4 - diag

计算 diag——方阵的对角线向量。

def diag_spec(a, out):
    for i in range(len(a)):
        out[i] = a[i][i]
        
def diag(a: TT["i", "i"]) -> TT["i"]:
    raise NotImplementedError


test_diag = make_test("diag", diag, diag_spec)

# run_test(test_diag)

谜题 5 - eye

计算 eye——单位矩阵。

def eye_spec(out):
    for i in range(len(out)):
        out[i][i] = 1
        
def eye(j: int) -> TT["j", "j"]:
    raise NotImplementedError
    
test_eye = make_test("eye", eye, eye_spec, add_sizes=["j"])

# run_test(test_eye)

谜题 6 - triu

计算 triu——上三角矩阵。

def triu_spec(out):
    for i in range(len(out)):
        for j in range(len(out)):
            if i <= j:
                out[i][j] = 1
            else:
                out[i][j] = 0
                
def triu(j: int) -> TT["j", "j"]:
    raise NotImplementedError


test_triu = make_test("triu", triu, triu_spec, add_sizes=["j"])

# run_test(test_triu)

谜题 7 - cumsum

计算 cumsum——累积和。

def cumsum_spec(a, out):
    total = 0
    for i in range(len(out)):
        out[i] = total + a[i]
        total += a[i]

def cumsum(a: TT["i"]) -> TT["i"]:
    raise NotImplementedError

test_cumsum = make_test("cumsum", cumsum, cumsum_spec)

# run_test(test_cumsum)

谜题 8 - diff

计算 diff——滚动差分。

def diff_spec(a, out):
    out[0] = a[0]
    for i in range(1, len(out)):
        out[i] = a[i] - a[i - 1]

def diff(a: TT["i"], i: int) -> TT["i"]:
    raise NotImplementedError

test_diff = make_test("diff", diff, diff_spec, add_sizes=["i"])

# run_test(test_diff)

谜题 9 - vstack

计算 vstack——两个向量组成的矩阵。

def vstack_spec(a, b, out):
    for i in range(len(out[0])):
        out[0][i] = a[i]
        out[1][i] = b[i]

def vstack(a: TT["i"], b: TT["i"]) -> TT[2, "i"]:
    raise NotImplementedError


test_vstack = make_test("vstack", vstack, vstack_spec)

# run_test(test_vstack)

谜题 10 - roll

计算 roll——向量循环右移一位。

def roll_spec(a, out):
    for i in range(len(out)):
        if i + 1 < len(out):
            out[i] = a[i + 1]
        else:
            out[i] = a[i + 1 - len(out)]
            
def roll(a: TT["i"], i: int) -> TT["i"]:
    raise NotImplementedError


test_roll = make_test("roll", roll, roll_spec, add_sizes=["i"])

# run_test(test_roll)

谜题 11 - flip

计算 flip——反转向量。

def flip_spec(a, out):
    for i in range(len(out)):
        out[i] = a[len(out) - i - 1]
        
def flip(a: TT["i"], i: int) -> TT["i"]:
    raise NotImplementedError


test_flip = make_test("flip", flip, flip_spec, add_sizes=["i"])

# run_test(test_flip)

谜题 12 - compress

计算 compress——仅保留掩码标记的元素（左对齐）。

def compress_spec(g, v, out):
    j = 0
    for i in range(len(g)):
        if g[i]:
            out[j] = v[i]
            j += 1
            
def compress(g: TT["i", bool], v: TT["i"], i:int) -> TT["i"]:
    raise NotImplementedError


test_compress = make_test("compress", compress, compress_spec, add_sizes=["i"])

# run_test(test_compress)

谜题 13 - pad_to

计算 pad_to——通过删除或添加零来改变向量的大小。

def pad_to_spec(a, out):
    for i in range(min(len(out), len(a))):
        out[i] = a[i]


def pad_to(a: TT["i"], i: int, j: int) -> TT["j"]:
    raise NotImplementedError


test_pad_to = make_test("pad_to", pad_to, pad_to_spec, add_sizes=["i", "j"])

# run_test(test_pad_to)

谜题 14 - sequence_mask

计算 sequence_mask——按批次填充到指定长度。

def sequence_mask_spec(values, length, out):
    for i in range(len(out)):
        for j in range(len(out[0])):
            if j < length[i]:
                out[i][j] = values[i][j]
            else:
                out[i][j] = 0
    
def sequence_mask(values: TT["i", "j"], length: TT["i", int]) -> TT["i", "j"]:
    raise NotImplementedError


def constraint_set_length(d):
    d["length"] = d["length"] % d["values"].shape[1]
    return d


test_sequence = make_test("sequence_mask",
    sequence_mask, sequence_mask_spec, constraint=constraint_set_length
)

# run_test(test_sequence)

谜题 15 - bincount

计算 bincount——统计每个元素出现的次数。

def bincount_spec(a, out):
    for i in range(len(a)):
        out[a[i]] += 1
        
def bincount(a: TT["i"], j: int) -> TT["j"]:
    raise NotImplementedError


def constraint_set_max(d):
    d["a"] = d["a"] % d["return"].shape[0]
    return d


test_bincount = make_test("bincount",
    bincount, bincount_spec, add_sizes=["j"], constraint=constraint_set_max
)

# run_test(test_bincount)

谜题 16 - scatter_add

计算 scatter_add——将指向同一位置的值相加。

def scatter_add_spec(values, link, out):
    for j in range(len(values)):
        out[link[j]] += values[j]
        
def scatter_add(values: TT["i"], link: TT["i"], j: int) -> TT["j"]:
    raise NotImplementedError


def constraint_set_max(d):
    d["link"] = d["link"] % d["return"].shape[0]
    return d


test_scatter_add = make_test("scatter_add",
    scatter_add, scatter_add_spec, add_sizes=["j"], constraint=constraint_set_max
)

# run_test(test_scatter_add)

谜题 17 - flatten

计算 flatten

def flatten_spec(a, out):
    k = 0
    for i in range(len(a)):
        for j in range(len(a[0])):
            out[k] = a[i][j]
            k += 1

def flatten(a: TT["i", "j"], i:int, j:int) -> TT["i * j"]:
    raise NotImplementedError

test_flatten = make_test("flatten", flatten, flatten_spec, add_sizes=["i", "j"])

# run_test(test_flatten)

谜题 18 - linspace

计算 linspace

def linspace_spec(i, j, out):
    for k in range(len(out)):
        out[k] = float(i + (j - i) * k / max(1, len(out) - 1))

def linspace(i: TT[1], j: TT[1], n: int) -> TT["n", float]:
    raise NotImplementedError

test_linspace = make_test("linspace", linspace, linspace_spec, add_sizes=["n"])

# run_test(test_linspace)

谜题 19 - heaviside

计算 heaviside

def heaviside_spec(a, b, out):
    for k in range(len(out)):
        if a[k] == 0:
            out[k] = b[k]
        else:
            out[k] = int(a[k] > 0)

def heaviside(a: TT["i"], b: TT["i"]) -> TT["i"]:
    raise NotImplementedError

test_heaviside = make_test("heaviside", heaviside, heaviside_spec)

# run_test(test_heaviside)

谜题 20 - repeat (1d)

计算 repeat

def repeat_spec(a, d, out):
    for i in range(d[0]):
        for k in range(len(a)):
            out[i][k] = a[k]

def constraint_set(d):
    d["d"][0] = d["return"].shape[0]
    return d
            
def repeat(a: TT["i"], d: TT[1]) -> TT["d", "i"]:
    raise NotImplementedError

test_repeat = make_test("repeat", repeat, repeat_spec, constraint=constraint_set)

谜题 21 - bucketize

计算 bucketize

def bucketize_spec(v, boundaries, out):
    for i, val in enumerate(v):
        out[i] = 0
        for j in range(len(boundaries)-1):
            if val >= boundaries[j]:
                out[i] = j + 1
        if val >= boundaries[-1]:
            out[i] = len(boundaries)


def constraint_set(d):
    d["boundaries"] = np.abs(d["boundaries"]).cumsum()
    return d
            
def bucketize(v: TT["i"], boundaries: TT["j"]) -> TT["i"]:
    raise NotImplementedError

test_bucketize = make_test("bucketize", bucketize, bucketize_spec,
                           constraint=constraint_set)

快速通关模式!

你能把每一个函数都做得多小呢？

import inspect
fns = (ones, sum, outer, diag, eye, triu, cumsum, diff, vstack, roll, flip,
       compress, pad_to, sequence_mask, bincount, scatter_add)

for fn in fns:
    lines = [l for l in inspect.getsource(fn).split("\n") if not l.strip().startswith("#")]
    
    if len(lines) > 3:
        print(fn.__name__, len(lines[2]), "(超过1行)")
    else:
        print(fn.__name__, len(lines[1]))

ones 29
sum 29
outer 29
diag 29
eye 29
triu 29
cumsum 29
diff 29
vstack 29
roll 29
flip 29
compress 29
pad_to 29
sequence_mask 29
bincount 29
scatter_add 29

Tensor-Puzzles 快速上手指南

Tensor-Puzzles 是一套由 Sasha Rush 开发的交互式编程练习，旨在帮助开发者深入理解 PyTorch 和 NumPy 中的广播机制（Broadcasting）。通过 21 个精心设计的谜题，你将学会在不依赖高级库函数（如 sum, view, squeeze 等）的情况下，仅利用基础运算、索引和广播规则来重构常见的张量操作。

环境准备

• 运行平台：推荐使用 Google Colab（无需本地配置，开箱即用）。
• 本地环境（可选）：
  • Python 3.8+
  • PyTorch
  • NumPy
  • 额外依赖：torchtyping, hypothesis, pytest, chalk

安装步骤

方案一：使用 Google Colab（推荐）

点击以下按钮直接在浏览器中打开并复制笔记本，即可开始练习：

方案二：本地安装

如果你希望在本地运行，请先克隆项目或下载核心文件，然后执行以下命令安装依赖：

!pip install -qqq torchtyping hypothesis pytest git+https://github.com/danoneata/chalk@srush-patch-1
!wget -q https://github.com/srush/Tensor-Puzzles/raw/main/lib.py

(注：国内用户若下载缓慢，可尝试手动下载 lib.py 文件至当前目录，或使用国内 pip 镜像源安装其他依赖)

基本使用

1. 导入环境

在 Notebook 单元格中运行以下代码以加载工具库和基础函数：

from lib import draw_examples, make_test, run_test
import torch
import numpy as np
from torchtyping import TensorType as TT
tensor = torch.tensor

2. 核心规则

在开始解题前，请牢记以下限制，这是本项目的核心训练目标：

1. 核心机制：所有谜题均围绕**广播（Broadcasting）**展开。
2. 代码限制：每个谜题必须在 1 行代码（少于 80 列）内完成。
3. 允许操作：矩阵乘法 (@)、算术运算、比较运算、.shape、任意索引（如 a[:j], a[:, None], a[arange(10)]）以及已完成的谜题函数。
4. 禁止操作：严禁使用 view, sum, take, squeeze, tensor 等“魔法”函数。

3. 解题示例

每个谜题都提供了规格说明函数（_spec，模拟循环逻辑）和你需要实现的函数。

示例：Puzzle 1 - ones 目标：创建一个全为 1 的向量，不能使用 torch.ones。

# 题目定义
def ones_spec(out):
    for i in range(len(out)):
        out[i] = 1
        
def ones(i: int) -> TT["i"]:
    raise NotImplementedError  # 在此处填写你的一行代码

# 生成测试用例
test_ones = make_test("one", ones, ones_spec, add_sizes=["i"])

# 运行测试验证答案
run_test(test_ones)

解题思路提示： 利用提供的辅助函数 arange 和广播机制。例如，可以通过让一个全 1 的标量与任意形状的张量进行广播运算，或者利用布尔值转换来实现。

4. 进阶辅助函数

项目中预定义了两个关键辅助函数，用于替代循环和条件判断：

• arange: 替代 for 循环生成序列。

  def arange(i: int):
      return torch.tensor(range(i))
  

• where: 替代 if-else 逻辑。

  def where(q, a, b):
      return (q * a) + (~q) * b
  

5. 开始挑战

依次完成从 Puzzle 1 到 Puzzle 21 的代码填空，每完成一个取消注释 run_test(...) 即可验证正确性。如需可视化参考，可观看作者的 YouTube 讲解视频。