- by Sasha Rush - srush_nlp (with Marcos Treviso)
When learning a tensor programming language like PyTorch or Numpy it is tempting to rely on the standard library (or more honestly StackOverflow) to find a magic function for everything. But in practice, the tensor language is extremely expressive, and you can do most things from first principles and clever use of broadcasting.
This is a collection of 21 tensor puzzles. Like chess puzzles these are not meant to simulate the complexity of a real program, but to practice in a simplified environment. Each puzzle asks you to reimplement one function in the NumPy standard library without magic.
I recommend running in Colab. Click here and copy the notebook to get start.
If you are interested, there is also a youtube walkthrough of the puzzles
!pipinstall-qqqtorchtypinghypothesispytestgit+https://github.com/danoneata/chalk@srush-patch-1
!wget-qhttps://github.com/srush/Tensor-Puzzles/raw/main/lib.pyfromlibimportdraw_examples, make_test, run_testimporttorchimportnumpyasnpfromtorchtypingimportTensorTypeasTTtensor=torch.tensor- These puzzles are about broadcasting. Know this rule.
Each puzzle needs to be solved in 1 line (<80 columns) of code.
You are allowed @, arithmetic, comparison,
shape, any indexing (e.g.a[:j], a[:, None], a[arange(10)]), and previous puzzle functions.You are not allowed anything else. No
view,sum,take,squeeze,tensor.You can start with these two functions:
defarange(i: int):
"Use this function to replace a for-loop."returntorch.tensor(range(i))
draw_examples("arange", [{"" : arange(i)} foriin [5, 3, 9]])
# Example of broadcasting.examples= [(arange(4), arange(5)[:, None]) ,
(arange(3)[:, None], arange(2))]
draw_examples("broadcast", [{"a": a, "b":b, "ret": a+b} fora, binexamples])
defwhere(q, a, b):
"Use this function to replace an if-statement."return (q*a) + (~q) *b# In diagrams, orange is positive/True, where is zero/False, and blue is negative.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)} forq, a, binexamples])
Compute ones - the vector of all ones.
defones_spec(out):
foriinrange(len(out)):
out[i] =1defones(i: int) ->TT["i"]:
raiseNotImplementedErrortest_ones=make_test("one", ones, ones_spec, add_sizes=["i"])
# run_test(test_ones)Compute sum - the sum of a vector.
defsum_spec(a, out):
out[0] =0foriinrange(len(a)):
out[0] +=a[i]
defsum(a: TT["i"]) ->TT[1]:
raiseNotImplementedErrortest_sum=make_test("sum", sum, sum_spec)
# run_test(test_sum)Compute outer - the outer product of two vectors.
defouter_spec(a, b, out):
foriinrange(len(out)):
forjinrange(len(out[0])):
out[i][j] =a[i] *b[j]
defouter(a: TT["i"], b: TT["j"]) ->TT["i", "j"]:
raiseNotImplementedErrortest_outer=make_test("outer", outer, outer_spec)
# run_test(test_outer)Compute diag - the diagonal vector of a square matrix.
defdiag_spec(a, out):
foriinrange(len(a)):
out[i] =a[i][i]
defdiag(a: TT["i", "i"]) ->TT["i"]:
raiseNotImplementedErrortest_diag=make_test("diag", diag, diag_spec)
# run_test(test_diag)Compute eye - the identity matrix.
defeye_spec(out):
foriinrange(len(out)):
out[i][i] =1defeye(j: int) ->TT["j", "j"]:
raiseNotImplementedErrortest_eye=make_test("eye", eye, eye_spec, add_sizes=["j"])
# run_test(test_eye)Compute triu - the upper triangular matrix.
deftriu_spec(out):
foriinrange(len(out)):
forjinrange(len(out)):
ifi<=j:
out[i][j] =1else:
out[i][j] =0deftriu(j: int) ->TT["j", "j"]:
raiseNotImplementedErrortest_triu=make_test("triu", triu, triu_spec, add_sizes=["j"])
# run_test(test_triu)Compute cumsum - the cumulative sum.
defcumsum_spec(a, out):
total=0foriinrange(len(out)):
out[i] =total+a[i]
total+=a[i]
defcumsum(a: TT["i"]) ->TT["i"]:
raiseNotImplementedErrortest_cumsum=make_test("cumsum", cumsum, cumsum_spec)
# run_test(test_cumsum)Compute diff - the running difference.
defdiff_spec(a, out):
out[0] =a[0]
foriinrange(1, len(out)):
out[i] =a[i] -a[i-1]
defdiff(a: TT["i"], i: int) ->TT["i"]:
raiseNotImplementedErrortest_diff=make_test("diff", diff, diff_spec, add_sizes=["i"])
# run_test(test_diff)Compute vstack - the matrix of two vectors
defvstack_spec(a, b, out):
foriinrange(len(out[0])):
out[0][i] =a[i]
out[1][i] =b[i]
defvstack(a: TT["i"], b: TT["i"]) ->TT[2, "i"]:
raiseNotImplementedErrortest_vstack=make_test("vstack", vstack, vstack_spec)
# run_test(test_vstack)Compute roll - the vector shifted 1 circular position.
defroll_spec(a, out):
foriinrange(len(out)):
ifi+1<len(out):
out[i] =a[i+1]
else:
out[i] =a[i+1-len(out)]
defroll(a: TT["i"], i: int) ->TT["i"]:
raiseNotImplementedErrortest_roll=make_test("roll", roll, roll_spec, add_sizes=["i"])
# run_test(test_roll)Compute flip - the reversed vector
defflip_spec(a, out):
foriinrange(len(out)):
out[i] =a[len(out) -i-1]
defflip(a: TT["i"], i: int) ->TT["i"]:
raiseNotImplementedErrortest_flip=make_test("flip", flip, flip_spec, add_sizes=["i"])
# run_test(test_flip)Compute compress - keep only masked entries (left-aligned).
defcompress_spec(g, v, out):
j=0foriinrange(len(g)):
ifg[i]:
out[j] =v[i]
j+=1defcompress(g: TT["i", bool], v: TT["i"], i:int) ->TT["i"]:
raiseNotImplementedErrortest_compress=make_test("compress", compress, compress_spec, add_sizes=["i"])
# run_test(test_compress)Compute pad_to - eliminate or add 0s to change size of vector.
defpad_to_spec(a, out):
foriinrange(min(len(out), len(a))):
out[i] =a[i]
defpad_to(a: TT["i"], i: int, j: int) ->TT["j"]:
raiseNotImplementedErrortest_pad_to=make_test("pad_to", pad_to, pad_to_spec, add_sizes=["i", "j"])
# run_test(test_pad_to)Compute sequence_mask - pad out to length per batch.
defsequence_mask_spec(values, length, out):
foriinrange(len(out)):
forjinrange(len(out[0])):
ifj<length[i]:
out[i][j] =values[i][j]
else:
out[i][j] =0defsequence_mask(values: TT["i", "j"], length: TT["i", int]) ->TT["i", "j"]:
raiseNotImplementedErrordefconstraint_set_length(d):
d["length"] =d["length"] %d["values"].shape[1]
returndtest_sequence=make_test("sequence_mask",
sequence_mask, sequence_mask_spec, constraint=constraint_set_length
)
# run_test(test_sequence)Compute bincount - count number of times an entry was seen.
defbincount_spec(a, out):
foriinrange(len(a)):
out[a[i]] +=1defbincount(a: TT["i"], j: int) ->TT["j"]:
raiseNotImplementedErrordefconstraint_set_max(d):
d["a"] =d["a"] %d["return"].shape[0]
returndtest_bincount=make_test("bincount",
bincount, bincount_spec, add_sizes=["j"], constraint=constraint_set_max
)
# run_test(test_bincount)Compute scatter_add - add together values that link to the same location.
defscatter_add_spec(values, link, out):
forjinrange(len(values)):
out[link[j]] +=values[j]
defscatter_add(values: TT["i"], link: TT["i"], j: int) ->TT["j"]:
raiseNotImplementedErrordefconstraint_set_max(d):
d["link"] =d["link"] %d["return"].shape[0]
returndtest_scatter_add=make_test("scatter_add",
scatter_add, scatter_add_spec, add_sizes=["j"], constraint=constraint_set_max
)
# run_test(test_scatter_add)Compute flatten
defflatten_spec(a, out):
k=0foriinrange(len(a)):
forjinrange(len(a[0])):
out[k] =a[i][j]
k+=1defflatten(a: TT["i", "j"], i:int, j:int) ->TT["i * j"]:
raiseNotImplementedErrortest_flatten=make_test("flatten", flatten, flatten_spec, add_sizes=["i", "j"])
# run_test(test_flatten)Compute linspace
deflinspace_spec(i, j, out):
forkinrange(len(out)):
out[k] =float(i+ (j-i) *k/max(1, len(out) -1))
deflinspace(i: TT[1], j: TT[1], n: int) ->TT["n", float]:
raiseNotImplementedErrortest_linspace=make_test("linspace", linspace, linspace_spec, add_sizes=["n"])
# run_test(test_linspace)Compute heaviside
defheaviside_spec(a, b, out):
forkinrange(len(out)):
ifa[k] ==0:
out[k] =b[k]
else:
out[k] =int(a[k] >0)
defheaviside(a: TT["i"], b: TT["i"]) ->TT["i"]:
raiseNotImplementedErrortest_heaviside=make_test("heaviside", heaviside, heaviside_spec)
# run_test(test_heaviside)Compute repeat
defrepeat_spec(a, d, out):
foriinrange(d[0]):
forkinrange(len(a)):
out[i][k] =a[k]
defconstraint_set(d):
d["d"][0] =d["return"].shape[0]
returnddefrepeat(a: TT["i"], d: TT[1]) ->TT["d", "i"]:
raiseNotImplementedErrortest_repeat=make_test("repeat", repeat, repeat_spec, constraint=constraint_set)
Compute bucketize
defbucketize_spec(v, boundaries, out):
fori, valinenumerate(v):
out[i] =0forjinrange(len(boundaries)-1):
ifval>=boundaries[j]:
out[i] =j+1ifval>=boundaries[-1]:
out[i] =len(boundaries)
defconstraint_set(d):
d["boundaries"] =np.abs(d["boundaries"]).cumsum()
returnddefbucketize(v: TT["i"], boundaries: TT["j"]) ->TT["i"]:
raiseNotImplementedErrortest_bucketize=make_test("bucketize", bucketize, bucketize_spec,
constraint=constraint_set)
What is the smallest you can make each of these?
importinspectfns= (ones, sum, outer, diag, eye, triu, cumsum, diff, vstack, roll, flip,
compress, pad_to, sequence_mask, bincount, scatter_add)
forfninfns:
lines= [lforlininspect.getsource(fn).split("\n") ifnotl.strip().startswith("#")]
iflen(lines) >3:
print(fn.__name__, len(lines[2]), "(more than 1 line)")
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