quimb.tensor.optimize

Support for optimizing tensor networks using automatic differentiation to automatically derive gradients for input to scipy optimizers.

Attributes

_DEFAULT_BACKEND
_REAL_CONVERSION
_COMPLEX_CONVERSION
_VARIABLE_TAG
variable_finder
_BACKEND_HANDLERS
_STOC_GRAD_METHODS

Classes

ArrayInfo	Simple container for recording size and dtype information about arrays.
Vectorizer	Object for mapping back and forth between any pytree of mixed
AutoGradHandler
JaxHandler
TensorFlowHandler
TorchHandler
MultiLossHandler
SGD	Stateful scipy.optimize.minimize compatible implementation of
RMSPROP	Stateful scipy.optimize.minimize compatible implementation of
ADAM	Stateful scipy.optimize.minimize compatible implementation of
CADAM	Cautious ADAM - https://arxiv.org/abs/2411.16085.
NADAM	Stateful scipy.optimize.minimize compatible implementation of
ADABELIEF	Stateful scipy.optimize.minimize compatible implementation of
MakeArrayFn	Class wrapper so picklable.
TNOptimizer	Globally optimize tensors within a tensor network with respect to any

Functions

_parse_opt_in(tn, tags, shared_tags, to_constant)	Parse a tensor network where tensors are assumed to be constant unless
_parse_opt_out(tn, constant_tags, to_constant)	Parse a tensor network where tensors are assumed to be variables unless
_parse_pytree_to_backend(x, to_constant)	Parse a arbitrary pytree, collecting variables. There is not opting in
parse_network_to_backend(tn, to_constant[, tags, ...])	Parse tensor network to:
_inject_variables_pytree(arrays, tree)
inject_variables(arrays, tn)	Given the list of optimized variables arrays and the target tensor
convert_raw_arrays(x, f)	Given a TensorNetwork, Tensor, or other possibly structured raw
convert_variables_to_numpy(x)
get_autograd()
get_tensorflow()
get_torch()
identity_fn(x)

Module Contents

quimb.tensor.optimize._DEFAULT_BACKEND = 'jax'

quimb.tensor.optimize._REAL_CONVERSION

quimb.tensor.optimize._COMPLEX_CONVERSION

class quimb.tensor.optimize.ArrayInfo(array)

Simple container for recording size and dtype information about arrays.

__slots__ = ('shape', 'size', 'dtype', 'iscomplex', 'real_size', 'equivalent_real_type', 'equivalent_complex_type')

shape

size

dtype

equivalent_real_type

equivalent_complex_type

iscomplex

real_size

__repr__()

class quimb.tensor.optimize.Vectorizer(tree)

Object for mapping back and forth between any pytree of mixed real/complex n-dimensional arrays to a single, real, double precision numpy vector, as required by scipy.optimize routines.

Parameters:

• tree (pytree of array) – Any nested container of arrays, which will be flattened and packed into a single float64 vector.

• is_leaf (callable, optional) – A function which takes a single argument and returns True if it is a leaf node in the tree and should be extracted, False otherwise. Defaults to everything that is not a tuple, list or dict.

infos = []

d = 0

ref_tree

pack(tree, name='vector')

Take arrays and pack their values into attribute .{name}, by default .vector.

unpack(vector=None)

Turn the single, flat vector into a sequence of arrays.

quimb.tensor.optimize._VARIABLE_TAG = '__VARIABLE{}__'

quimb.tensor.optimize.variable_finder

quimb.tensor.optimize._parse_opt_in(tn, tags, shared_tags, to_constant)

Parse a tensor network where tensors are assumed to be constant unless tagged.

quimb.tensor.optimize._parse_opt_out(tn, constant_tags, to_constant)

Parse a tensor network where tensors are assumed to be variables unless tagged.

quimb.tensor.optimize._parse_pytree_to_backend(x, to_constant)

Parse a arbitrary pytree, collecting variables. There is not opting in or out, all networks, tensors and raw arrays are considered variables.

quimb.tensor.optimize.parse_network_to_backend(tn, to_constant, tags=None, shared_tags=None, constant_tags=None)

Parse tensor network to:

• identify the dimension of the optimisation space and the initial point of the optimisation from the current values in the tensor network,

• add variable tags to individual tensors so that optimisation vector values can be efficiently reinserted into the tensor network.

There are two different modes:

• ‘opt in’ : tags (and optionally shared_tags) are specified and only these tensor tags will be optimised over. In this case constant_tags is ignored if it is passed,

• ‘opt out’ : tags is not specified. In this case all tensors will be optimised over, unless they have one of constant_tags tags.

Parameters:

• tn (TensorNetwork) – The initial tensor network to parse.

• to_constant (Callable) – Function that fixes a tensor as constant.

• tags (str, or sequence of str, optional) – Set of opt-in tags to optimise.

• shared_tags (str, or sequence of str, optional) – Subset of opt-in tags to joint optimise i.e. all tensors with tag s in shared_tags will correspond to the same optimisation variables.

• constant_tags (str, or sequence of str, optional) – Set of opt-out tags if tags not passed.

Returns:

• tn_ag (TensorNetwork) – Tensor network tagged for reinsertion of optimisation variable values.

• variables (list) – List of variables extracted from tn.

quimb.tensor.optimize._inject_variables_pytree(arrays, tree)

quimb.tensor.optimize.inject_variables(arrays, tn)

Given the list of optimized variables arrays and the target tensor network or pytree tn, inject the variables back in.

quimb.tensor.optimize.convert_raw_arrays(x, f)

Given a TensorNetwork, Tensor, or other possibly structured raw array, return a copy where the underyling data has had f applied to it. Structured raw arrays should implement the tree = get_params() and set_params(tree) methods which get or set their underlying data using an arbitrary pytree.

quimb.tensor.optimize.convert_variables_to_numpy(x)

quimb.tensor.optimize.get_autograd()

class quimb.tensor.optimize.AutoGradHandler(device='cpu')

to_variable(x)

to_constant(x)

setup_fn(fn)

value(arrays)

value_and_grad(arrays)

class quimb.tensor.optimize.JaxHandler(jit_fn=True, device=None)

jit_fn = True

device = None

to_variable(x)

to_constant(x)

setup_fn(fn)

_setup_hessp(fn)

value(arrays)

value_and_grad(arrays)

hessp(primals, tangents)

quimb.tensor.optimize.get_tensorflow()

class quimb.tensor.optimize.TensorFlowHandler(jit_fn=False, autograph=False, experimental_compile=False, device=None)

jit_fn = False

autograph = False

experimental_compile = False

device = None

to_variable(x)

to_constant(x)

setup_fn(fn)

value(arrays)

value_and_grad(arrays)

quimb.tensor.optimize.get_torch()

class quimb.tensor.optimize.TorchHandler(jit_fn=False, device=None, dtype=None)

jit_fn = False

device = None

dtype = None

to_variable(x)

to_constant(x)

setup_fn(fn)

_setup_backend_fn(arrays)

value(arrays)

value_and_grad(arrays)

quimb.tensor.optimize._BACKEND_HANDLERS

class quimb.tensor.optimize.MultiLossHandler(autodiff_backend, executor=None, **backend_opts)

autodiff_backend

backend_opts

executor = None

handlers

to_constant

setup_fn(funcs)

_value_seq(arrays)

_value_par_seq(arrays)

value(arrays)

_value_and_grad_seq(arrays)

_value_and_grad_par(arrays)

value_and_grad(arrays)

class quimb.tensor.optimize.SGD

Stateful scipy.optimize.minimize compatible implementation of stochastic gradient descent with momentum.

Adapted from autograd/misc/optimizers.py.

OptimizeResult

_i = 0

_velocity = None

get_velocity(x)

__call__(fun, x0, jac, args=(), learning_rate=0.1, mass=0.9, maxiter=1000, callback=None, bounds=None, **kwargs)

class quimb.tensor.optimize.RMSPROP

Stateful scipy.optimize.minimize compatible implementation of root mean squared prop: See Adagrad paper for details.

Adapted from autograd/misc/optimizers.py.

OptimizeResult

_i = 0

_avg_sq_grad = None

get_avg_sq_grad(x)

__call__(fun, x0, jac, args=(), learning_rate=0.1, gamma=0.9, eps=1e-08, maxiter=1000, callback=None, bounds=None, **kwargs)

class quimb.tensor.optimize.ADAM(cautious=False)

Stateful scipy.optimize.minimize compatible implementation of ADAM - http://arxiv.org/pdf/1412.6980.pdf.

Adapted from autograd/misc/optimizers.py.

OptimizeResult

_i = 0

_m = None

_v = None

_cautious = False

get_m(x)

get_v(x)

__call__(fun, x0, jac, args=(), learning_rate=0.001, beta1=0.9, beta2=0.999, eps=1e-08, maxiter=1000, callback=None, bounds=None, **kwargs)

class quimb.tensor.optimize.CADAM

Bases: ADAM

Cautious ADAM - https://arxiv.org/abs/2411.16085.

class quimb.tensor.optimize.NADAM

Stateful scipy.optimize.minimize compatible implementation of NADAM - [Dozat - http://cs229.stanford.edu/proj2015/054_report.pdf].

Adapted from autograd/misc/optimizers.py.

OptimizeResult

_i = 0

_m = None

_v = None

_mus = None

get_m(x)

get_v(x)

get_mus(beta1)

__call__(fun, x0, jac, args=(), learning_rate=0.001, beta1=0.9, beta2=0.999, eps=1e-08, maxiter=1000, callback=None, bounds=None, **kwargs)

class quimb.tensor.optimize.ADABELIEF

Stateful scipy.optimize.minimize compatible implementation of ADABELIEF - https://arxiv.org/abs/2010.07468.

Adapted from autograd/misc/optimizers.py.

OptimizeResult

_i = 0

_m = None

_s = None

get_m(x)

get_s(x)

__call__(fun, x0, jac, args=(), learning_rate=0.001, beta1=0.9, beta2=0.999, eps=1e-08, maxiter=1000, callback=None, bounds=None, **kwargs)

quimb.tensor.optimize._STOC_GRAD_METHODS

class quimb.tensor.optimize.MakeArrayFn(tn_opt, loss_fn, norm_fn, autodiff_backend)

Class wrapper so picklable.

__name__ = 'MakeArrayFn'

tn_opt

loss_fn

norm_fn

autodiff_backend

__call__(arrays)

quimb.tensor.optimize.identity_fn(x)

class quimb.tensor.optimize.TNOptimizer(tn, loss_fn, norm_fn=None, loss_constants=None, loss_kwargs=None, tags=None, shared_tags=None, constant_tags=None, loss_target=None, optimizer='L-BFGS-B', progbar=True, bounds=None, autodiff_backend='AUTO', executor=None, callback=None, **backend_opts)

Globally optimize tensors within a tensor network with respect to any loss function via automatic differentiation. If parametrized tensors are used, optimize the parameters rather than the raw arrays.

Parameters:

• tn (TensorNetwork) – The core tensor network structure within which to optimize tensors.

• loss_fn (callable or sequence of callable) – The function that takes tn (as well as loss_constants and loss_kwargs) and returns a single real ‘loss’ to be minimized. For Hamiltonians which can be represented as a sum over terms, an iterable collection of terms (e.g. list) can be given instead. In that case each term is evaluated independently and the sum taken as loss_fn. This can reduce the total memory requirements or allow for parallelization (see executor).

• norm_fn (callable, optional) – A function to call before loss_fn that prepares or ‘normalizes’ the raw tensor network in some way.

• loss_constants (dict, optional) – Extra tensor networks, tensors, dicts/list/tuples of arrays, or arrays which will be supplied to loss_fn but also converted to the correct backend array type.

• loss_kwargs (dict, optional) – Extra options to supply to loss_fn (unlike loss_constants these are assumed to be simple options that don’t need conversion).

• tags (str, or sequence of str, optional) – If supplied, only optimize tensors with any of these tags.

• shared_tags (str, or sequence of str, optional) – If supplied, each tag in shared_tags corresponds to a group of tensors to be optimized together.

• constant_tags (str, or sequence of str, optional) – If supplied, skip optimizing tensors with any of these tags. This ‘opt-out’ mode is overridden if either tags or shared_tags is supplied.

• loss_target (float, optional) – Stop optimizing once this loss value is reached.

• optimizer (str, optional) – Which scipy.optimize.minimize optimizer to use (the 'method' kwarg of that function). In addition, quimb implements a few custom optimizers compatible with this interface that you can reference by name - {'adam', 'nadam', 'rmsprop', 'sgd'}.

• executor (None or Executor, optional) – To be used with term-by-term Hamiltonians. If supplied, this executor is used to parallelize the evaluation. Otherwise each term is evaluated in sequence. It should implement the basic concurrent.futures (PEP 3148) interface.

• progbar (bool, optional) – Whether to show live progress.

• bounds (None or (float, float), optional) – Constrain the optimized tensor entries within this range (if the scipy optimizer supports it).

• autodiff_backend ({'jax', 'autograd', 'tensorflow', 'torch'}, optional) – Which backend library to use to perform the automatic differentation (and computation).

• callback (callable, optional) –

  A function to call after each optimization step. It should take the current TNOptimizer instance as its only argument. Information such as the current loss and number of evaluations can then be accessed:

  def callback(tnopt):
      print(tnopt.nevals, tnopt.loss)
  

• backend_opts – Supplied to the backend function compiler and array handler. For example jit_fn=True or device='cpu' .

progbar = True

tags = None

shared_tags = None

constant_tags = None

_autodiff_backend = 'AUTO'

_multiloss

norm_fn = None

loss_constants

loss_kwargs

property bounds

property optimizer

The underlying optimizer that works with the vectorized functions.

callback = None

_set_tn(tn)

_reset_tracking_info(loss_target=None)

reset(tn=None, clear_info=True, loss_target=None)

Reset this optimizer without losing the compiled loss and gradient functions.

Parameters:

• tn (TensorNetwork, optional) – Set this tensor network as the current state of the optimizer, it must exactly match the original tensor network.

• clear_info (bool, optional) – Clear the tracked losses and iterations.

_maybe_init_pbar(n)

_maybe_update_pbar()

_maybe_close_pbar()

_check_loss_target()

_maybe_call_callback()

vectorized_value(x)

The value of the loss function at vector x.

vectorized_value_and_grad(x)

The value and gradient of the loss function at vector x.

vectorized_hessp(x, p)

The action of the hessian at point x on vector p.

__repr__()

property d

property nevals

The number of gradient evaluations.

get_tn_opt()

Extract the optimized tensor network, this is a three part process:

1. inject the current optimized vector into the target tensor network or pytree,

2. run it through norm_fn,

3. drop any tags used to identify variables.

Returns:

tn_opt

Return type:

TensorNetwork

optimize(n, tol=None, jac=True, hessp=False, optlib='scipy', **options)

Run the optimizer for n function evaluations, using by default scipy.optimize.minimize() as the driver for the vectorized computation. Supplying the gradient and hessian vector product is controlled by the jac and hessp options respectively.

Parameters:

• n (int) – Notionally the maximum number of iterations for the optimizer, note that depending on the optimizer being used, this may correspond to number of function evaluations rather than just iterations.

• tol (None or float, optional) – Tolerance for convergence, note that various more specific tolerances can usually be supplied to options, depending on the optimizer being used.

• jac (bool, optional) – Whether to supply the jacobian, i.e. gradient, of the loss function.

• hessp (bool, optional) – Whether to supply the hessian vector product of the loss function.

• optlib ({'scipy', 'nlopt'}, optional) – Which optimization library to use.

• options – Supplied to scipy.optimize.minimize() or whichever optimizer is being used.

Returns:

tn_opt

Return type:

TensorNetwork

optimize_scipy(n, tol=None, jac=True, hessp=False, **options)

Scipy based optimization, see optimize() for details.

optimize_basinhopping(n, nhop, temperature=1.0, jac=True, hessp=False, **options)

Run the optimizer for using scipy.optimize.basinhopping() as the driver for the vectorized computation. This performs nhop local optimization each with n iterations.

Parameters:

• n (int) – Number of iterations per local optimization.

• nhop (int) – Number of local optimizations to hop between.

• temperature (float, optional) – H

• options – Supplied to the inner scipy.optimize.minimize() call.

Returns:

tn_opt

Return type:

TensorNetwork

optimize_nlopt(n, tol=None, jac=True, hessp=False, ftol_rel=None, ftol_abs=None, xtol_rel=None, xtol_abs=None)

Run the optimizer for n function evaluations, using nlopt as the backend library to run the optimization. Whether the gradient is computed depends on which optimizer is selected, see valid options at https://nlopt.readthedocs.io/en/latest/NLopt_Algorithms/.

The following scipy optimizer options are automatically translated to the corresponding nlopt algorithms: {“l-bfgs-b”, “slsqp”, “tnc”, “cobyla”}.

Parameters:

• n (int) – The maximum number of iterations for the optimizer.

• tol (None or float, optional) – Tolerance for convergence, here this is taken to be the relative tolerance for the loss (ftol_rel below overrides this).

• jac (bool, optional) – Whether to supply the jacobian, i.e. gradient, of the loss function.

• hessp (bool, optional) – Whether to supply the hessian vector product of the loss function.

• ftol_rel (float, optional) – Set relative tolerance on function value.

• ftol_abs (float, optional) – Set absolute tolerance on function value.

• xtol_rel (float, optional) – Set relative tolerance on optimization parameters.

• xtol_abs (float, optional) – Set absolute tolerances on optimization parameters.

Returns:

tn_opt

Return type:

TensorNetwork

optimize_ipopt(n, tol=None, **options)

Run the optimizer for n function evaluations, using ipopt as the backend library to run the optimization via the python package cyipopt.

Parameters:

n (int) – The maximum number of iterations for the optimizer.

Returns:

tn_opt

Return type:

TensorNetwork

optimize_nevergrad(n)

Run the optimizer for n function evaluations, using nevergrad as the backend library to run the optimization. As the name suggests, the gradient is not required for this method.

Parameters:

n (int) – The maximum number of iterations for the optimizer.

Returns:

tn_opt

Return type:

TensorNetwork

plot(xscale='symlog', xscale_linthresh=20, zoom='auto', hlines=())

Plot the loss function as a function of the number of iterations.

Parameters:

• xscale (str, optional) – The scale of the x-axis. Default is "symlog", i.e. linear for the first part of the plot, and logarithmic for the rest, changing at xscale_linthresh.

• xscale_linthresh (int, optional) – The threshold for the change from linear to logarithmic scale, if xscale is "symlog". Default is 20.

• zoom (None or int, optional) – If not None, show an inset plot of the last zoom iterations.

• hlines (dict, optional) – A dictionary of horizontal lines to plot. The keys are the labels of the lines, and the values are the y-values of the lines.

Returns:

• fig (matplotlib.figure.Figure) – The figure object.

• ax (matplotlib.axes.Axes) – The axes object.