Batching and broadcasting in Boulder Opal
Approaches to handle multidimensional data efficiently in graphs
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
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
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
(2, 2), and
(3, 2) along
axis=0, you will obtain a new matrix of shape
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.
tensor creates a new tensor by extracting only the first element of the left axis of the
tensor have one less axis than the original
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 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,
tensor[:, 2] will take only the third element of the second axis (since the index begins at 0), but keep the entire first axis.
tensor was originally a matrix of shape
(2, 3), then
tensor[:, 2] will be a vector with shape
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
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.
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.
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".
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 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:
- 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
- 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
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
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:
- If one tensor's shape has fewer axes than the other, a new axis with length
1is added to the left of its shape, until both tensors have the same number of axes.
- For each axis with length
1elements 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.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 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
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
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
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
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
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_stf, or the parameter
When you do that, all the axes that are to the left of the
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
batch_shape=(7, 1), you obtain a new PWC with
value_shape=(3, 2) and
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
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