Understanding Memory Aliasing for Speed and Correctness#
The aggressive reuse of memory is one of the ways through which PyTensor makes code fast, and it is important for the correctness and speed of your program that you understand how PyTensor might alias buffers.
This section describes the principles based on which PyTensor handles memory, and explains when you might want to alter the default behaviour of some functions and methods for faster performance.
The Memory Model: Two Spaces#
There are some simple principles that guide PyTensor’s handling of memory. The main idea is that there is a pool of memory managed by PyTensor, and PyTensor tracks changes to values in that pool.
PyTensor manages its own memory space, which typically does not overlap with the memory of normal Python variables that non-PyTensor code creates.
PyTensor functions only modify buffers that are in PyTensor’s memory space.
PyTensor’s memory space includes the buffers allocated to store
sharedvariables and the temporaries used to evaluate functions.
Physically, PyTensor’s memory space may be spread across the host, a GPU device(s), and in the future may even include objects on a remote machine.
The memory allocated for a
sharedvariable buffer is unique: it is never aliased to another
PyTensor’s managed memory is constant while PyTensor functions are not running and PyTensor’s library code is not running.
The default behaviour of a function is to return user-space values for outputs, and to expect user-space values for inputs.
The distinction between PyTensor-managed memory and user-managed memory can be
broken down by some PyTensor functions (e.g.
get_value and the
Out) by using a
This can make those methods faster (by avoiding copy operations) at the expense
of risking subtle bugs in the overall program (by aliasing memory).
The rest of this section is aimed at helping you to understand when it is safe
to use the
borrow=True argument and reap the benefits of faster code.
Borrowing when Constructing Function Objects#
borrow argument can also be provided to the
that control how
pytensor.function handles its argument[s] and return value[s].
import pytensor import pytensor.tensor as at from pytensor.compile.io import In, Out x = at.matrix() y = 2 * x f = pytensor.function([In(x, borrow=True)], Out(y, borrow=True))
Borrowing an input means that PyTensor will treat the argument you provide as if
it were part of PyTensor’s pool of temporaries. Consequently, your input
may be reused as a buffer (and overwritten!) during the computation of other variables in the
course of evaluating that function (e.g.
Borrowing an output means that PyTensor will not insist on allocating a fresh
output buffer every time you call the function. It will possibly reuse the same one as
on a previous call, and overwrite the old content. Consequently, it may overwrite
old return values through side-effect.
Those return values may also be overwritten in
the course of evaluating another compiled function (for example, the output
may be aliased to a
shared variable). So be careful to use a borrowed return
value right away before calling any more PyTensor functions.
The default is of course to not borrow internal results.
It is also possible to pass a
return_internal_type=True flag to the
variable which has the same interpretation as the
get_value function. Unlike
borrow=True arguments to
Out() are not guaranteed to avoid copying an output value. They are just
hints that give more flexibility to the compilation and rewriting of the
Take home message:
When an input
x to a function is not needed after the function
returns and you would like to make it available to PyTensor as
additional workspace, then consider marking it with
In(x, borrow=True). It
may make the function faster and reduce its memory requirement. When a return
y is large (in terms of memory footprint), and you only need to read
from it once, right away when it’s returned, then consider marking it with an