PRobability Optimizer with TEnsor Sampling

This is an implementation of PROTES probabilistic optimization algorithm [1] from scratch. Its original implementation is available at anabatsh/protes and published implementation is at anabatsh/TT_pro. The package can be directly installed from GitHub as follows.

pip install git+https://github.com/daskol/protes.git

As an example, let's solve a toy problem. Consider a minimization problem of summation 10 non-negative numbers, not greater than 5. Then this problem can be easily solved with PROTES.

from protes import minimize

(x, y), pdf = minimize(
    fn=lambda ix: ix.sum(axis=1),
    shape=(5, ) * 10,
    max_trials=5000,
)

print('x =', x)  # x = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
print('y =', y)  # y = 0

All objects provided protes are actual pytree's and could be used with jax.jit. So, it is possible to pickle and unpickle them. With the snippet above, one can easily save density estimator pdf to and load it from a file.

from pickle import dump, load

with open('pdf.pkl', 'wb') as fout:
    dump(pdf, fout)

with open('pdf.pkl', 'rb') as fin:
    pdf = load(fin)

key = jax.random.PRNGKey(42)
batch = pdf.sample(key, (2, 1))
assert batch.ndim == 3
assert batch.shape = (2, 1, 10)

Also, the package provides an implementation of PROTES as a stateful object TensorTrainSampler for optimization in the manner of active learning. Again, everything can be jax.jit'ed.

from functools import partial
from protes import TensorTrainSampler
import jax


@partial(jax.jit, static_argnames=('fn', ))
def step(key, sampler, fn):
    indices = sampler.sample(key, (2, ))
    values = fn(indices)
    sampler, _ = sampler.submit(indices, values)
    return sampler

# Instantiate sampler, make a step, and save sampler.
key = jax.random.PRNGKey(42)
sampler = TensorTrainSampler((5, ) * 10)
sampler = step(key, sampler, lambda ix: ix.sum(axis=1))
sampler.save('sampler.pkl')

# Load sampler and continue optimization.
sampler = TensorTrainSampler.load('sampler.pkl')
sampler = step(key, sampler, lambda ix: ix.sum(axis=1))
sampler.save('sampler.pkl')

Note, optax implements optimizers as closures. This results in inability to pickle it with pickle in standard library. Hopefully, there is a cloudpickle which is able to serialize and deserialize closures and lambdas. So, TensorTrainSampler has methods load and save in order to hide these details. Other objects are picklable with standard one.

@article{batsheva2023protes,
    author    = {Batsheva, Anastasia and Chertkov, Andrei  and Ryzhakov, Gleb and Oseledets, Ivan},
    year      = {2023},
    title     = {PROTES: Probabilistic Optimization with Tensor Sampling},
    journal   = {arXiv preprint arXiv:2301.12162},
    url       = {https://arxiv.org/pdf/2301.12162.pdf}
}