quimb.tensor.optimize#

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

Module Contents#

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

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_network_to_backend(tn, to_constant[, tags, ...])

Parse tensor network to:

convert_raw_arrays(x, f)

Given a TensorNetwork, Tensor, or other possibly structured raw

get_autograd()

get_tensorflow()

get_torch()

inject_(arrays, tn)

identity_fn(x)

Attributes#

quimb.tensor.optimize._DEFAULT_BACKEND = 'jax'#
quimb.tensor.optimize._REAL_CONVERSION#
quimb.tensor.optimize._COMPLEX_CONVERSION#
class quimb.tensor.optimize.ArrayInfo(array)[source]#

Simple container for recording size and dtype information about arrays.

__slots__ = ('shape', 'size', 'dtype', 'iscomplex', 'real_size', 'equivalent_real_type', 'equivalent_complex_type')#
__repr__()[source]#

Return repr(self).

class quimb.tensor.optimize.Vectorizer(tree)[source]#

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.

pack(tree, name='vector')[source]#

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

unpack(vector=None)[source]#

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)[source]#

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

quimb.tensor.optimize._parse_opt_out(tn, constant_tags, to_constant)[source]#

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

quimb.tensor.optimize.parse_network_to_backend(tn, to_constant, tags=None, shared_tags=None, constant_tags=None)[source]#

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.convert_raw_arrays(x, f)[source]#

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.get_autograd()[source]#
class quimb.tensor.optimize.AutoGradHandler(device='cpu')[source]#
to_variable(x)[source]#
to_constant(x)[source]#
setup_fn(fn)[source]#
value(arrays)[source]#
value_and_grad(arrays)[source]#
class quimb.tensor.optimize.JaxHandler(jit_fn=True, device=None)[source]#
to_variable(x)[source]#
to_constant(x)[source]#
setup_fn(fn)[source]#
_setup_hessp(fn)[source]#
value(arrays)[source]#
value_and_grad(arrays)[source]#
hessp(primals, tangents)[source]#
quimb.tensor.optimize.get_tensorflow()[source]#
class quimb.tensor.optimize.TensorFlowHandler(jit_fn=False, autograph=False, experimental_compile=False, device=None)[source]#
to_variable(x)[source]#
to_constant(x)[source]#
setup_fn(fn)[source]#
value(arrays)[source]#
value_and_grad(arrays)[source]#
quimb.tensor.optimize.get_torch()[source]#
class quimb.tensor.optimize.TorchHandler(jit_fn=False, device=None)[source]#
to_variable(x)[source]#
to_constant(x)[source]#
setup_fn(fn)[source]#
_setup_backend_fn(arrays)[source]#
value(arrays)[source]#
value_and_grad(arrays)[source]#
quimb.tensor.optimize._BACKEND_HANDLERS#
class quimb.tensor.optimize.MultiLossHandler(autodiff_backend, executor=None, **backend_opts)[source]#
setup_fn(funcs)[source]#
_value_seq(arrays)[source]#
_value_par_seq(arrays)[source]#
value(arrays)[source]#
_value_and_grad_seq(arrays)[source]#
_value_and_grad_par(arrays)[source]#
value_and_grad(arrays)[source]#
quimb.tensor.optimize.inject_(arrays, tn)[source]#
class quimb.tensor.optimize.SGD[source]#

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

Adapted from autograd/misc/optimizers.py.

get_velocity(x)[source]#
__call__(fun, x0, jac, args=(), learning_rate=0.1, mass=0.9, maxiter=1000, callback=None, bounds=None, **kwargs)[source]#
class quimb.tensor.optimize.RMSPROP[source]#

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

Adapted from autograd/misc/optimizers.py.

get_avg_sq_grad(x)[source]#
__call__(fun, x0, jac, args=(), learning_rate=0.1, gamma=0.9, eps=1e-08, maxiter=1000, callback=None, bounds=None, **kwargs)[source]#
class quimb.tensor.optimize.ADAM[source]#

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

Adapted from autograd/misc/optimizers.py.

get_m(x)[source]#
get_v(x)[source]#
__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)[source]#
class quimb.tensor.optimize.NADAM[source]#

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

Adapted from autograd/misc/optimizers.py.

get_m(x)[source]#
get_v(x)[source]#
get_mus(beta1)[source]#
__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)[source]#
class quimb.tensor.optimize.ADABELIEF[source]#

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

Adapted from autograd/misc/optimizers.py.

get_m(x)[source]#
get_s(x)[source]#
__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)[source]#
quimb.tensor.optimize._STOC_GRAD_METHODS#
class quimb.tensor.optimize.MakeArrayFn(tn_opt, loss_fn, norm_fn, autodiff_backend)[source]#

Class wrapper so picklable.

__name__ = 'MakeArrayFn'#
__call__(arrays)[source]#
quimb.tensor.optimize.identity_fn(x)[source]#
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)[source]#

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

property d#
property nevals#

The number of gradient evaluations.

property optimizer#

The underlying optimizer that works with the vectorized functions.

property bounds#
_set_tn(tn)[source]#
_reset_tracking_info(loss_target=None)[source]#
reset(tn=None, clear_info=True, loss_target=None)[source]#

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)[source]#
_maybe_update_pbar()[source]#
_maybe_close_pbar()[source]#
_check_loss_target()[source]#
_maybe_call_callback()[source]#
vectorized_value(x)[source]#

The value of the loss function at vector x.

vectorized_value_and_grad(x)[source]#

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

vectorized_hessp(x, p)[source]#

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

__repr__()[source]#

Return repr(self).

get_tn_opt()[source]#

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

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

  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)[source]#

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)[source]#

Scipy based optimization, see optimize() for details.

optimize_basinhopping(n, nhop, temperature=1.0, jac=True, hessp=False, **options)[source]#

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)[source]#

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)[source]#

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)[source]#

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=())[source]#

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.