{
"cells": [
{
"cell_type": "markdown",
"id": "9c8e1ba6-6b35-4dec-830e-686299698e2b",
"metadata": {},
"source": [
"# Basics"
]
},
{
"cell_type": "markdown",
"id": "ebe10c20-484e-4299-bd2d-255724be621b",
"metadata": {},
"source": [
"This page describes the basic ways to set up a contraction and an explicit optimizer and defines some common terms.\n",
"The focus of `cotengra` is *exact contraction*. That is, taking a collection of tensors with indices described by a tensor network or einsum equation and then:\n",
"\n",
"1. finding the best sequence of pairwise contractions that reduces them to a single output tensor\n",
"2. performing this contraction using a mix of `tensordot` and potentially `einsum` calls\n",
"\n",
"```{note}\n",
"`cotengra` doesn't involve itself with building, modifying, simplifying or decomposing tensors and tensor networks etc.\n",
"```\n",
"\n",
"The minimal information you need to describe such a contraction is:\n",
"\n",
"* the index labels for each tensor\n",
"* the output index labels\n",
"* the size of each index\n",
"\n",
"Here's a very small example of such information involving 4 tensors:"
]
},
{
"cell_type": "code",
"execution_count": 1,
"id": "bafad29a-5ba0-471a-a3e9-5e91847fe02c",
"metadata": {},
"outputs": [],
"source": [
"%config InlineBackend.figure_formats = ['svg']\n",
"import cotengra as ctg\n",
"\n",
"inputs = [\n",
" ('a', 'b', 'x'),\n",
" ('b', 'c', 'd'),\n",
" ('c', 'e', 'y'),\n",
" ('e', 'a', 'd'),\n",
"]\n",
"\n",
"output = ('x', 'y')\n",
"\n",
"size_dict = {'x': 2, 'y': 3, 'a': 4, 'b': 5, 'c': 6, 'd': 7, 'e': 8}"
]
},
{
"cell_type": "markdown",
"id": "878b0d5b-1daa-4644-add0-e38bc2a8873f",
"metadata": {},
"source": [
"This is equivalent to describing an einsum equation and array shapes (such as for `numpy.einsum` or `opt_einsum.contract`):"
]
},
{
"cell_type": "code",
"execution_count": 2,
"id": "e0b47b06-0a19-494c-b639-6104daf8d309",
"metadata": {},
"outputs": [],
"source": [
"import numpy as np\n",
"\n",
"eq = 'abx,bcd,cey,ead->xy'\n",
"shapes = [(4, 5, 2), (5, 6, 7), (6, 8, 3), (8, 4, 7)]\n",
"arrays = [np.ones(s) for s in shapes]"
]
},
{
"cell_type": "markdown",
"id": "904673db-bf08-4abc-ae47-84776ea8565e",
"metadata": {},
"source": [
"The actual names of indices here are only relevant for defining the *geometry*\n",
"within this contraction. \n",
"\n",
"```{note}\n",
"For readability, generally a single unicode character is used per index.\n",
"See `ctg.get_symbol(i)` if you need to generate a usable set of these. \n",
"```\n",
"\n",
"Each index can be thought of as an ***edge*** in a tensor network, \n",
"with each tensor being a ***node***. If every index appears exactly twice in either\n",
"the inputs or output, then the underlying geometry is described as a \n",
"[***graph***](https://en.wikipedia.org/wiki/Graph_(discrete_mathematics)), \n",
"otherwise it is a [***hypergraph***](https://en.wikipedia.org/wiki/Hypergraph). \n",
"You can visualize an input contraction with:"
]
},
{
"cell_type": "code",
"execution_count": 3,
"id": "1bb0f0ef-00ba-46dd-b845-1166f45caa31",
"metadata": {},
"outputs": [
{
"data": {
"image/svg+xml": [
"\n",
"\n",
"\n"
],
"text/plain": [
""
]
},
"metadata": {
"needs_background": "light"
},
"output_type": "display_data"
},
{
"data": {
"text/plain": [
"(, )"
]
},
"execution_count": 3,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"ctg.HyperGraph(inputs, output, size_dict).plot()"
]
},
{
"cell_type": "markdown",
"id": "748b5ef2-02a0-4397-94be-80d0359b7b56",
"metadata": {},
"source": [
"Usually one of these representations is very easy to produce, or a library like `quimb` will do it for you. In any case, the next step is to create an *optimizer*."
]
},
{
"cell_type": "markdown",
"id": "d5ec1d35-8e66-4444-8adc-abec97400b61",
"metadata": {},
"source": [
"## `HyperOptimizer`\n",
"\n",
"The main driver is the [`HyperOptimizer`](cotengra.hyper.HyperOptimizer) class, which optimizes a single contraction: "
]
},
{
"cell_type": "code",
"execution_count": 4,
"id": "0fcf5b6b-fe3a-44ac-ab41-9efb6dd9e02a",
"metadata": {},
"outputs": [],
"source": [
"opt = ctg.HyperOptimizer()"
]
},
{
"cell_type": "markdown",
"id": "4c45918f-b2ae-4285-a69d-7f2ce1137845",
"metadata": {},
"source": [
"The most flexible way to use this is to call the [`search`](cotengra.hyper.HyperOptimizer.search) method which directly returns a\n",
"[`ContractionTree`](cotengra.core.ContractionTree). This is a ***rooted \n",
"binary tree***:"
]
},
{
"cell_type": "code",
"execution_count": 5,
"id": "e586ea8d-5da3-43d6-9685-8a2f08b6a1a3",
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
""
]
},
"execution_count": 5,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"# direct use\n",
"tree = opt.search(inputs, output, size_dict)\n",
"tree"
]
},
{
"cell_type": "code",
"execution_count": 23,
"id": "947587ee",
"metadata": {},
"outputs": [
{
"data": {
"image/svg+xml": [
"\n",
"\n",
"\n"
],
"text/plain": [
""
]
},
"metadata": {
"needs_background": "light"
},
"output_type": "display_data"
}
],
"source": [
"tree.plot_flat();"
]
},
{
"cell_type": "markdown",
"id": "f2ec336e-1db6-434e-be91-d116c1f784d2",
"metadata": {},
"source": [
"The tree (which also has the mathematical name, \n",
"['carving decomposition'](https://en.wikipedia.org/wiki/Branch-decomposition#Carving_width)) contains all the crucial information about costs and sizes of\n",
"intermediate tensors:"
]
},
{
"cell_type": "code",
"execution_count": 6,
"id": "9c06c75f-464e-492b-a1d2-20c8c3b63aab",
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"(8.39231742277876, 4656.0)"
]
},
"execution_count": 6,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"tree.contraction_width(), tree.contraction_cost()"
]
},
{
"cell_type": "markdown",
"id": "dbeb57ba-a73b-4d8e-b7e8-6b69b6fdcf6d",
"metadata": {},
"source": [
"* the ***contraction width***, $W$, is $log_2$ the size of the largest intermediate tensor\n",
"* the ***contraction cost***, $C$, is the total number of scalar multiplications\n",
"\n",
"The tree can be used to perform the actual contraction too:"
]
},
{
"cell_type": "code",
"execution_count": 7,
"id": "cbf3ec5a-d204-4ff4-843a-c17f6c7879ec",
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"array([[6720., 6720., 6720.],\n",
" [6720., 6720., 6720.]])"
]
},
"execution_count": 7,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"tree.contract(arrays)"
]
},
{
"cell_type": "markdown",
"id": "98e7a525-f979-4433-aa5f-12b90264f9f6",
"metadata": {},
"source": [
"A tree combined with a specific traversal ordering is known as a ***path***: "
]
},
{
"cell_type": "code",
"execution_count": 8,
"id": "d0aa12c6-b61d-427b-87b6-ebd81cf4803d",
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"((0, 1), (1, 2), (0, 1))"
]
},
"execution_count": 8,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"path = tree.get_path()\n",
"path"
]
},
{
"cell_type": "markdown",
"id": "074140f1-f095-4671-b227-adba4dcba382",
"metadata": {},
"source": [
"Such paths can be supplied to `opt_einsum` and `quimb` functions, *or*\n",
"you can supply the [`HyperOptimizer`](cotengra.hyper.HyperOptimizer) \n",
"itself, in which case it will first run a search and then pass on a path."
]
},
{
"cell_type": "markdown",
"id": "7a784d01-fc8b-44e3-8204-ee11db4f924f",
"metadata": {
"tags": []
},
"source": [
"#### With ``quimb``"
]
},
{
"cell_type": "code",
"execution_count": 9,
"id": "42aa390b-65c5-4779-bb0f-df1f2c592e05",
"metadata": {},
"outputs": [
{
"data": {
"text/html": [
"Tensor(shape=(2, 3), inds=[x, y], tags={}),backend=numpy, dtype=float64, data=array([[6720., 6720., 6720.],\n",
" [6720., 6720., 6720.]])"
],
"text/plain": [
"Tensor(shape=(2, 3), inds=('x', 'y'), tags=oset([]))"
]
},
"execution_count": 9,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"import quimb.tensor as qtn\n",
"\n",
"tn = qtn.TensorNetwork([\n",
" qtn.Tensor(array, inds)\n",
" for array, inds in zip(arrays, inputs)\n",
"])\n",
"\n",
"tn.contract(..., optimize=ctg.HyperOptimizer())"
]
},
{
"cell_type": "markdown",
"id": "839765b6-8451-499b-85cd-f753f556ea64",
"metadata": {},
"source": [
"Note for non-hyper graphs `quimb` will figure out the output indices for you,\n",
"else you will need to supply `output_inds`. `quimb` also knows how to return\n",
"the [`ContractionTree`](cotengra.core.ContractionTree) directly with:"
]
},
{
"cell_type": "code",
"execution_count": 10,
"id": "31bfd2a9-9128-4c74-98a2-064179e13d37",
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
""
]
},
"execution_count": 10,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"tn.contraction_tree(path)"
]
},
{
"cell_type": "markdown",
"id": "b858c274-1dc7-428a-95f5-bbd741bdb023",
"metadata": {},
"source": [
"And many other methods and algorithms take a `optimize=path_or_optimizer` option."
]
},
{
"cell_type": "markdown",
"id": "cf9fec71-71f8-4394-87c7-a3336c91af41",
"metadata": {},
"source": [
"#### With `opt_einsum`"
]
},
{
"cell_type": "markdown",
"id": "28c7ea55-c4d2-4d2d-a65c-db982b766bc5",
"metadata": {},
"source": [
"You can supply a `path` or `HyperOptimizer` to all the functions of `opt_einsum`\n",
"which take an `optimize` kwarg:"
]
},
{
"cell_type": "code",
"execution_count": 11,
"id": "c4235c8d-47c5-4a15-ab25-e58ab98ac559",
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"array([[6720., 6720., 6720.],\n",
" [6720., 6720., 6720.]])"
]
},
"execution_count": 11,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"import opt_einsum as oe\n",
"\n",
"oe.contract(eq, *arrays, optimize=path)"
]
},
{
"cell_type": "markdown",
"id": "476158f6-4dfb-488e-9df9-e9e0be1c10a0",
"metadata": {},
"source": [
"```{hint}\n",
"For convenience `cotengra` also registers a few presets which can be used like `optimize='hyper'`, these can also be created.\n",
"```"
]
},
{
"cell_type": "markdown",
"id": "9cded446-465a-4ced-a133-6dc55d1eebf7",
"metadata": {},
"source": [
"A single [`HyperOptimizer`](cotengra.hyper.HyperOptimizer) instance can only be used for a single contraction - everytime you supply it, as long as the contraction matches, it will simply continue it's search."
]
},
{
"cell_type": "markdown",
"id": "d77aee38-47b8-4197-93f1-8b4ba41e5869",
"metadata": {},
"source": [
"## `ReusableHyperOptimizer`\n",
"\n",
"Often, instead you want a single optimizer with maybe customized settings to use for many different contractions - the answer is to use a [`ReusableHyperOptimizer`](cotengra.hyper.ReusableHyperOptimizer).\n",
"Everytime this is supplied to a *new* contraction it runs a search and stores the resulting path. The next time it sees this contraction it simply returns this cached path."
]
},
{
"cell_type": "code",
"execution_count": 12,
"id": "74a31688-7048-44f2-8c2c-e860689ef295",
"metadata": {},
"outputs": [],
"source": [
"opt = ctg.ReusableHyperOptimizer(progbar=True)"
]
},
{
"cell_type": "code",
"execution_count": 13,
"id": "029f66c6-76e4-490c-8db1-27f737fd6308",
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"log2[SIZE]: 8.39 log10[FLOPs]: 3.67: 23%|██▎ | 29/128 [00:00<00:01, 70.01it/s] "
]
},
{
"name": "stderr",
"output_type": "stream",
"text": [
"log2[SIZE]: 8.39 log10[FLOPs]: 3.67: 100%|██████████| 128/128 [00:02<00:00, 61.43it/s]\n"
]
},
{
"data": {
"text/plain": [
""
]
},
"execution_count": 13,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"opt.search(inputs, output, size_dict)"
]
},
{
"cell_type": "code",
"execution_count": 14,
"id": "cfb5758d-dccd-41bb-842a-ce69274815a4",
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
""
]
},
"execution_count": 14,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"opt.search(inputs, output, size_dict)"
]
},
{
"cell_type": "markdown",
"id": "44b3e161-f06e-47ed-a7b4-ef10a2dabc6f",
"metadata": {},
"source": [
"Note how the second call didn't display a progress bar as it used the cached result.\n",
"\n",
"```{hint}\n",
"The contractions are not cached using full (hyper) graph isomoprhism, which would not be scalable. Instead, the inputs have to be in the same order to produce the same hash key.\n",
"```"
]
},
{
"cell_type": "markdown",
"id": "2274aa4f-69db-4f90-b22a-61071872aa95",
"metadata": {},
"source": [
"### Disk persistence\n",
"\n",
"If you want to persist contraction paths across python sessions (i.e. don't want to explicitly save the `tree` or `path` objects yourself), then you can supply the `directory` kwarg:"
]
},
{
"cell_type": "code",
"execution_count": 15,
"id": "cfb5c741-176a-4fe2-bc57-fe80ac6eb112",
"metadata": {},
"outputs": [],
"source": [
"opt = ctg.ReusableHyperOptimizer(directory='cotengra_cache')"
]
},
{
"cell_type": "code",
"execution_count": 16,
"id": "74328ba8-826b-44d1-9cd9-d16809b88ab8",
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
""
]
},
"execution_count": 16,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"opt.search(inputs, output, size_dict)"
]
},
{
"cell_type": "markdown",
"id": "74b06b17-a7af-4464-8409-729f983d0de5",
"metadata": {},
"source": [
"The directory contains a single pickled file per contraction it has seen:"
]
},
{
"cell_type": "code",
"execution_count": 17,
"id": "4c7e466e-c7bf-43dd-a198-888fdc7334a0",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"0489b9ab17a455bebcf3cc86c8ea5642f518ae95\n"
]
}
],
"source": [
"!ls cotengra_cache/"
]
},
{
"cell_type": "code",
"execution_count": 18,
"id": "0db8e920-1d77-4e15-b729-e1e1d4204584",
"metadata": {},
"outputs": [],
"source": [
"# clean it up for now\n",
"!rm -rf cotengra_cache/"
]
},
{
"cell_type": "markdown",
"id": "45f14770",
"metadata": {},
"source": [
"If you supply `directory=True` then the cache name will be generated from a \n",
"hash of the relevant path finding options supplied to the optimizer, meaning \n",
"you don't need to manually change the name in order to use different caches \n",
"for different settings. "
]
},
{
"cell_type": "markdown",
"id": "d9c6309b",
"metadata": {},
"source": [
"## What next?\n",
"\n",
"* The ['Advanced'](advanced) page, describes how to customize a [`HyperOptimizer`](cotengra.hyper.HyperOptimizer) in detail, for example to include subtree reconfiguration and dynamic slicing\n",
"* The ['Visualization'](visualization) page details many functions for plotting the contraction, tree and optimizer trials\n",
"* The ['Contraction'](contraction) page contains more information about actually performing the contraction, for example using a GPU\n",
"* The ['Tree Surgery'](trees) page describes the design of the central [`ContractionTree`](cotengra.core.ContractionTree) object and ways to manipulate it\n",
"\n",
"## Quick-start\n",
"\n",
"If you just want to get going, the following illustrate some common customizations of the optimizers. "
]
},
{
"cell_type": "markdown",
"id": "4dfeec6c",
"metadata": {},
"source": [
"### High quality sliced optimizer\n",
"\n",
"The following is an example of a high quality optimizer you might use to search\n",
"for a single contraction, where you are memory bound to width $W=30$ and thus \n",
"need to use slicing: "
]
},
{
"cell_type": "code",
"execution_count": 19,
"id": "e9a47175",
"metadata": {},
"outputs": [],
"source": [
"opt_hq_W30 = ctg.HyperOptimizer(\n",
" # do extra runs\n",
" max_repeats=1024,\n",
" # use dynamic slicing to target a width of 30\n",
" slicing_reconf_opts={'target_size': 2**30},\n",
" # use the nevergrad space searcher - good with large trial budget\n",
" optlib='nevergrad',\n",
" # terminate search if no change for 128 trials\n",
" max_time='equil:128',\n",
" # show live progress\n",
" progbar=True,\n",
")"
]
},
{
"cell_type": "markdown",
"id": "2d98d65b",
"metadata": {},
"source": [
"* Everytime you supply this optimizer instance it will *continue* its search, so you \n",
"can interactively run it until you are happy or it seems to have converged.\n",
"\n",
"* While a few hundred runs is usually sufficient when no slicing is needed, very large contractions requiring heavy slicing might benefit from a few thousand runs."
]
},
{
"cell_type": "markdown",
"id": "d4207442",
"metadata": {},
"source": [
"### Economical optimizer\n",
"\n",
"The following is an example of a reusable optimizer that is cheap to run and \n",
"requires no extra depedencies (i.e. `kahypar` or a Bayesian optimizer), \n",
"but will still yield much better results than simple algorithms. \n",
"Useful if you have many smallish contractions:"
]
},
{
"cell_type": "code",
"execution_count": 20,
"id": "1a9b54ea",
"metadata": {},
"outputs": [],
"source": [
"opt_eco = ctg.ReusableHyperOptimizer(\n",
" # just do a few runs\n",
" max_repeats=32,\n",
" # only use the basic greedy optimizer ...\n",
" methods=['greedy'],\n",
" # ... but pair it with reconfiguration\n",
" reconf_opts={},\n",
" # just uniformly sample the space\n",
" optlib='random',\n",
" # terminate search if contraction is cheap\n",
" max_time='rate:1e6',\n",
" # account for both flops and write - usually wise for practical performance\n",
" minimize='combo-64',\n",
" # persist paths found in here\n",
" directory='cotengra_cache_eco',\n",
")"
]
},
{
"cell_type": "code",
"execution_count": 21,
"id": "70a1551d-7a15-4f68-b250-79f789a0a095",
"metadata": {},
"outputs": [],
"source": [
"# clean up the cache for now\n",
"!rm -rf cotengra_cache_eco"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.11.3"
}
},
"nbformat": 4,
"nbformat_minor": 5
}