output(i, j) +=! A(i, kk) * B(kk, j)
}
"""

In this example, we define a function `matmul` which takes two input `A` and `B` of shapes `M x N` and `N x K` and returns a single `output`. The shape of `output` is automatically inferred by the TC language (discussed below).

Let's look at this line:

```python
output(i, j) +=! A(i, kk) * B(kk, j)

It says:

output(i, j) means output is 2D.
for each location output(i, j), we add (+=) A(i, kk) * B(kk, j).
i is well-defined as all locations in A dim=0, i.e. i in range(0, M)
j is well-defined as all locations in B dim=1, i.e. j in range(0, K)
kk is inferred as all locations from 0 to N

The shape of output is inferred from the maximum values i and j can take, which is M and K, so output is of size M x K.
The ! symbol initializes output with 0.0. It is equivalent to:
output(i, j) = 0
output(i, j) += A(i, kk) * B(kk, j)

Scalar inputs and range constraints: implement AvgPool2d
```python