# 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 `(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:

- 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`

. - 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:

- 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. - 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`

.