Skip to content

# Custom Path Optimizers

If you want to implement or just experiment with custom contaction paths then you can easily by subclassing the `opt_einsum.paths.PathOptimizer` object. For example, imagine we want to test the path that just blindly contracts the first pair of tensors again and again. We would implement this as:

    import opt_einsum as oe

    class MyOptimizer(oe.paths.PathOptimizer):

        def __call__(self, inputs, output, size_dict, memory_limit=None):
            return [(0, 1)] * (len(inputs) - 1)

Once defined we can use this as:

    import numpy as np

    # set-up a random contraction
    eq, shapes = oe.helpers.rand_equation(10, 3, seed=42)
    arrays = list(map(np.ones, shapes))

    # set-up our optimizer and use it
    optimizer = MyOptimizer()
    path, path_info = oe.contract_path(eq, *arrays, optimize=optimizer)

    print(path)
    #> [(0, 1), (0, 1), (0, 1), (0, 1), (0, 1), (0, 1), (0, 1), (0, 1), (0, 1)]

    print(path_info.speedup)
    #> 133.21363671496357

Note that though we still get a considerable speedup over `einsum` this is of course not a good strategy to take in general.

## Custom Random Optimizers

If your custom path optimizer is inherently random, then you can reuse all the machinery of the random-greedy approach. Namely:

- A max-repeats or max-time approach
- Minimization with respect to total flops or largest intermediate size
- Parallelization using a pool-executor

This is done by subclassing the `opt_einsum.paths.RandomOptimizer` object and implementing a `setup` method. Here's an example where we just randomly select any path (again, although we get a considerable speedup over `einsum` this is not a good strategy to take in general):

    from opt_einsum.path_random import ssa_path_compute_cost

    class MyRandomOptimizer(oe.path_random.RandomOptimizer):

        @staticmethod
        def random_path(r, n, inputs, output, size_dict):
            """Picks a completely random contraction order.
            """
            np.random.seed(r)
            ssa_path = []
            remaining = set(range(n))
            while len(remaining) > 1:
                i, j = np.random.choice(list(remaining), size=2, replace=False)
                remaining.add(n + len(ssa_path))
                remaining.remove(i)
                remaining.remove(j)
                ssa_path.append((i, j))
            cost, size = ssa_path_compute_cost(ssa_path, inputs, output, size_dict)
            return ssa_path, cost, size

        def setup(self, inputs, output, size_dict):
            """Prepares the function and arguments to repeatedly call.
            """
            n = len(inputs)
            trial_fn = self.random_path
            trial_args = (n, inputs, output, size_dict)
            return trial_fn, trial_args

Which we can now instantiate using various other options:

    optimizer = MyRandomOptimizer(max_repeats=1000, max_time=10,
                                  parallel=True, minimize='size')
    path, path_info = oe.contract_path(eq, *arrays, optimize=optimizer)

    print(path)
    #> [(3, 4), (1, 3), (0, 3), (3, 5), (3, 4), (3, 4), (1, 0), (0, 1), (0, 1)]

    print(path_info.speedup)
    #> 712829.9451056132

There are a few things to note here:

1. The core function (`MyRandomOptimizer.random_path` here), should take a trial number `r` as it first argument
2. It should return a ssa_path (see `opt_einsum.paths.ssa_to_linear` and `opt_einsum.paths.linear_to_ssa`) as well as a flops-cost and max-size.
3. The `setup` method prepares this function, as well as any input to it, so that the trials will look roughly like `[trial_fn(r, *trial_args) for r in range(max_repeats)]`. If you need to parse the standard arguments (into a network for example), it thus only needs to be done once per optimization

More details about `opt_einsum.paths.RandomOptimizer` options can be found in `RandomGreedyPathPage` section.