\n",
" \n",
"\n",
""
],
"text/plain": [
""
]
},
"metadata": {},
"output_type": "display_data"
}
],
"source": [
"def where(q, a, b):\n",
" \"Use this statement to replace an if-statement.\"\n",
" return (q * a) + (~q) * b\n",
"\n",
"# In my diagrams, orange is positive/True, where is zero/False, and green is negative.\n",
"\n",
"examples = [(tensor([False]), tensor([10]), tensor([0])),\n",
" (tensor([False, True]), tensor([1, 1]), tensor([-10, 0])),\n",
" (tensor([False, True]), tensor([1]), tensor([-10, 0])),\n",
" (tensor([[False, True], [True, False]]), tensor([1]), tensor([-10, 0])),\n",
" (tensor([[False, True], [True, False]]), tensor([[0], [10]]), tensor([-10, 0])),\n",
" ]\n",
" \n",
"\n",
"draw_examples(\"where\", [{\"q\": q, \"a\":a, \"b\":b, \"ret\": where(q, a, b)} for q, a, b in examples])\n"
]
},
{
"cell_type": "markdown",
"id": "ce1a49b3",
"metadata": {
"id": "ce1a49b3"
},
"source": [
"### Anti-Rules\n",
"\n",
"1. Nothing else. No `view`, `sum`, `take`, `squeeze`, `tensor`.\n",
"2. No cheating. Stackoverflow is great, but this is about first-principles.\n",
"3. Hint... these puzzles are mostly about Broadcasting. Make sure you understand this rule."
]
},
{
"cell_type": "markdown",
"id": "8648caeb",
"metadata": {
"id": "8648caeb"
},
"source": [
""
]
},
{
"cell_type": "code",
"execution_count": 17,
"id": "3a269693",
"metadata": {
"id": "3a269693",
"outputId": "48f08f78-6e92-4e0b-bcc6-27abc859029b"
},
"outputs": [
{
"data": {
"text/html": [
"\n",
"\n",
"\n",
"\n",
"\n",
"
\n",
"
\n",
" \n",
"\n",
""
],
"text/plain": [
""
]
},
"metadata": {},
"output_type": "display_data"
}
],
"source": [
"examples = [(arange(4), arange(5)[:, None]) ,\n",
" (arange(3)[:, None], arange(2))]\n",
"draw_examples(\"broadcast\", [{\"a\": a, \"b\":b, \"ret\": a + b} for a, b in examples])\n"
]
},
{
"cell_type": "markdown",
"id": "2fe8507e",
"metadata": {},
"source": [
"# Dims: a primer\n",
"\n",
"Dims are python objects that represent a dimension in a tensor, or a in a set of tensors that share the same dimension. Dimensions are created with the `dims` function and bound to variables:\n"
]
},
{
"cell_type": "code",
"execution_count": 18,
"id": "905f17b1",
"metadata": {},
"outputs": [],
"source": [
"from dim import dims\n",
"from torch import where"
]
},
{
"cell_type": "code",
"execution_count": 19,
"id": "3b6f9072",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"i j k\n"
]
}
],
"source": [
"i, j = dims()\n",
"k = dims()\n",
"print(i, j, k)"
]
},
{
"cell_type": "markdown",
"id": "d8130b77",
"metadata": {},
"source": [
"`dims` always returns the number of dimensions that are assigned to it. This is done using some C Python trickery, but is is the only \"magical\" thing about the API. It allows the dimensions to be created cheaply, and to provide the right names to the dim objects without redundancy.\n",
"\n",
"Dims can be bound to particular size on creation, by setting their `.size` property later, or _binding_ the dimension to the dimension of an existing Tensor:"
]
},
{
"cell_type": "code",
"execution_count": 20,
"id": "2d49de38",
"metadata": {
"scrolled": true
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"tensor([9, 1, 7, 6])\n",
"with dims=(j,) torch.Size([4])\n",
"0\n"
]
}
],
"source": [
"i.size = 3\n",
"x, y = dims(5, 6)\n",
"A = torch.randint(10, (4,))\n",
"a = A[j]\n",
"\n",
"print(a)\n",
"print(a.ndim)"
]
},
{
"cell_type": "markdown",
"id": "e2c28b29",
"metadata": {},
"source": [
"Tensors like `a` which has a bound dimension `j` are implicitly batched over that dimension. Similiar to xmap, operations on tensors with bound dimensions will result in a tensor with the union of all of their dimensions so operations can be batched across multiple bound dimensions:"
]
},
{
"cell_type": "code",
"execution_count": 21,
"id": "e0ef25cd",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"tensor([300, 200, 700])\n",
"with dims=(i,) torch.Size([3])\n",
"tensor([[309, 209, 709],\n",
" [301, 201, 701],\n",
" [307, 207, 707],\n",
" [306, 206, 706]])\n",
"with dims=(j, i) torch.Size([4, 3])\n"
]
}
],
"source": [
"B = torch.randint(10, (3,)) * 100\n",
"b = B[i]\n",
"print(b)\n",
"print(a + b)"
]
},
{
"cell_type": "markdown",
"id": "0a059422",
"metadata": {},
"source": [
"Operations that normally take numbers to specify a dimension can how take a dimension object to act on that dimension:"
]
},
{
"cell_type": "code",
"execution_count": 22,
"id": "5d1cdb89",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"tensor(23)\n"
]
}
],
"source": [
"print(a.sum(j))"
]
},
{
"cell_type": "markdown",
"id": "40aa46db",
"metadata": {},
"source": [
"To 'unbind' a dimension, returning it back to a positional dimension, call the `positional` function, which will prepend that dimension to the positional dimensions of a tensor:"
]
},
{
"cell_type": "code",
"execution_count": 23,
"id": "6d9885e9",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"tensor([9, 1, 7, 6])\n"
]
}
],
"source": [
"print(a.positional(j))"
]
},
{
"cell_type": "markdown",
"id": "92254761",
"metadata": {},
"source": [
"Tensors can have a mix of bound and unbound dimensions. When pretty printed, the `dims=` listing will indicate which physical dimensions of the tensor are bound to dimensions. Negative numbers indicate the positional index of unbound dimensions indexed from the right. It is printed this way because broadcast rules first right align the tensors, so the broadcasting of a tensor with mixed bound/unbound dimensions will be the union of the bound dimensions and negative integers in all the inputs."
]
},
{
"cell_type": "code",
"execution_count": 24,
"id": "f5bc4fd9",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"tensor([[0.7208, 0.7437, 0.5185, 0.4025],\n",
" [0.4226, 0.1594, 0.6546, 0.4325],\n",
" [0.6184, 0.2549, 0.7148, 0.5443]])\n",
"with dims=(i, -1) torch.Size([3, 4])\n",
"tensor([[0.8196, 0.5994, 0.3172, 0.8566, 0.4267],\n",
" [0.4182, 0.0527, 0.9666, 0.4099, 0.6620],\n",
" [0.8302, 0.0130, 0.7828, 0.0457, 0.2299],\n",
" [0.3119, 0.1769, 0.6141, 0.5387, 0.6078]])\n",
"with dims=(-1, k) torch.Size([4, 5])\n",
"tensor([[[1.5404, 1.3202, 1.0380, 1.5774, 1.1475],\n",
" [1.1619, 0.7964, 1.7103, 1.1536, 1.4057],\n",
" [1.3487, 0.5315, 1.3013, 0.5642, 0.7484],\n",
" [0.7144, 0.5794, 1.0166, 0.9412, 1.0102]],\n",
"\n",
" [[1.2422, 1.0220, 0.7398, 1.2792, 0.8493],\n",
" [0.5776, 0.2120, 1.1260, 0.5693, 0.8214],\n",
" [1.4848, 0.6676, 1.4374, 0.7003, 0.8845],\n",
" [0.7444, 0.6094, 1.0466, 0.9712, 1.0402]],\n",
"\n",
" [[1.4379, 1.2178, 0.9356, 1.4750, 1.0451],\n",
" [0.6732, 0.3076, 1.2215, 0.6649, 0.9169],\n",
" [1.5450, 0.7278, 1.4977, 0.7606, 0.9447],\n",
" [0.8562, 0.7212, 1.1584, 1.0830, 1.1520]]])\n",
"with dims=(i, -1, k) torch.Size([3, 4, 5])\n"
]
}
],
"source": [
"C = torch.rand(3, 4)\n",
"D = torch.rand(4, 5)\n",
"print(C[i])\n",
"print(D[:, k])\n",
"print(C[i] + D[:, k])"
]
},
{
"cell_type": "markdown",
"id": "4abfa9d7",
"metadata": {},
"source": [
"The `A[i]` indexing which has been used to bind dimensions is just the advanced indexing operator for Torch's tensors extended to work with dimension objects. When a dimension object is directly used as an index, it binds the object to that dimension of the tensor, and sets the size of the dimension object to the Tensor's size. However, it is also possible to index a tensor with other tensors that have bound dimensions. This behaves just like any other operator that has bound dimensions. The output tensor will have a union of all the bound dimensions of the inputs. A very useful special case of this behavior is when the index tensors have no unbound dimensions. In this case, the indexing operation behaves the same as a similarly written loop nest but without doing the loop in Python:"
]
},
{
"cell_type": "code",
"execution_count": 25,
"id": "c3a0e16b",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"tensor([0.9071, 0.6241])\n",
"with dims=(i,) torch.Size([2])\n",
"[tensor(0.9071), tensor(0.6241)]\n"
]
}
],
"source": [
"A = torch.rand(5)\n",
"indices = torch.tensor([1, 3])\n",
"i = dims()\n",
"\n",
"print(A[indices[i]])\n",
"\n",
"print([A[indices[x]] for x in range(2)])"
]
},
{
"cell_type": "markdown",
"id": "130d5bef",
"metadata": {},
"source": [
"Dimension objects can be used as tensor values as well where they are the index along a dimension (concretely, they have the value `torch.arange(i.size)[i]`). This can be useful for computing masks:"
]
},
{
"cell_type": "code",
"execution_count": 60,
"id": "58bd7469",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"tensor([[ True, False, False],\n",
" [False, True, False],\n",
" [False, False, True]])\n",
"with dims=(i, j) torch.Size([3, 3])\n",
"tensor([[ 0, 1, 2],\n",
" [100, 101, 102],\n",
" [200, 201, 202]])\n",
"with dims=(i, j) torch.Size([3, 3])\n"
]
}
],
"source": [
"i, j = dims(3, 3)\n",
"print(i.eq(j))\n",
"print(100*i + j)"
]
},
{
"cell_type": "markdown",
"id": "d17a742e",
"metadata": {},
"source": [
"Dims have more functionality: they can be used to split, merge, and reorder dimesions and `einsums` are automatically discovered by `sum` with `*`, but this information is sufficient to understand the answers to the puzzels below."
]
},
{
"cell_type": "markdown",
"id": "0e9cb801",
"metadata": {},
"source": [
"# Dim: New Rules\n",
"* Each solution can have 1 line to declare the dimensions, and 1 line to compute the result.\n",
"* pointwise arithmetic including where and numeric constants are allowed.\n",
"* You can use torch.tensor to lift a constant into a zero dim tensor, but not for any other use.\n",
"* @ (matrix multiply) is not allowed, but sum is allowed.\n",
"* arange is not allowed."
]
},
{
"cell_type": "markdown",
"id": "ff35419a",
"metadata": {
"id": "ff35419a"
},
"source": [
"## Running puzzles\n",
"\n",
"Each example, corresponds to a unit test which will randomly\n",
"try to break your code based on the spec. The spec is written in\n",
"standard python with lists.\n",
"\n",
"[](https://colab.research.google.com/github/srush/Tensor-Puzzles/blob/main/Tensor%20Puzzlers.ipynb)\n",
"\n",
"After you convince yourself it is right, uncomment the test for each example. If the test succeeds, you will get a puppy. "
]
},
{
"cell_type": "markdown",
"id": "dd21d735",
"metadata": {
"id": "dd21d735"
},
"source": [
"## Puzzle 1 - ones\n",
"\n",
"Compute [ones](https://numpy.org/doc/stable/reference/generated/numpy.ones.html) - the vector of all ones."
]
},
{
"cell_type": "code",
"execution_count": 27,
"id": "8cac10f4",
"metadata": {
"id": "8cac10f4",
"outputId": "d5c809bc-92ed-4435-daec-4ced8228924d"
},
"outputs": [
{
"data": {
"text/html": [
"\n",
"\n",
"\n",
"\n",
"\n",
"
\n",
"
\n",
" \n",
"\n",
""
],
"text/plain": [
""
]
},
"metadata": {},
"output_type": "display_data"
}
],
"source": [
"def ones_spec(out):\n",
" for i in range(len(out)):\n",
" out[i] = 1\n",
" \n",
"def ones(i: int) -> TT[\"i\"]:\n",
" i = dims(int(i))\n",
" return torch.tensor(1).expand(i).positional(i)\n",
"\n",
" # the `i` passed by the test framework is actually a tensor\n",
" # so it has to be cast back to an int, normally you could just write dims(x) to create a dim\n",
" # of size x.\n",
"\n",
" # One would rarely create\n",
" # specifically sized tensors of all ones, instead\n",
" # preferring to just use the constant 1 in arithmetic work with the constant 1\n",
"\n",
"\n",
"test_ones = make_test(\"one\", ones, ones_spec, add_sizes=[\"i\"])"
]
},
{
"cell_type": "code",
"execution_count": 28,
"id": "4a592143",
"metadata": {
"id": "4a592143",
"lines_to_next_cell": 2
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Correct!\n"
]
},
{
"data": {
"text/html": [
"\n",
" \n",
" "
],
"text/plain": [
""
]
},
"execution_count": 28,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"run_test(test_ones)"
]
},
{
"cell_type": "markdown",
"id": "df28c975",
"metadata": {
"id": "df28c975"
},
"source": [
"## Puzzle 2 - sum\n",
"\n",
"Compute [sum](https://numpy.org/doc/stable/reference/generated/numpy.sum.html) - the sum of a vector."
]
},
{
"cell_type": "code",
"execution_count": 29,
"id": "cbfa0c91",
"metadata": {
"id": "cbfa0c91",
"outputId": "8ab0c016-f6b6-4f34-8999-ad2b0db934a7"
},
"outputs": [
{
"data": {
"text/html": [
"\n",
"\n",
"\n",
"\n",
"\n",
"
\n",
"
\n",
" \n",
"\n",
""
],
"text/plain": [
""
]
},
"metadata": {},
"output_type": "display_data"
}
],
"source": [
"def sum_spec(a, out):\n",
" out[0] = 0\n",
" for i in range(len(a)):\n",
" out[0] += a[i]\n",
" \n",
"def sum(a: TT[\"i\"]) -> TT[1]:\n",
" # sum makes more sense as a primitive than matrix multiply\n",
" # so this puzzle no longer makes sense in this setup\n",
" return a.sum()\n",
"\n",
"\n",
"test_sum = make_test(\"sum\", sum, sum_spec)"
]
},
{
"cell_type": "code",
"execution_count": 30,
"id": "e4616f13",
"metadata": {
"id": "e4616f13",
"lines_to_next_cell": 2
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Correct!\n"
]
},
{
"data": {
"text/html": [
"\n",
" \n",
" "
],
"text/plain": [
""
]
},
"execution_count": 30,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"run_test(test_sum)"
]
},
{
"cell_type": "markdown",
"id": "8f5f4e49",
"metadata": {
"id": "8f5f4e49"
},
"source": [
"## Puzzle 3 - outer\n",
"\n",
"Compute [outer](https://numpy.org/doc/stable/reference/generated/numpy.outer.html) - the outer product of two vectors."
]
},
{
"cell_type": "code",
"execution_count": 31,
"id": "a162f170",
"metadata": {
"id": "a162f170",
"outputId": "a7591571-9773-4243-cf21-cde858eed269"
},
"outputs": [
{
"data": {
"text/html": [
"\n",
"\n",
"\n",
"\n",
"\n",
"