srush
One piece of information missing: Practicality
I just finished all Tensor Puzzles - they were certainly very fun and made me think in a different way. But one question which lingers in my mind: how practical are these puzzles? Many of these take simple clear linear(complexity) code and make it into a tricky and quadratic(complexity) code. Most notably, something like "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"]:
return (eye(i) - 1*(arange(i)[:,None] - arange(i) == 1)) @ a
I know Python is super slow, but probably still not as slow to up the asymptotic complexity by introducing almost-all-zeroes matrixes. Do I understand correctly that (most) puzzles are primarily intended just to get comfortable with tensor operations and not to use their solutions in real-life scenarios?
P.S.: shorter and simpler "compress"
def compress(g: TT["i", bool], v: TT["i"], i:int) -> TT["i"]:
return pad_to(v[g == True], None, i)
(for this to work, Task 13 (pad_to) needs to be moved before Task 12, and written in a way that it ignores the second parameter (current size) and deduces current size from the given array)
They actually are surprisingly practical! Most things on GPUs are speed bound by the memory of loading memory. If the "quadratic" factor is just the size of the matrix you are already working with, then it is way better to just do more compute in exchange for fewer loads.
I agree that the specific example of "diff" could be done better with other techniques particularly in a high sparsity environment, but it looks quite similar to other techniques that people work with in things like transformer attention masking.
