Batching and broadcasting in Boulder Opal

Boulder Opal uses graphs to represent the custom simulations or optimizations that you want to calculate. In many problems of quantum control, you might want to build a graph that can handle large amounts of data, such as experimental data that you're analyzing, or different theoretical conditions to which you want to subject your system. Accordingly, you can structure your data into multidimensional objects that allow you to operate on all of it simultaneously, instead of building loops where an operation is repeated in each portion of the data. This technique of vectorizing or batching is recommended to improve the performance of your calculation, and is one of the tips listed in the topic Improving calculation performance in graphs.

This topic explains how higher-dimensional vectorized data structures are created in Boulder Opal by batching together lower-dimensional objects with the same shape. It also explains how operations between batches of different shapes work, in a process known as broadcasting. Together, these are important tools to make your calculations simpler and faster.

Shapes of tensors

The tensors that you use in Boulder Opal graphs can have multiple elements, which you can organize into different axes. A tensor with four elements, for example, can be a vector where all the four elements are arranged sequentially, or a matrix where the elements are organized in two rows and columns. Each of these directions in which we can organize the elements is called an axis, or dimension: matrices have two axes, vectors have one axis, scalars don't have axes. Boulder Opal allows you to create tensors with as many axes as you need.

Just like NumPy arrays, tensors from Boulder Opal have a shape attribute that specifies the number of elements in each of its axes. This attribute takes the form of a tuple of integer values, where each integer is the number of elements in each axis—in other words, the length of the axis. A $2\times 2$ square matrix, for example, has shape (2, 2), while a four-element vector has shape (4,). The shape of scalars is simply ().

Manipulating shapes of tensors

Boulder Opal offers several graph operations that you can use to manipulate the shape of a tensor. The main operation for this purpose is graph.reshape, which returns a tensor with the same elements as your input tensor (and in the same order) but with the new shape that you specify. For example, if your tensor is a $4 \times 4$ matrix and hence has shape (4, 4), you could reshape it into a $2 \times 8$ matrix with shape (2, 8), or into a four-axed tensor with shape (2, 2, 2, 2). Notice that the number of elements of the tensor has to stay the same: you can't reshape a (4, 4) tensor into a (3, 3) tensor, for example.

There are also other operations capable of changing the shape of the tensor in more particular ways. One example is graph.transpose, which by default creates a new tensor where the order of the axes is reversed: if you pass a matrix with shape (2, 3) to this node, it will output a matrix of shape (3, 2). Another example is the Boulder Opal operation graph.concatenate, which you can use to join together smaller tensors and create a larger one with all the information. For example, if you concatenate matrices of shape (5, 2), (2, 2), and (3, 2) along axis=0, you will obtain a new matrix of shape (10, 2). Boulder Opal also has a graph.sum operation, which allows you to sum all elements along the axes that you choose, resulting in an output tensor where the axes are no longer present.

Besides using nodes to manipulate the shape of a tensor, you can also use the same basic slicing and indexing rules from NumPy. In its simplest form, slicing is done by writing the indices of the elements that you want to extract in square brackets, after the name of the tensor. Note that slicing a tensor in Boulder Opal doesn't simply change the shape of the original tensor, but rather creates a new tensor with the requested shape. For example, tensor[0] creates a new tensor by extracting only the first element of the left axis of the tensor. This makes tensor[0] have one less axis than the original tensor. You can also specify more than one index at a time: tensor[0, 2] will create a new tensor with two less axes than tensor. If tensor was originally a matrix with shape (2, 3), then tensor[0, 2] will be a scalar with shape ().

Slicing also allows you to select entire axes, instead of just specific indices from it. To do this, replace the index with a colon, :. For example, tensor[:, 2] will take only the third element of the second axis (since the index begins at 0), but keep the entire first axis. If the tensor was originally a matrix of shape (2, 3), then tensor[:, 2] will be a vector with shape (2,).

Additionally, you can use the notation start:stop:step to extract a slice of an axis that starts at the element start, ends before the element stop, and only takes elements spaced by step. For example, tensor[2:10:2] will extract the elements 2, 4, 6, and 8 from the left axis (keep in mind that Python indexes the elements beginning from 0). Note that if you omit the step it will be assumed to be 1, so the slice tensor[2:10] will contain all the elements from 2 to 9.

Adding extra axes

A particularly important shape manipulation that you can do with indexing is adding new empty axes, which you can then use to create batches and improve the performance of your calculation (read more in the "Batching" section of this notebook). You can add new axes by using the keyword None in the position where you want to include the new axis. For example, tensor[None] adds one extra axis with length 1 to the left side of the tensor's shape, while tensor[:, None] adds a new axis between the first and the second left axes. If tensor originally had shape (2, 3, 2), then tensor[None] will have shape (1, 2, 3, 2) and tensor[:, None] will have shape (2, 1, 3, 2).

Another tip for adding new axes is that you can use an ellipsis (...) as a convenient shorthand meaning "all the remaining axes". For example, tensor[..., None] adds one extra axis with 1 element to the right side of the shape, while tensor[None, ...] is equivalent to tensor[None] in adding an extra axis to the left side of the shape. You can also use ... together with : in the same indexing. For example, if you want to turn a tensor of shape (2, 3, 2) into a tensor of shape (2, 3, 1, 2), you can write tensor[..., None, :] to insert a new axis between the last and the-second-to-last initial axes.

Broadcasting

Broadcasting extends operations between tensors (such as addition) to support tensors that have different shapes. Arithmetic operations, for example, are performed element-wise when the tensors have the same shape. When the tensors have different but compatible shapes, the operation proceeds as if the dimensions of the operands had been expanded until they were of the same size. The broadcasting rules of Boulder Opal determine when two tensors are considered compatible, and what the shape of the output should be.

Broadcasting rules for tensors

The rules for broadcasting tensors in Boulder Opal are the same as the rules for broadcasting NumPy arrays. To see if two shapes are broadcastable, follow these two steps:

1. If one shape has fewer axes than the other, add axes to its left until both shapes have the same number of axes. Let these new axes have length 1.
2. Sequentially compare the length of each axis of the two shapes. For the shapes to be broadcastable, each pair of axis lengths must either be equal, or one of them must be 1.

For example, the shapes (1, 2, 3) and (3, 2, 1) are broadcastable, but the shapes (1, 2, 3) and (1, 2, 2) are not—their last axes have neither the same length, nor length 1. Similarly, the shapes (2, 3) and (3, 2, 1) are broadcastable, while the shapes (2, 3) and (1, 2, 2) are not.

If the shapes of two tensors are broadcastable, the operation proceeds as if it were applied to tensors whose dimensions had been expanded in the following manner:

1. If one tensor's shape has fewer axes than the other, a new axis with length 1 is added to the left of its shape, until both tensors have the same number of axes.
2. For each axis with length 1 elements in one tensor but not in the other, copies of the smaller tensor are concatenated along that axis until it has the same length as the larger tensor.

This means that a simple arithmetic operation between a tensor with values [1, 2] and another tensor with values [[10, 20], [30, 40]] happens as if the first tensor had been expanded to become [[1, 2], [1, 2]]—the number of axes is increased to match the second tensor, and the axis with only one element is duplicated until it has as many elements as the largest one. In this example, the result of adding a tensor with values [1, 2] and a tensor with values [[10, 20], [30, 40]] is a new tensor with values [[11, 22], [31, 42]], while a multiplication results in [[10, 40], [30, 80]].

Using broadcasting in practice

Besides the arithmetic operations, other types of nodes will attempt to perform broadcasting whenever possible, and whenever it makes sense. For example, nodes such as graph.kron and graph.matmul perform broadcasting in all but the two right dimensions, which obey the rules of Kronecker products and matrix multiplication instead. To learn about the specifics of each operation, consult our reference documentation.

For more concrete examples of broadcasting in practice, see the How to evaluate control susceptibility to quasi-static noise and the How to import and use pulses from the Q-CTRL Open Controls library user guides.

Batching

Batching uses higher-dimensional numerical objects to simultaneously perform multiple computations on lower-dimensional objects. Doing this allows you to reduce the number of operations in your calculation, which can speed it up. For example, instead of individually calculating the traces of five matrices with shape (2, 2), you can create a batch of matrices with shape (5, 2, 2) and then calculate each of their traces in a single operation. In general, you should attempt to replace loops with batch operations whenever possible to improve the efficiency of your code.

Batch dimensions in tensors

In Boulder Opal, the left axes of a tensor's shape are treated as their batch dimensions. To create a batch, it is easiest to start by creating a tensor with the parameter that changes for each element of the batch. For example, if you want to create a batch of 40 sigma_x matrices multiplied by different values of amplitude noise, you can start by creating a tensor with the 40 values of noise, moving the values of the noise to the left axis with tensor[:, None, None], and then multiplying them by the matrices with tensor[:, None, None] * sigma_x. In this example, tensor has shape (40,), tensor[:, None, None] has shape (40, 1, 1), and tensor[:, None, None] * sigma_x has shape (40, 2, 2).

Alternatively, you can create a batch tensor by explicitly specifying the values of all its elements, without parameterizing it first. In this case, the batch is created in the same way that would create a single tensor, but with more axes present. However, note that your code will be more compact if you parameterize the batch creation whenever possible.

Finally, depending on how your data is structured, you may want to give your tensor multiple batch axes. Consider again the example of creating a batch of sigma_x matrices multiplied by 40 values of amplitude noise, but assume that you are also interested in all its possible combinations with 30 values of additive noise. In this case, it can be convenient to keep noise values of different origin in different axes, which results in a batch of shape (40, 30, 2, 2).

Using batches in practice

The user guides in the Boulder Opal documentation provide examples of how to use batching in practice to simplify calculations. To illustrate the example in the previous subsection, you can read about how to create batches of Hamiltonians with different quasi-static noise values in the How to evaluate control susceptibility to quasi-static noise user guide. Other examples include the use of batches for automated optimization, simulation of multiple circuits, and system identification for small and large amounts of data.

Batching and broadcasting for PWCs and STFs

Piecewise-constant (PWC) and sampleable tensor-valued functions (STF) of time are special kinds of data structures used mainly to represent controls. Although they take tensor values in time, they also associate time durations to those values, which requires some special attention.

When considering operations with PWCs and STFs, it is important to keep in mind the types of your operands. Most operations, including the basic arithmetic ones, do not allow mixing PWCs and STFs—operating between them requires converting one into the other with graph.discretize_stf or graph.convolve_pwc. However, you can mix either PWCs or STFs with tensors in arithmetic operations: the result will be a PWC or STF, respectively.

Batch dimensions in PWCs and STFs

Objects like PWCs and STFs do not have a single shape attribute, but rather separate batch_shape and value_shape attributes. This means that these objects keep track separately of the axes that pertain to the batches and the axes that pertain to the value object itself. This is different from the case of tensors, for which the batch shape and the value shape are just the left and right sides of the same shape attribute. For example, if a tensor representing a batch of $2\times 2$ matrices has shape=(40, 30, 2, 2), then the corresponding PWC or STF would have batch_shape=(40, 30) and value_shape=(2, 2).

When creating a batched PWC or STF, you have to specify which axes correspond to the batch and which axes correspond to the values. You can do this by specifying the parameter batch_dimension_count in nodes such as graph.constant_pwc or graph.constant_stf, or the parameter time_dimension in graph.pwc. When you do that, all the axes that are to the left of the batch_dimension_count (or time_dimension) are considered batch dimensions. Note that by default the batch_dimension_count is set to zero, which means that no batch is created unless you explicitly request it.

Broadcasting rules for PWCs and STFs

As PWCs and STFs have two separate shape attributes, broadcasting for them is performed separately. For example if you're operating on a PWC with value_shape=(3, 2) and batch_shape=(3,) and another PWC with value_shape=(2,) and batch_shape=(7, 1), you obtain a new PWC with value_shape=(3, 2) and batch_shape=(7, 3).

When operating between tensors and PWCs or STFs, the tensors are treated as PWCs or STFs whose value_shape corresponds to the shape of the tensor. For example, if you have a PWC with value_shape=(2, 3) and batch_shape=(4, 5), then adding to it a tensor with shape (4, 2, 1) results in a new PWC with value_shape=(4, 2, 3) but whose batch_shape is still (4, 5). If you want to treat a tensor as a batch in an operation with a PWC or STF, you need to convert it first using either graph.constant_pwc or graph.constant_stf.