Trajectree 0.0.0__py3-none-any.whl → 0.0.1__py3-none-any.whl

trajectree/quimb/quimb/tensor/optimize.py

@@ -0,0 +1,1818 @@
+"""Support for optimizing tensor networks using automatic differentiation to
+automatically derive gradients for input to scipy optimizers.
+"""
+
+import functools
+import importlib.util
+import re
+import warnings
+from collections.abc import Iterable
+
+import numpy as np
+import tqdm
+from autoray import astype, get_dtype_name, to_numpy
+
+from ..core import prod
+from ..utils import (
+    default_to_neutral_style,
+    ensure_dict,
+    tree_flatten,
+    tree_map,
+    tree_unflatten,
+)
+from .contraction import contract_backend
+from .interface import get_jax
+from .tensor_core import (
+    TensorNetwork,
+    tags_to_oset,
+)
+
+if importlib.util.find_spec("jax") is not None:
+    _DEFAULT_BACKEND = "jax"
+elif importlib.util.find_spec("tensorflow") is not None:
+    _DEFAULT_BACKEND = "tensorflow"
+elif importlib.util.find_spec("torch") is not None:
+    _DEFAULT_BACKEND = "torch"
+else:
+    _DEFAULT_BACKEND = "autograd"
+
+
+_REAL_CONVERSION = {
+    "float32": "float32",
+    "float64": "float64",
+    "complex64": "float32",
+    "complex128": "float64",
+}
+
+_COMPLEX_CONVERSION = {
+    "float32": "complex64",
+    "float64": "complex128",
+    "complex64": "complex64",
+    "complex128": "complex128",
+}
+
+
+class ArrayInfo:
+    """Simple container for recording size and dtype information about arrays."""
+
+    __slots__ = (
+        "shape",
+        "size",
+        "dtype",
+        "iscomplex",
+        "real_size",
+        "equivalent_real_type",
+        "equivalent_complex_type",
+    )
+
+    def __init__(self, array):
+        self.shape = array.shape
+        self.size = prod(self.shape)
+        self.dtype = get_dtype_name(array)
+        self.equivalent_real_type = _REAL_CONVERSION[self.dtype]
+        self.equivalent_complex_type = _COMPLEX_CONVERSION[self.dtype]
+        self.iscomplex = "complex" in self.dtype
+        self.real_size = self.size * (2 if self.iscomplex else 1)
+
+    def __repr__(self):
+        return (
+            "ArrayInfo("
+            f"shape={self.shape}, "
+            f"size={self.size}, "
+            f"dtype={self.dtype}"
+            ")"
+        )
+
+
+class Vectorizer:
+    """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.
+    """
+
+    def __init__(self, tree):
+        arrays = []
+        self.infos = []
+        self.d = 0
+
+        def extracter(x):
+            arrays.append(x)
+            info = ArrayInfo(x)
+            self.infos.append(info)
+            self.d += info.real_size
+            return info
+
+        self.ref_tree = tree_map(extracter, tree)
+        self.pack(arrays)
+
+    def pack(self, tree, name="vector"):
+        """Take ``arrays`` and pack their values into attribute `.{name}`, by
+        default `.vector`.
+        """
+        arrays = tree_flatten(tree)
+
+        # create the vector if it doesn't exist yet
+        if not hasattr(self, name):
+            setattr(self, name, np.empty(self.d, "float64"))
+        x = getattr(self, name)
+
+        i = 0
+        for array, info in zip(arrays, self.infos):
+            if not isinstance(array, np.ndarray):
+                array = to_numpy(array)
+            # flatten
+            if info.iscomplex:
+                # view as real array of double the length
+                real_view = array.reshape(-1).view(info.equivalent_real_type)
+            else:
+                real_view = array.reshape(-1)
+            # pack into our vector
+            f = i + info.real_size
+            x[i:f] = real_view
+            i = f
+
+        return x
+
+    def unpack(self, vector=None):
+        """Turn the single, flat ``vector`` into a sequence of arrays."""
+        if vector is None:
+            vector = self.vector
+
+        i = 0
+        arrays = []
+        for info in self.infos:
+            # get the linear slice
+            f = i + info.real_size
+            array = vector[i:f]
+            i = f
+            if info.iscomplex:
+                # view as complex array of half the length
+                array = array.view(np.complex128)
+            # reshape (inplace)
+            array.shape = info.shape
+            if get_dtype_name(array) != info.dtype:
+                # cast as original dtype
+                array = astype(array, info.dtype)
+            arrays.append(array)
+
+        return tree_unflatten(
+            arrays, self.ref_tree, lambda x: isinstance(x, ArrayInfo)
+        )
+
+
+_VARIABLE_TAG = "__VARIABLE{}__"
+variable_finder = re.compile(r"__VARIABLE(\d+)__")
+
+
+def _parse_opt_in(tn, tags, shared_tags, to_constant):
+    """Parse a tensor network where tensors are assumed to be constant unless
+    tagged.
+    """
+    tn_ag = tn.copy()
+    variables = []
+
+    # tags where each individual tensor should get a separate variable
+    individual_tags = tags - shared_tags
+
+    # handle tagged tensors that are not shared
+    for t in tn_ag.select_tensors(individual_tags, "any"):
+        # append the raw data but mark the corresponding tensor
+        # for reinsertion
+        data = t.get_params()
+        variables.append(data)
+        t.add_tag(_VARIABLE_TAG.format(len(variables) - 1))
+
+    # handle shared tags
+    for tag in shared_tags:
+        var_name = _VARIABLE_TAG.format(len(variables))
+        test_data = None
+
+        for t in tn_ag.select_tensors(tag):
+            data = t.get_params()
+
+            # detect that this tensor is already variable tagged and skip
+            # if it is
+            if any(variable_finder.match(tag) for tag in t.tags):
+                warnings.warn(
+                    "TNOptimizer warning, tensor tagged with"
+                    " multiple `tags` or `shared_tags`."
+                )
+                continue
+
+            if test_data is None:
+                # create variable and store data
+                variables.append(data)
+                test_data = data
+            else:
+                # check that the shape of the variable's data matches the
+                # data of this new tensor
+                if test_data.shape != data.shape:
+                    raise ValueError(
+                        "TNOptimizer error, a `shared_tags` tag "
+                        "covers tensors with different numbers of"
+                        " params."
+                    )
+
+            # mark the corresponding tensor for reinsertion
+            t.add_tag(var_name)
+
+    # iterate over tensors which *don't* have any of the given tags
+    for t in tn_ag.select_tensors(tags, which="!any"):
+        t.modify(apply=to_constant)
+
+    return tn_ag, variables
+
+
+def _parse_opt_out(
+    tn,
+    constant_tags,
+    to_constant,
+):
+    """Parse a tensor network where tensors are assumed to be variables unless
+    tagged.
+    """
+    tn_ag = tn.copy()
+    variables = []
+
+    for t in tn_ag:
+        if t.tags & constant_tags:
+            t.modify(apply=to_constant)
+            continue
+
+        # append the raw data but mark the corresponding tensor
+        # for reinsertion
+        data = t.get_params()
+        variables.append(data)
+        t.add_tag(_VARIABLE_TAG.format(len(variables) - 1))
+
+    return tn_ag, variables
+
+
+def _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.
+    """
+    variables = []
+
+    def collect(x):
+        if hasattr(x, "get_params"):
+
+            if hasattr(x, "apply_to_arrays"):
+                x.apply_to_arrays(to_constant)
+
+            # variables can be a pytree
+            variables.append(x.get_params())
+            return x.copy()
+        else:
+            # raw array
+            variables.append(x)
+            return None
+
+    ref = tree_map(collect, x)
+    return ref, variables
+
+
+def 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``.
+    """
+    tags = tags_to_oset(tags)
+    shared_tags = tags_to_oset(shared_tags)
+    constant_tags = tags_to_oset(constant_tags)
+
+    if not isinstance(tn, TensorNetwork):
+        if any((tags, shared_tags, constant_tags)):
+            raise ValueError(
+                "TNOptimizer error, if `tags`, `shared_tags`, or"
+                " `constant_tags` are specified then `tn` must be a"
+                " TensorNetwork, rather than a general pytree."
+            )
+        return _parse_pytree_to_backend(tn, to_constant=to_constant)
+
+    if tags | shared_tags:
+        # opt_in
+        if not (tags & shared_tags) == shared_tags:
+            tags = tags | shared_tags
+            warnings.warn(
+                "TNOptimizer warning, some `shared_tags` are missing"
+                " from `tags`. Automatically adding these missing"
+                " `shared_tags` to `tags`."
+            )
+        if constant_tags:
+            warnings.warn(
+                "TNOptimizer warning, if `tags` or `shared_tags` are"
+                " specified then `constant_tags` is ignored - "
+                "consider instead untagging those tensors."
+            )
+        return _parse_opt_in(
+            tn,
+            tags,
+            shared_tags,
+            to_constant,
+        )
+
+    # opt-out
+    return _parse_opt_out(
+        tn,
+        constant_tags,
+        to_constant,
+    )
+
+
+def _inject_variables_pytree(arrays, tree):
+    arrays = iter(arrays)
+
+    def inject(x):
+        if hasattr(x, "set_params"):
+            x.set_params(next(arrays))
+            return x
+        else:
+            return next(arrays)
+
+    return tree_map(inject, tree)
+
+
+def inject_variables(arrays, tn):
+    """Given the list of optimized variables ``arrays`` and the target tensor
+    network or pytree ``tn``, inject the variables back in.
+    """
+    if not isinstance(tn, TensorNetwork):
+        return _inject_variables_pytree(arrays, tn)
+
+    tn = tn.copy()
+    for t in tn:
+        for tag in t.tags:
+            match = variable_finder.match(tag)
+            if match is not None:
+                i = int(match.groups(1)[0])
+                t.set_params(arrays[i])
+                break
+    return tn
+
+
+def 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.
+    """
+    try:
+        # Tensor, TensorNetwork...
+        x = x.copy()
+        x.apply_to_arrays(f)
+        return x
+    except AttributeError:
+        pass
+
+    try:
+        # raw structured arrays that provide the {get|set}_params interface
+        x = x.copy()
+        x.set_params(tree_map(f, x.get_params()))
+        return x
+    except AttributeError:
+        pass
+
+    # other raw arrays
+    return f(x)
+
+
+def convert_variables_to_numpy(x):
+    if hasattr(x, "apply_to_arrays"):
+        x.apply_to_arrays(to_numpy)
+        return x
+    elif hasattr(x, "get_params"):
+        old_params = x.get_params()
+        new_params = tree_map(to_numpy, old_params)
+        x.set_params(new_params)
+        return x
+    else:
+        return to_numpy(x)
+
+
+@functools.lru_cache(1)
+def get_autograd():
+    import autograd
+
+    return autograd
+
+
+class AutoGradHandler:
+    def __init__(self, device="cpu"):
+        if device != "cpu":
+            raise ValueError(
+                "`autograd` currently is only " "backed by cpu, numpy arrays."
+            )
+
+    def to_variable(self, x):
+        return np.asarray(x)
+
+    def to_constant(self, x):
+        return np.asarray(x)
+
+    def setup_fn(self, fn):
+        autograd = get_autograd()
+        self._backend_fn = fn
+        self._value_and_grad = autograd.value_and_grad(fn)
+        self._hvp = autograd.hessian_vector_product(fn)
+
+    def value(self, arrays):
+        return self._backend_fn(arrays)
+
+    def value_and_grad(self, arrays):
+        loss, grads = self._value_and_grad(arrays)
+        return loss, tree_map(lambda x: x.conj(), grads)
+
+
+class JaxHandler:
+    def __init__(self, jit_fn=True, device=None):
+        self.jit_fn = jit_fn
+        self.device = device
+
+    def to_variable(self, x):
+        jax = get_jax()
+        return jax.numpy.asarray(x)
+
+    def to_constant(self, x):
+        jax = get_jax()
+        return jax.numpy.asarray(x)
+
+    def setup_fn(self, fn):
+        jax = get_jax()
+        if self.jit_fn:
+            self._backend_fn = jax.jit(fn, backend=self.device)
+            self._value_and_grad = jax.jit(
+                jax.value_and_grad(fn), backend=self.device
+            )
+        else:
+            self._backend_fn = fn
+            self._value_and_grad = jax.value_and_grad(fn)
+
+        self._setup_hessp(fn)
+
+    def _setup_hessp(self, fn):
+        jax = get_jax()
+
+        def hvp(primals, tangents):
+            return jax.jvp(jax.grad(fn), (primals,), (tangents,))[1]
+
+        if self.jit_fn:
+            hvp = jax.jit(hvp, device=self.device)
+
+        self._hvp = hvp
+
+    def value(self, arrays):
+        jax_arrays = tree_map(self.to_constant, arrays)
+        return to_numpy(self._backend_fn(jax_arrays))
+
+    def value_and_grad(self, arrays):
+        loss, grads = self._value_and_grad(arrays)
+        return (
+            loss,
+            tree_map(lambda x: to_numpy(x.conj()), grads),
+        )
+
+    def hessp(self, primals, tangents):
+        jax_arrays = self._hvp(primals, tangents)
+        return tree_map(to_numpy, jax_arrays)
+
+
+@functools.lru_cache(1)
+def get_tensorflow():
+    import tensorflow
+
+    return tensorflow
+
+
+class TensorFlowHandler:
+    def __init__(
+        self,
+        jit_fn=False,
+        autograph=False,
+        experimental_compile=False,
+        device=None,
+    ):
+        self.jit_fn = jit_fn
+        self.autograph = autograph
+        self.experimental_compile = experimental_compile
+        self.device = device
+
+    def to_variable(self, x):
+        tf = get_tensorflow()
+        if self.device is None:
+            return tf.Variable(x)
+        with tf.device(self.device):
+            return tf.Variable(x)
+
+    def to_constant(self, x):
+        tf = get_tensorflow()
+        if self.device is None:
+            return tf.constant(x)
+        with tf.device(self.device):
+            return tf.constant(x)
+
+    def setup_fn(self, fn):
+        tf = get_tensorflow()
+        if self.jit_fn:
+            self._backend_fn = tf.function(
+                fn,
+                autograph=self.autograph,
+                experimental_compile=self.experimental_compile,
+            )
+        else:
+            self._backend_fn = fn
+
+    def value(self, arrays):
+        tf_arrays = tree_map(self.to_constant, arrays)
+        return to_numpy(self._backend_fn(tf_arrays))
+
+    def value_and_grad(self, arrays):
+        tf = get_tensorflow()
+
+        variables = tree_map(self.to_variable, arrays)
+
+        with tf.GradientTape() as t:
+            result = self._backend_fn(variables)
+
+        tf_grads = t.gradient(
+            result,
+            variables,
+            # want to return zeros for unconnected gradients
+            unconnected_gradients=tf.UnconnectedGradients.ZERO,
+        )
+
+        grads = tree_map(to_numpy, tf_grads)
+        loss = to_numpy(result)
+        return loss, grads
+
+
+@functools.lru_cache(1)
+def get_torch():
+    import torch
+
+    return torch
+
+
+class TorchHandler:
+    def __init__(self, jit_fn=False, device=None):
+        torch = get_torch()
+        self.jit_fn = jit_fn
+        if device is None:
+            device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.device = device
+
+    def to_variable(self, x):
+        torch = get_torch()
+        return torch.tensor(x).to(self.device).requires_grad_()
+
+    def to_constant(self, x):
+        torch = get_torch()
+        return torch.tensor(x).to(self.device)
+
+    def setup_fn(self, fn):
+        self._fn = fn
+        self._backend_fn = None
+
+    def _setup_backend_fn(self, arrays):
+        torch = get_torch()
+        if self.jit_fn:
+            example_inputs = (tree_map(self.to_variable, arrays),)
+            with warnings.catch_warnings():
+                warnings.filterwarnings(
+                    action="ignore",
+                    message=".*can't record the data flow of Python values.*",
+                )
+                self._backend_fn = torch.jit.trace(
+                    self._fn, example_inputs=example_inputs
+                )
+        else:
+            self._backend_fn = self._fn
+
+    def value(self, arrays):
+        if self._backend_fn is None:
+            self._setup_backend_fn(arrays)
+        torch_arrays = tree_map(self.to_constant, arrays)
+        return to_numpy(self._backend_fn(torch_arrays))
+
+    def value_and_grad(self, arrays):
+        if self._backend_fn is None:
+            self._setup_backend_fn(arrays)
+
+        variables = tree_map(self.to_variable, arrays)
+        result = self._backend_fn(variables)
+
+        def get_gradient_from_torch(t):
+            if t.grad is None:
+                return np.zeros(t.shape, dtype=get_dtype_name(t))
+            return to_numpy(t.grad).conj()
+
+        result.backward()
+        grads = tree_map(get_gradient_from_torch, variables)
+
+        loss = to_numpy(result)
+        return loss, grads
+
+
+_BACKEND_HANDLERS = {
+    "numpy": AutoGradHandler,
+    "autograd": AutoGradHandler,
+    "jax": JaxHandler,
+    "tensorflow": TensorFlowHandler,
+    "torch": TorchHandler,
+}
+
+
+class MultiLossHandler:
+    def __init__(self, autodiff_backend, executor=None, **backend_opts):
+        self.autodiff_backend = autodiff_backend
+        self.backend_opts = backend_opts
+        self.executor = executor
+
+        # start just with one, as we don't don't know how many functions yet
+        h0 = _BACKEND_HANDLERS[autodiff_backend](**backend_opts)
+        self.handlers = [h0]
+        # ... but we do need access to `to_constant`
+        self.to_constant = h0.to_constant
+
+    def setup_fn(self, funcs):
+        fn0, *fns = funcs
+        self.handlers[0].setup_fn(fn0)
+        for fn in fns:
+            h = _BACKEND_HANDLERS[self.autodiff_backend](**self.backend_opts)
+            h.setup_fn(fn)
+            self.handlers.append(h)
+
+    def _value_seq(self, arrays):
+        return sum(h.value(arrays) for h in self.handlers)
+
+    def _value_par_seq(self, arrays):
+        futures = [
+            self.executor.submit(h.value, arrays) for h in self.handlers
+        ]
+        return sum(f.result() for f in futures)
+
+    def value(self, arrays):
+        if self.executor is not None:
+            return self._value_par(arrays)
+        return self._value_seq(arrays)
+
+    def _value_and_grad_seq(self, arrays):
+        h0, *hs = self.handlers
+        loss, grads = h0.value_and_grad(arrays)
+        # need to make arrays writeable for efficient inplace sum
+        grads = list(map(np.array, grads))
+        for h in hs:
+            loss_i, grads_i = h.value_and_grad(arrays)
+            loss = loss + loss_i
+            for i, g_i in enumerate(grads_i):
+                grads[i] += g_i
+        return loss, grads
+
+    def _value_and_grad_par(self, arrays):
+        futures = [
+            self.executor.submit(h.value_and_grad, arrays)
+            for h in self.handlers
+        ]
+        results = (f.result() for f in futures)
+
+        # get first result
+        loss, grads = next(results)
+        grads = list(map(np.array, grads))
+
+        # process remaining results
+        for loss_i, grads_i in results:
+            loss = loss + loss_i
+            for i, g_i in enumerate(grads_i):
+                grads[i] += g_i
+
+        return loss, grads
+
+    def value_and_grad(self, arrays):
+        if self.executor is not None:
+            return self._value_and_grad_par(arrays)
+        return self._value_and_grad_seq(arrays)
+
+
+class SGD:
+    """Stateful ``scipy.optimize.minimize`` compatible implementation of
+    stochastic gradient descent with momentum.
+
+    Adapted from ``autograd/misc/optimizers.py``.
+    """
+
+    def __init__(self):
+        from scipy.optimize import OptimizeResult
+
+        self.OptimizeResult = OptimizeResult
+        self._i = 0
+        self._velocity = None
+
+    def get_velocity(self, x):
+        if self._velocity is None:
+            self._velocity = np.zeros_like(x)
+        return self._velocity
+
+    def __call__(
+        self,
+        fun,
+        x0,
+        jac,
+        args=(),
+        learning_rate=0.1,
+        mass=0.9,
+        maxiter=1000,
+        callback=None,
+        bounds=None,
+        **kwargs,
+    ):
+        x = x0
+        velocity = self.get_velocity(x)
+
+        for _ in range(maxiter):
+            self._i += 1
+
+            g = jac(x)
+
+            if callback and callback(x):
+                break
+
+            velocity = mass * velocity - (1.0 - mass) * g
+            x = x + learning_rate * velocity
+
+            if bounds is not None:
+                x = np.clip(x, bounds[:, 0], bounds[:, 1])
+
+        # save for restart
+        self._velocity = velocity
+
+        return self.OptimizeResult(
+            x=x, fun=fun(x), jac=g, nit=self._i, nfev=self._i, success=True
+        )
+
+
+class RMSPROP:
+    """Stateful ``scipy.optimize.minimize`` compatible implementation of
+    root mean squared prop: See Adagrad paper for details.
+
+    Adapted from ``autograd/misc/optimizers.py``.
+    """
+
+    def __init__(self):
+        from scipy.optimize import OptimizeResult
+
+        self.OptimizeResult = OptimizeResult
+        self._i = 0
+        self._avg_sq_grad = None
+
+    def get_avg_sq_grad(self, x):
+        if self._avg_sq_grad is None:
+            self._avg_sq_grad = np.ones_like(x)
+        return self._avg_sq_grad
+
+    def __call__(
+        self,
+        fun,
+        x0,
+        jac,
+        args=(),
+        learning_rate=0.1,
+        gamma=0.9,
+        eps=1e-8,
+        maxiter=1000,
+        callback=None,
+        bounds=None,
+        **kwargs,
+    ):
+        x = x0
+        avg_sq_grad = self.get_avg_sq_grad(x)
+
+        for _ in range(maxiter):
+            self._i += 1
+
+            g = jac(x)
+
+            if callback and callback(x):
+                break
+
+            avg_sq_grad = avg_sq_grad * gamma + g**2 * (1 - gamma)
+            x = x - learning_rate * g / (np.sqrt(avg_sq_grad) + eps)
+
+            if bounds is not None:
+                x = np.clip(x, bounds[:, 0], bounds[:, 1])
+
+        # save for restart
+        self._avg_sq_grad = avg_sq_grad
+
+        return self.OptimizeResult(
+            x=x, fun=fun(x), jac=g, nit=self._i, nfev=self._i, success=True
+        )
+
+
+class ADAM:
+    """Stateful ``scipy.optimize.minimize`` compatible implementation of
+    ADAM - http://arxiv.org/pdf/1412.6980.pdf.
+
+    Adapted from ``autograd/misc/optimizers.py``.
+    """
+
+    def __init__(self):
+        from scipy.optimize import OptimizeResult
+
+        self.OptimizeResult = OptimizeResult
+        self._i = 0
+        self._m = None
+        self._v = None
+
+    def get_m(self, x):
+        if self._m is None:
+            self._m = np.zeros_like(x)
+        return self._m
+
+    def get_v(self, x):
+        if self._v is None:
+            self._v = np.zeros_like(x)
+        return self._v
+
+    def __call__(
+        self,
+        fun,
+        x0,
+        jac,
+        args=(),
+        learning_rate=0.001,
+        beta1=0.9,
+        beta2=0.999,
+        eps=1e-8,
+        maxiter=1000,
+        callback=None,
+        bounds=None,
+        **kwargs,
+    ):
+        x = x0
+        m = self.get_m(x)
+        v = self.get_v(x)
+
+        for _ in range(maxiter):
+            self._i += 1
+
+            g = jac(x)
+
+            if callback and callback(x):
+                break
+
+            m = (1 - beta1) * g + beta1 * m  # first  moment estimate.
+            v = (1 - beta2) * (g**2) + beta2 * v  # second moment estimate.
+            mhat = m / (1 - beta1**self._i)  # bias correction.
+            vhat = v / (1 - beta2**self._i)
+            x = x - learning_rate * mhat / (np.sqrt(vhat) + eps)
+
+            if bounds is not None:
+                x = np.clip(x, bounds[:, 0], bounds[:, 1])
+
+        # save for restart
+        self._m = m
+        self._v = v
+
+        return self.OptimizeResult(
+            x=x, fun=fun(x), jac=g, nit=self._i, nfev=self._i, success=True
+        )
+
+
+class NADAM:
+    """Stateful ``scipy.optimize.minimize`` compatible implementation of
+    NADAM - [Dozat - http://cs229.stanford.edu/proj2015/054_report.pdf].
+
+    Adapted from ``autograd/misc/optimizers.py``.
+    """
+
+    def __init__(self):
+        from scipy.optimize import OptimizeResult
+
+        self.OptimizeResult = OptimizeResult
+        self._i = 0
+        self._m = None
+        self._v = None
+        self._mus = None
+
+    def get_m(self, x):
+        if self._m is None:
+            self._m = np.zeros_like(x)
+        return self._m
+
+    def get_v(self, x):
+        if self._v is None:
+            self._v = np.zeros_like(x)
+        return self._v
+
+    def get_mus(self, beta1):
+        if self._mus is None:
+            self._mus = [1, beta1 * (1 - 0.5 * 0.96**0.004)]
+        return self._mus
+
+    def __call__(
+        self,
+        fun,
+        x0,
+        jac,
+        args=(),
+        learning_rate=0.001,
+        beta1=0.9,
+        beta2=0.999,
+        eps=1e-8,
+        maxiter=1000,
+        callback=None,
+        bounds=None,
+        **kwargs,
+    ):
+        x = x0
+        m = self.get_m(x)
+        v = self.get_v(x)
+        mus = self.get_mus(beta1)
+
+        for _ in range(maxiter):
+            self._i += 1
+
+            # this is ``mu[t + 1]`` -> already computed ``mu[t]``
+            self._mus.append(
+                beta1 * (1 - 0.5 * 0.96 ** (0.004 * (self._i + 1)))
+            )
+
+            g = jac(x)
+
+            if callback and callback(x):
+                break
+
+            gd = g / (1 - np.prod(self._mus[:-1]))
+            m = beta1 * m + (1 - beta1) * g
+            md = m / (1 - np.prod(self._mus))
+            v = beta2 * v + (1 - beta2) * g**2
+            vd = v / (1 - beta2**self._i)
+            mhat = (1 - self._mus[self._i]) * gd + self._mus[self._i + 1] * md
+
+            x = x - learning_rate * mhat / (np.sqrt(vd) + eps)
+
+            if bounds is not None:
+                x = np.clip(x, bounds[:, 0], bounds[:, 1])
+
+        # save for restart
+        self._m = m
+        self._v = v
+        self._mus = mus
+
+        return self.OptimizeResult(
+            x=x, fun=fun(x), jac=g, nit=self._i, nfev=self._i, success=True
+        )
+
+
+class ADABELIEF:
+    """Stateful ``scipy.optimize.minimize`` compatible implementation of
+    ADABELIEF - https://arxiv.org/abs/2010.07468.
+
+    Adapted from ``autograd/misc/optimizers.py``.
+    """
+
+    def __init__(self):
+        from scipy.optimize import OptimizeResult
+
+        self.OptimizeResult = OptimizeResult
+        self._i = 0
+        self._m = None
+        self._s = None
+
+    def get_m(self, x):
+        if self._m is None:
+            self._m = np.zeros_like(x)
+        return self._m
+
+    def get_s(self, x):
+        if self._s is None:
+            self._s = np.zeros_like(x)
+        return self._s
+
+    def __call__(
+        self,
+        fun,
+        x0,
+        jac,
+        args=(),
+        learning_rate=0.001,
+        beta1=0.9,
+        beta2=0.999,
+        eps=1e-8,
+        maxiter=1000,
+        callback=None,
+        bounds=None,
+        **kwargs,
+    ):
+        x = x0
+        m = self.get_m(x)
+        s = self.get_s(x)
+
+        for _ in range(maxiter):
+            self._i += 1
+
+            g = jac(x)
+
+            if callback and callback(x):
+                break
+
+            m = (1 - beta1) * g + beta1 * m
+            s = (1 - beta2) * (g - m) ** 2 + beta2 * s + eps
+            # bias correction
+            mhat = m / (1 - beta1**self._i)
+            shat = s / (1 - beta2**self._i)
+            x = x - learning_rate * mhat / (np.sqrt(shat) + eps)
+
+            if bounds is not None:
+                x = np.clip(x, bounds[:, 0], bounds[:, 1])
+
+        # save for restart
+        self._m = m
+        self._s = s
+
+        return self.OptimizeResult(
+            x=x, fun=fun(x), jac=g, nit=self._i, nfev=self._i, success=True
+        )
+
+
+_STOC_GRAD_METHODS = {
+    "sgd": SGD,
+    "rmsprop": RMSPROP,
+    "adam": ADAM,
+    "nadam": NADAM,
+    "adabelief": ADABELIEF,
+}
+
+
+class MakeArrayFn:
+    """Class wrapper so picklable."""
+
+    __name__ = "MakeArrayFn"
+
+    def __init__(self, tn_opt, loss_fn, norm_fn, autodiff_backend):
+        self.tn_opt = tn_opt
+        self.loss_fn = loss_fn
+        self.norm_fn = norm_fn
+        self.autodiff_backend = autodiff_backend
+
+    def __call__(self, arrays):
+        tn_compute = inject_variables(arrays, self.tn_opt)
+
+        # set backend explicitly as maybe mixing with numpy arrays
+        with contract_backend(self.autodiff_backend):
+            return self.loss_fn(self.norm_fn(tn_compute))
+
+
+def identity_fn(x):
+    return x
+
+
+class TNOptimizer:
+    """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'`` .
+    """
+
+    def __init__(
+        self,
+        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,
+    ):
+        self.progbar = progbar
+        self.tags = tags
+        self.shared_tags = shared_tags
+        self.constant_tags = constant_tags
+
+        if autodiff_backend.upper() == "AUTO":
+            autodiff_backend = _DEFAULT_BACKEND
+        self._autodiff_backend = autodiff_backend
+        self._multiloss = isinstance(loss_fn, Iterable)
+
+        # the object that handles converting to backend + computing gradient
+        if self._multiloss:
+            # special meta-handler if loss function is sequence to sum
+            backend_opts["executor"] = executor
+            self.handler = MultiLossHandler(autodiff_backend, **backend_opts)
+        else:
+            self.handler = _BACKEND_HANDLERS[autodiff_backend](**backend_opts)
+
+        # use identity if no nomalization required
+        if norm_fn is None:
+            norm_fn = identity_fn
+        self.norm_fn = norm_fn
+
+        self.reset(tn, loss_target=loss_target)
+
+        # convert constant raw arrays ahead of time to correct backend
+        self.loss_constants = tree_map(
+            functools.partial(convert_raw_arrays, f=self.handler.to_constant),
+            ensure_dict(loss_constants),
+        )
+        self.loss_kwargs = ensure_dict(loss_kwargs)
+        kws = {**self.loss_constants, **self.loss_kwargs}
+
+        # inject these constant options to the loss function(s)
+        if self._multiloss:
+            # loss is a sum of independent terms
+            self.loss_fn = [functools.partial(fn, **kws) for fn in loss_fn]
+        else:
+            # loss is all in one
+            self.loss_fn = functools.partial(loss_fn, **kws)
+
+        # first we wrap the function to convert from array args to TN arg
+        #     (i.e. to autodiff library compatible form)
+        if self._multiloss:
+            array_fn = [
+                MakeArrayFn(self._tn_opt, fn, self.norm_fn, autodiff_backend)
+                for fn in self.loss_fn
+            ]
+        else:
+            array_fn = MakeArrayFn(
+                self._tn_opt, self.loss_fn, self.norm_fn, autodiff_backend
+            )
+
+        # then we pass it to the handler which generates a function that
+        # computes both the value and gradients (still in array form)
+        self.handler.setup_fn(array_fn)
+
+        # options to do with the minimizer
+        self.bounds = bounds
+        self.optimizer = optimizer
+        self.callback = callback
+
+    def _set_tn(self, tn):
+        # work out which tensors to optimize and get the underlying data
+        self._tn_opt, variables = parse_network_to_backend(
+            tn,
+            tags=self.tags,
+            shared_tags=self.shared_tags,
+            constant_tags=self.constant_tags,
+            to_constant=self.handler.to_constant,
+        )
+        # handles storing and packing / unpacking many arrays as a vector
+        self.vectorizer = Vectorizer(variables)
+
+    def _reset_tracking_info(self, loss_target=None):
+        # tracking info
+        self.loss = float("inf")
+        self.loss_best = float("inf")
+        self.loss_target = loss_target
+        self.losses = []
+        self._n = 0
+        self._pbar = None
+
+    def reset(self, 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.
+        """
+        if tn is not None:
+            self._set_tn(tn)
+        if clear_info:
+            self._reset_tracking_info(loss_target=loss_target)
+
+    def _maybe_init_pbar(self, n):
+        if self.progbar:
+            self._pbar = tqdm.tqdm(total=n)
+
+    def _maybe_update_pbar(self):
+        if self._pbar is not None:
+            self._pbar.update()
+            self.loss_best = min(self.loss_best, self.loss)
+            msg = f"{self.loss:+.12f} [best: {self.loss_best:+.12f}] "
+            self._pbar.set_description(msg)
+
+    def _maybe_close_pbar(self):
+        if self._pbar is not None:
+            self._pbar.close()
+            self._pbar = None
+
+    def _check_loss_target(self):
+        if (self.loss_target is not None) and (self.loss <= self.loss_target):
+            # for scipy terminating optimizer with callback doesn't work
+            raise KeyboardInterrupt
+
+    def _maybe_call_callback(self):
+        if self.callback is not None:
+            self.callback(self)
+
+    def vectorized_value(self, x):
+        """The value of the loss function at vector ``x``."""
+        self.vectorizer.vector[:] = x
+        arrays = self.vectorizer.unpack()
+        self.loss = self.handler.value(arrays).item()
+        self.losses.append(self.loss)
+        self._n += 1
+        self._maybe_update_pbar()
+        self._check_loss_target()
+        self._maybe_call_callback()
+        return self.loss
+
+    def vectorized_value_and_grad(self, x):
+        """The value and gradient of the loss function at vector ``x``."""
+        self.vectorizer.vector[:] = x
+        arrays = self.vectorizer.unpack()
+        result, grads = self.handler.value_and_grad(arrays)
+        self._n += 1
+        self.loss = result.item()
+        self.losses.append(self.loss)
+        vec_grad = self.vectorizer.pack(grads, "grad")
+        self._maybe_update_pbar()
+        self._check_loss_target()
+        self._maybe_call_callback()
+        return self.loss, vec_grad
+
+    def vectorized_hessp(self, x, p):
+        """The action of the hessian at point ``x`` on vector ``p``."""
+        primals = self.vectorizer.unpack(x)
+        tangents = self.vectorizer.unpack(p)
+        hp_arrays = self.handler.hessp(primals, tangents)
+        self._n += 1
+        self.losses.append(self.loss)
+        self._maybe_update_pbar()
+        return self.vectorizer.pack(hp_arrays, "hp")
+
+    def __repr__(self):
+        return (
+            f"<TNOptimizer(d={self.d}, " f"backend={self._autodiff_backend})>"
+        )
+
+    @property
+    def d(self):
+        return int(self.vectorizer.d)
+
+    @property
+    def nevals(self):
+        """The number of gradient evaluations."""
+        return self._n
+
+    @property
+    def optimizer(self):
+        """The underlying optimizer that works with the vectorized functions."""
+        return self._optimizer
+
+    @optimizer.setter
+    def optimizer(self, x):
+        self._optimizer = x
+        if self.optimizer in _STOC_GRAD_METHODS:
+            self._method = _STOC_GRAD_METHODS[self.optimizer]()
+        else:
+            self._method = self.optimizer
+
+    @property
+    def bounds(self):
+        return self._bounds
+
+    @bounds.setter
+    def bounds(self, x):
+        if x is not None:
+            self._bounds = np.array((x,) * self.vectorizer.d)
+        else:
+            self._bounds = None
+
+    def get_tn_opt(self):
+        """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 : TensorNetwork
+        """
+        arrays = tree_map(
+            self.handler.to_constant,
+            self.vectorizer.unpack(),
+        )
+        tn = inject_variables(arrays, self._tn_opt)
+        tn = self.norm_fn(tn)
+
+        if isinstance(tn, TensorNetwork):
+            tn.drop_tags(t for t in tn.tags if variable_finder.match(t))
+
+        return tree_map(convert_variables_to_numpy, tn)
+
+    def optimize(
+        self, n, tol=None, jac=True, hessp=False, optlib="scipy", **options
+    ):
+        """Run the optimizer for ``n`` function evaluations, using by default
+        :func:`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 :func:`scipy.optimize.minimize` or whichever optimizer
+            is being used.
+
+        Returns
+        -------
+        tn_opt : TensorNetwork
+        """
+        return {
+            "scipy": self.optimize_scipy,
+            "nlopt": self.optimize_nlopt,
+        }[optlib](n=n, tol=tol, jac=jac, hessp=hessp, **options)
+
+    def optimize_scipy(self, n, tol=None, jac=True, hessp=False, **options):
+        """Scipy based optimization, see
+        :meth:`~quimb.tensor.optimize.TNOptimizer.optimize` for details.
+        """
+        from scipy.optimize import minimize
+
+        if jac:
+            fun = self.vectorized_value_and_grad
+        else:
+            fun = self.vectorized_value
+
+        if self._method in ("l-bfgs-b", "tnc"):
+            options.setdefault("maxfun", n)
+
+        try:
+            self._maybe_init_pbar(n)
+            self.res = minimize(
+                fun=fun,
+                jac=jac,
+                hessp=self.vectorized_hessp if hessp else None,
+                x0=self.vectorizer.vector,
+                tol=tol,
+                bounds=self.bounds,
+                method=self._method,
+                options=dict(maxiter=n, **options),
+            )
+            self.vectorizer.vector[:] = self.res.x
+        except KeyboardInterrupt:
+            pass
+        finally:
+            self._maybe_close_pbar()
+
+        return self.get_tn_opt()
+
+    def optimize_basinhopping(
+        self, n, nhop, temperature=1.0, jac=True, hessp=False, **options
+    ):
+        """Run the optimizer for using :func:`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 :func:`scipy.optimize.minimize` call.
+
+        Returns
+        -------
+        tn_opt : TensorNetwork
+        """
+        from scipy.optimize import basinhopping
+
+        if jac:
+            fun = self.vectorized_value_and_grad
+        else:
+            fun = self.vectorized_value
+
+        try:
+            self._maybe_init_pbar(n * nhop)
+            self.res = basinhopping(
+                func=fun,
+                x0=self.vectorizer.vector,
+                niter=nhop,
+                minimizer_kwargs=dict(
+                    jac=jac,
+                    hessp=self.vectorized_hessp if hessp else None,
+                    method=self._method,
+                    bounds=self.bounds,
+                    options=dict(maxiter=n, **options),
+                ),
+                T=temperature,
+            )
+            self.vectorizer.vector[:] = self.res.x
+
+        except KeyboardInterrupt:
+            pass
+        finally:
+            self._maybe_close_pbar()
+
+        return self.get_tn_opt()
+
+    def optimize_nlopt(
+        self,
+        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 : TensorNetwork
+        """
+        import nlopt
+
+        if not (jac is True and hessp is False):
+            raise NotImplementedError(
+                "Only gradient based optimizers are "
+                "supported with nlopt currently."
+            )
+
+        if tol == 0.0:
+            # assume no stopping criteria
+            ftol_rel == 0.0 if ftol_rel is None else ftol_rel
+            ftol_abs == 0.0 if ftol_abs is None else ftol_abs
+            xtol_rel == 0.0 if xtol_rel is None else xtol_rel
+            xtol_abs == 0.0 if xtol_abs is None else xtol_abs
+
+        elif (tol is not None) and (ftol_rel is None):
+            # assume relative loss tolerance is specified
+            ftol_rel = tol
+
+        # translate directly comparable algorithms
+        optimizer = {
+            "l-bfgs-b": "LD_LBFGS",
+            "slsqp": "LD_SLSQP",
+            "tnc": "LD_TNEWTON_PRECOND_RESTART",
+            "cobyla": "LN_COBYLA",
+        }.get(self.optimizer.lower(), self.optimizer)
+
+        try:
+            self._maybe_init_pbar(n)
+
+            def f(x, grad):
+                self.vectorizer.vector[:] = x
+                arrays = self.vectorizer.unpack()
+                if grad.size > 0:
+                    result, grads = self.handler.value_and_grad(arrays)
+                    grad[:] = self.vectorizer.pack(grads, "grad")
+                else:
+                    result = self.handler.value(arrays)
+                self._n += 1
+                self.loss = result.item()
+                self.losses.append(self.loss)
+                self._maybe_update_pbar()
+                return self.loss
+
+            opt = nlopt.opt(getattr(nlopt, optimizer), self.d)
+            opt.set_min_objective(f)
+            opt.set_maxeval(n)
+
+            if self.bounds is not None:
+                opt.set_lower_bounds(self.bounds[:, 0])
+                opt.set_upper_bounds(self.bounds[:, 1])
+
+            if self.loss_target is not None:
+                opt.set_stopval(self.loss_target)
+            if ftol_rel is not None:
+                opt.set_ftol_rel(ftol_rel)
+            if ftol_abs is not None:
+                opt.set_ftol_abs(ftol_abs)
+            if xtol_rel is not None:
+                opt.set_xtol_rel(xtol_rel)
+            if xtol_abs is not None:
+                opt.set_xtol_abs(xtol_abs)
+
+            self.vectorizer.vector[:] = opt.optimize(self.vectorizer.vector)
+
+        except (KeyboardInterrupt, RuntimeError):
+            pass
+        finally:
+            self._maybe_close_pbar()
+
+        return self.get_tn_opt()
+
+    def optimize_ipopt(self, 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 : TensorNetwork
+        """
+        from cyipopt import minimize_ipopt
+
+        try:
+            self._maybe_init_pbar(n)
+            self.res = minimize_ipopt(
+                fun=self.vectorized_value_and_grad,
+                jac=True,
+                x0=self.vectorizer.vector,
+                tol=tol,
+                bounds=self.bounds,
+                method=self._method,
+                options=dict(maxiter=n, **options),
+            )
+            self.vectorizer.vector[:] = self.res.x
+        except KeyboardInterrupt:
+            pass
+        finally:
+            self._maybe_close_pbar()
+
+        return self.get_tn_opt()
+
+    def optimize_nevergrad(self, 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 : TensorNetwork
+        """
+        import nevergrad as ng
+
+        opt = getattr(ng.optimizers, self.optimizer)(
+            parametrization=ng.p.Array(
+                init=self.vectorizer.vector,
+                lower=self.bounds[:, 0] if self.bounds is not None else None,
+                upper=self.bounds[:, 1] if self.bounds is not None else None,
+            ),
+            budget=n,
+        )
+
+        try:
+            self._maybe_init_pbar(n)
+            for _ in range(n):
+                x = opt.ask()
+                loss = self.vectorized_value(*x.args, **x.kwargs)
+                opt.tell(x, loss)
+                if self.loss_target is not None:
+                    if self.loss < self.loss_target:
+                        break
+
+        except KeyboardInterrupt:
+            pass
+        finally:
+            self._maybe_close_pbar()
+
+        # recommendation = opt.minimize(self.vectorized_value)
+        recommendation = opt.provide_recommendation()
+        self.vectorizer.vector[:] = recommendation.value
+
+        return self.get_tn_opt()
+
+    @default_to_neutral_style
+    def plot(
+        self,
+        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.
+        """
+        import matplotlib.pyplot as plt
+        from matplotlib.colors import hsv_to_rgb
+
+        ys = np.array(self.losses)
+        xs = np.arange(ys.size)
+
+        fig, ax = plt.subplots()
+        ax.plot(xs, ys, ".-")
+        if xscale == "symlog":
+            ax.set_xscale(xscale, linthresh=xscale_linthresh)
+            ax.axvline(xscale_linthresh, color=(.5, .5, .5), ls="-", lw=0.5)
+        else:
+            ax.set_xscale(xscale)
+        ax.set_xlabel("Iteration")
+        ax.set_ylabel("Loss")
+
+        if hlines:
+            hlines = dict(hlines)
+            for i, (label, value) in enumerate(hlines.items()):
+                color = hsv_to_rgb([(0.1 * i) % 1.0, 0.9, 0.9])
+                ax.axhline(value, color=color, ls="--", label=label)
+                ax.text(1, value, label, color=color, va="bottom", ha="left")
+
+        if zoom is not None:
+            if zoom == "auto":
+                zoom = min(50, ys.size // 2)
+
+            iax = ax.inset_axes([0.5, 0.5, 0.5, 0.5])
+            iax.plot(xs[-zoom:], ys[-zoom:], ".-")
+
+        return fig, ax