2️⃣ Batched Operations

Learning Objectives

Learn about some important concepts related to batched operations, e.g. broadcasting and logical reductions

Understand and use the einops library

Apply this knowledge to create & work with a batch of rays

In this section, we'll move onto batched operations. First, it's necessary to cover some important tips for working effectively with PyTorch tensors. If you've gone through the prerequisite material then several of these should already be familiar to you.

Tensor Operations - 5 Tips

Tip (1/5) - Elementwise Logical Operations on Tensors

For regular booleans, the keywords and, or, and not are used to do logical operations and the operators &, |, and ~ do and, or and not on each bit of the input numbers. Analogously, we use the operators &, | ~ on tensors to perform these operations on each element of the tensor, e.g. x & y returns the tensor with elements x[i] and y[i].

A few important gotchas here:

Don't try to use and, or or not on tensors, since Python will try to coerce the tensors to booleans, and you'll get an exception.
Remember about operator precedence! For instance, v >= 0 & v <= 1 will actually throw an error because it's actually evaluated as the chained comparison (v >= (0 & v)) and ((0 & v) <= 1) (because & binds more tightly than the comparisons) and 0 & v is not a valid operation. When in doubt, use parentheses to force the correct parsing: (v >= 0) & (v <= 1).

Tip (2/5) - `einops`

Einops is a useful library which we'll dive deeper with tomorrow. For now, the only important function you'll need to know is einops.repeat. This takes as arguments a tensor and a string, and returns a new tensor which has been repeated along the specified dimensions. For example, the following code shows how we can repeat a 2D tensor along the last dimension:

x = t.randn(4, 3)
x_repeated = einops.repeat(x, 'a b -> a b c', c=2) # copies x along a new dimension at the end

assert x_repeated.shape == (4, 3, 2)
t.testing.assert_close(x_repeated[:, :, 0], x)
t.testing.assert_close(x_repeated[:, :, 1], x)

Tip (3/5) - Logical Reductions

In plain Python, if you have a list of lists and want to know if any element in a row is True, you could use a list comprehension like [any(row) for row in rows]. The efficient way to do this in PyTorch is with torch.any() or equivalently the .any() method of a tensor, which accept the dimension to reduce over. Similarly, torch.all() or .all() method. Both of these methods accept a dim argument, which is the dimension to reduce over.

Aside - tensor methods

Most functions like torch.any(tensor, ...) (which take a tensor as first argument) have an equivalent tensor method tensor.any(...). We'll see many more examples of functions like this as we go through the course.

Tip (4/5) - Broadcasting

Broadcasting is what happens when you perform an operation on two tensors, and one is a smaller size, but is copied along the dimensions of the larger one in order to apply to it. There are some complicated broadcasting rules which we'll get into later in the course (you can also review the broadcasting section in the prerequisite material which comes before this chapter), but for our purposes the only thing you need to know is this: if you perform an operation on 2 tensors A, B where A.shape == B.shape[-A.ndim:] (i.e. A matches the last dimensions of B), then A will be copied along its leading dimensions until it has the same shape as B. For example:

B = t.ones(4, 3, 2)
A = t.ones(3, 2)
C = A + B
print(C.shape) # torch.Size([4, 3, 2]), the size of the broadcasted tensor
print(C.max(), C.min()) # tensor(2.) tensor(2.), because all elements are 2

Tip (5/5) - Indexing

Indexing is a pretty deep topic in PyTorch, and it takes a while to get fully comfortable with it. However, we'll specifically cover some features which will be useful for the following exercises.

Using ellipses. You can use an ellipsis ... in an indexing expression to avoid repeated : and to write indexing expressions that work on varying numbers of input dimensions. For example, x[..., 0] is equivalent to x[:, :, 0] if x is 3D, and equivalent to x[:, :, :, 0] if x is 4D.

Boolean indexing. Boolean indexing allows you to select elements of a tensor based on a corresponding boolean tensor. Only the elements where the boolean value is True are selected. The following example shows boolean indexing and broadcasting:

D = t.ones(2)  # Tensor with shape (2,)
E = t.zeros(3, 2)  # Tensor with shape (3, 2)
E[[True, False, True], :].shape  # Shape: (2, 2)
E[[True, False, True], :] = D  # Assign values from D to selected rows
print(E)
# output:
# tensor([[1., 1.],
#         [0., 0.],
#         [1., 1.]])

Note that E gets modified by the assignment above. This is not because E[[True, False, True], :] is a view of E - boolean indexing is advanced indexing, so reading it gives you a copy. The assignment works because writing to E[...] is a separate operation which writes back into E itself. We'll look a bit more at this topic later.

Multi-dimensional indexing. You can also apply boolean indexing to multiple dimensions simultaneously. This allows you to select elements across several dimensions based on a boolean condition. The following example shows multi-dimensional boolean indexing and broadcasting:

D = t.ones(2)  # Tensor with shape (2,)
F = t.zeros(2, 2, 2)  # Tensor with shape (2, 2, 2)
F[[[True, True], [False, True]], :].shape  # Shape: (3, 2)
F[[[True, True], [False, True]], :] = D
print(F)
# output:
# tensor([[[1., 1.],
#          [1., 1.]],

#         [[0., 0.],
#          [1., 1.]]])

This works because our indexing creates a tensor of shape (3, 2), formed by stacking the three tensors F[0, 0], F[0, 1] and F[1, 1] (those where the corresponding index value is True) along a new leading dimension.

Summary of tips

Use ... to avoid repeated : in indexing expressions

Use &, |, and ~ for elementwise logical operations on tensors

Use parentheses to force the correct operator precedence

Use torch.any() or .any() to do logical reductions (you can do this over a single dimension, with the dim argument)

Tensors can broadcast along leading dimensions to operate together

e.g. t.ones(2) + t.ones(3, 2) works because the former tensor gets broadcasted to shape (3, 2)

We can index with boolean tensors to select multiple elements

e.g. if A.shape = (3, n1, n2, ...) then A[[True, False, True]] has shape (2, n1, n2, ...) and contains the first & third rows of A

We can index with multi-dimensional boolean tensors too

e.g. if A.shape = (2, 2, n1, n2, ...) then A[[[True, True], [False, True]], :] has shape (3, n1, n2, ...) and contains the elements A[0, 0], A[0, 1] and A[1, 1], stacked along a new leading dimension

Batched Ray-Segment Intersection

We'll now move on to implementing a batched version of our previous function which takes multiple rays, multiple line segments, and returns a boolean for each ray indicating whether any segment intersects with that ray.

Note - in the batched version, we don't want the solver to throw an exception just because some of the matrices are singular. As in the previous exercise, treat these pairs as non-intersections by returning False, including collinear overlaps.

Exercise - implement `intersect_rays_1d`

Difficulty: 🔴🔴🔴🔴⚪

Importance: 🔵🔵🔵🔵⚪

You should spend up to 25-40 minutes on this exercise. This will probably be one of the hardest exercises you'll complete today.

One part you might find difficult is dealing with the zero determinant cases. Previously we dealt with those by using try / except, but here we can't do that because we want to perform all the operations at once. Instead, we can use the following clever trick to find all the pairs of rays and segment that intersect uniquely at a valid point:

Figure out which matrices have zero determinant (e.g. with determinants.abs() < 1e-6)
Replace those matrices with the identity matrix t.eye(2), since this will certainly not raise an error when solving
Find the matrices s.t. u, v are in the required range and our original matrix was non-singular

This way, we've identified all pairs of rays and segments where an intersection point exists, is unique, and is valid (i.e. it's actually on the positive side of the ray (in the direction the camera is facing), and somewhere in the middle of the segment).

Once we have this 2D array of booleans representing whether each ray intersects with each segment, we can reduce using the torch function t.any to find the rays which intersect any segment.

def intersect_rays_1d(
    rays: Float[Tensor, "nrays 2 3"], segments: Float[Tensor, "nsegments 2 3"]
) -> Bool[Tensor, " nrays"]:
    """
    For each ray, return True if it intersects any segment.
    """
    raise NotImplementedError()


tests.test_intersect_rays_1d(intersect_rays_1d)
tests.test_intersect_rays_1d_special_case(intersect_rays_1d)

Help - I don't know how to set all my matrices to the identity

You should have some variable mat which is a batch of matrices, i.e. it has shape (n_rays, n_segments, 2, 2). You can then define is_singular as a boolean tensor of shape (n_rays, n_segments) which is true wherever the matrix is singular. Now, indexing mat[is_singular] returns a tensor of shape (N, 2, 2) where the [i, :, :]-th element is the i-th singular matrix. Thanks to broadcasting rules, you can set mat[is_singular] = t.eye(2), and the identity matrix with shape (2, 2) will get broadcasted to the left hand shape (N, 2, 2).

Solution

def intersect_rays_1d(
    rays: Float[Tensor, "nrays 2 3"], segments: Float[Tensor, "nsegments 2 3"]
) -> Bool[Tensor, " nrays"]:
    """
    For each ray, return True if it intersects any segment.
    """
    NR = rays.size(0)
    NS = segments.size(0)

    # Get just the x and y coordinates
    rays = rays[..., :2]
    segments = segments[..., :2]

    # Repeat rays and segments so that we can compuate the intersection of every (ray, segment) pair
    rays = einops.repeat(rays, "nrays p d -> nrays nsegments p d", nsegments=NS)
    segments = einops.repeat(segments, "nsegments p d -> nrays nsegments p d", nrays=NR)

    # Each element of `rays` is [[Ox, Oy], [Dx, Dy]]
    O = rays[:, :, 0]
    D = rays[:, :, 1]
    assert O.shape == (NR, NS, 2)

    # Each element of `segments` is [[L1x, L1y], [L2x, L2y]]
    L_1 = segments[:, :, 0]
    L_2 = segments[:, :, 1]
    assert L_1.shape == (NR, NS, 2)

    # Define matrix on left hand side of equation
    mat = t.stack([D, L_1 - L_2], dim=-1)
    # Get boolean of where matrix is singular, and replace it with the identity in these positions
    dets = t.linalg.det(mat)
    is_singular = dets.abs() < 1e-6
    assert is_singular.shape == (NR, NS)
    mat[is_singular] = t.eye(2)

    # Define vector on the right hand side of equation
    vec = L_1 - O

    # Solve equation, get results
    sol = t.linalg.solve(mat, vec)
    u = sol[..., 0]
    v = sol[..., 1]

    # Return boolean of (matrix is nonsingular, and soln is in correct range implying intersection)
    return ((u >= 0) & (v >= 0) & (v <= 1) & ~is_singular).any(dim=-1)

Using LLMs to accelerate learning

This section updated May 2026. If you're reading this in 2027, it's probably out of date already.

LLMs have come a long way, and the advice that used to be here sorely needed an update. If you're still using LLMs mostly by copy-pasting code back and forth into the browser, you're doing it wrong.

The main uplift has been coding agents: Tools like Claude Code or Codex allow an LLM to interact with your computer via the CLI, use tools to read/write to files, issue commands to run scripts and read the output. For working through the ARENA material itself this is less useful (and often harmful as it spoils exercises!) but for more tedious tasks or when you're working on more open-ended tasks they are invaluable: I don't want to see a single student writing a training loop from scratch when we're in project week!

This blog post goes into more detail on making effective use of coding agents, and this one covers recent developments in the space (as well as the dangers of letting coding agents loose on your computer! Best to do that on a remote box instead, like the ones we provide at ARENA).

Is using LLMs in this way cheating? It can be, if your first instinct is to reach for one rather than trying to understand the code yourself. But it's important here to bring up the distinction of playing in easy mode vs playing in hard mode.

There are situations where it's valuable for you to think about a problem for a while before moving forward because that deliberation will directly lead to you becoming a better researcher or engineer (e.g. when you're thinking of a hypothesis for how a circuit works while doing mechanistic interpretability on a transformer, or you're pondering which datastructure best fits your use case while implementing some RL algorithm). But there are also situations (like this one) where you'll get more value from speedrunning towards an understanding of certain code or concepts, and apply your understanding in subsequent exercises. It's important to find a balance! The main takeaway is to use your time efficiently: don't spend too long on an exercise if you're stuck, make use of LLMs to increase the speed at which you can learn.

2D Rays

Now we're going to make use of the z dimension and have rays emitted from the origin in both y and z dimensions.

Exercise - implement `make_rays_2d`

Difficulty: 🔴🔴🔴⚪⚪

Importance: 🔵🔵⚪⚪⚪

You should spend up to 10-20 minutes on this exercise.

Implement make_rays_2d analogously to make_rays_1d. The result should look like a pyramid with the tip at the origin.

def make_rays_2d(num_pixels_y: int, num_pixels_z: int, y_limit: float, z_limit: float) -> Float[Tensor, "nrays 2 3"]:
    """
    num_pixels_y: The number of pixels in the y dimension
    num_pixels_z: The number of pixels in the z dimension

    y_limit: At x=1, the rays should extend from -y_limit to +y_limit, inclusive of both.
    z_limit: At x=1, the rays should extend from -z_limit to +z_limit, inclusive of both.

    Returns: shape (num_rays=num_pixels_y * num_pixels_z, num_points=2, num_dims=3).
    """
    raise NotImplementedError()


tests.test_make_rays_2d(make_rays_2d)

rays_2d = make_rays_2d(10, 10, 0.3, 0.3)
render_lines_with_plotly(rays_2d)

Click to see the expected output

Help - I'm not sure how to implement this function.

Don't write it as a function right away. The most efficient way is to write and test each line individually in the REPL to verify it does what you expect before proceeding.

You can either build up the output tensor using torch.stack, or you can initialize the output tensor to its final size and then assign to slices like rays[:, 1, 1] = .... It's good practice to be able to do it both ways.

Each y coordinate needs a ray with each corresponding z coordinate - in other words this is an outer product. The most elegant way to do this is with two calls to einops.repeat. You can also accomplish this with unsqueeze, expand, and reshape combined.