[
{
"path": ".codecov.yml",
"content": "codecov:\n require_ci_to_pass: yes\n\ncoverage:\n range: 50..100\n status:\n project:\n default:\n informational: true\n patch:\n default:\n informational: true\n changes: false\n\ncomment: off\n"
},
{
"path": ".gitattributes",
"content": "# Auto detect text files and perform LF normalization\n* text=auto\n\n# Standard to msysgit\n*.doc\t diff=astextplain\n*.DOC\t diff=astextplain\n*.docx diff=astextplain\n*.DOCX diff=astextplain\n*.dot diff=astextplain\n*.DOT diff=astextplain\n*.pdf diff=astextplain\n*.PDF\t diff=astextplain\n*.rtf\t diff=astextplain\n*.RTF\t diff=astextplain\n\n# include the version number in git archive\ncotengra/_version.py export-subst\n\n# make cotengra appear as a python project on github\n*.ipynb linguist-language=Python\n\n# SCM syntax highlighting & preventing 3-way merges\npixi.lock merge=binary linguist-language=YAML linguist-generated=true\n"
},
{
"path": ".github/dependabot.yml",
"content": "# To get started with Dependabot version updates, you'll need to specify which\n# package ecosystems to update and where the package manifests are located.\n# Please see the documentation for all configuration options:\n# https://docs.github.com/code-security/dependabot/dependabot-version-updates/configuration-options-for-the-dependabot.yml-file\n\nversion: 2\nupdates:\n # Enable Dependabot for GitHub Actions\n - package-ecosystem: \"github-actions\"\n directory: \"/\" # Location of your workflows, typically the root folder\n schedule:\n interval: \"daily\" # Frequency of update checks (daily, weekly, or monthly)\n"
},
{
"path": ".github/workflows/pypi-release.yml",
"content": "name: Build and Upload cotengra to PyPI\non:\n release:\n types:\n - published\n push:\n tags:\n - 'v*'\n\njobs:\n build-artifacts:\n runs-on: ubuntu-latest\n if: github.repository == 'jcmgray/cotengra'\n steps:\n - uses: actions/checkout@v6\n with:\n fetch-depth: 0\n - uses: actions/setup-python@v6\n name: Install Python\n with:\n python-version: \"3.12\"\n\n - name: Install dependencies\n run: |\n python -m pip install --upgrade pip\n python -m pip install build twine\n\n - name: Build tarball and wheels\n run: |\n git clean -xdf\n git restore -SW .\n python -m build\n\n - name: Check built artifacts\n run: |\n python -m twine check --strict dist/*\n pwd\n if [ -f dist/cotengra-0.0.0.tar.gz ]; then\n echo \"❌ INVALID VERSION NUMBER\"\n exit 1\n else\n echo \"✅ Looks good\"\n fi\n - uses: actions/upload-artifact@v7\n with:\n name: releases\n path: dist\n\n test-built-dist:\n needs: build-artifacts\n runs-on: ubuntu-latest\n steps:\n - uses: actions/setup-python@v6\n name: Install Python\n with:\n python-version: \"3.12\"\n - uses: actions/download-artifact@v8\n with:\n name: releases\n path: dist\n - name: List contents of built dist\n run: |\n ls -ltrh\n ls -ltrh dist\n\n - name: Verify the built dist/wheel is valid\n if: github.event_name == 'push'\n run: |\n python -m pip install --upgrade pip\n python -m pip install dist/cotengra*.whl\n\n upload-to-test-pypi:\n needs: test-built-dist\n if: github.event_name == 'push'\n runs-on: ubuntu-latest\n\n environment:\n name: pypi\n url: https://test.pypi.org/p/cotengra\n permissions:\n id-token: write\n\n steps:\n - uses: actions/download-artifact@v8\n with:\n name: releases\n path: dist\n - name: Publish package to TestPyPI\n if: github.event_name == 'push'\n uses: pypa/gh-action-pypi-publish@v1.14.0\n with:\n repository-url: https://test.pypi.org/legacy/\n verbose: true\n\n\n upload-to-pypi:\n needs: test-built-dist\n if: github.event_name == 'release'\n runs-on: ubuntu-latest\n\n environment:\n name: pypi\n url: https://pypi.org/p/cotengra\n permissions:\n id-token: write\n\n steps:\n - uses: actions/download-artifact@v8\n with:\n name: releases\n path: dist\n - name: Publish package to PyPI\n uses: pypa/gh-action-pypi-publish@v1.14.0\n with:\n verbose: true"
},
{
"path": ".github/workflows/test.yml",
"content": "name: Tests\n\non:\n workflow_dispatch:\n push:\n pull_request:\n\ndefaults:\n run:\n shell: bash -l {0}\n\njobs:\n run-tests:\n runs-on: ${{ matrix.os }}\n strategy:\n matrix:\n include:\n - os: ubuntu-latest\n pixi_environment: testminimal\n pixi_task: test\n\n - os: ubuntu-latest\n pixi_environment: testpyold\n pixi_task: test\n\n - os: ubuntu-latest\n pixi_environment: testpynew\n pixi_task: test\n\n - os: macos-latest\n pixi_environment: testpymid\n pixi_task: test\n\n - os: windows-latest\n pixi_environment: testpymid\n pixi_task: test\n\n - os: ubuntu-latest\n pixi_environment: testtorch\n pixi_task: test-backends\n\n - os: ubuntu-latest\n pixi_environment: testjax\n pixi_task: test-backends\n\n - os: ubuntu-latest\n pixi_environment: testtensorflow\n pixi_task: test-backends\n\n steps:\n - uses: actions/checkout@v6\n\n - uses: prefix-dev/setup-pixi@v0.9.5\n with:\n environments: ${{ matrix.pixi_environment }}\n\n - name: Test with pytest\n run: pixi run -e ${{ matrix.pixi_environment }} ${{ matrix.pixi_task }}\n\n - name: Report to codecov\n uses: codecov/codecov-action@v6\n with:\n token: ${{ secrets.CODECOV_TOKEN }}\n"
},
{
"path": ".gitignore",
"content": "# Byte-compiled / optimized / DLL files\n__pycache__/\n*.py[cod]\n*$py.class\n\n# C extensions\n*.so\n\n# Distribution / packaging\n.Python\nenv/\nbuild/\ndevelop-eggs/\ndist/\ndownloads/\neggs/\n.eggs/\nlib/\nlib64/\nparts/\nsdist/\nvar/\n*.egg-info/\n.installed.cfg\n*.egg\n\n# PyInstaller\n# Usually these files are written by a python script from a template\n# before PyInstaller builds the exe, so as to inject date/other infos into it.\n*.manifest\n*.spec\n\n# Installer logs\npip-log.txt\npip-delete-this-directory.txt\n\n# Unit test / coverage reports\nhtmlcov/\n.tox/\n.coverage\n.coverage.*\n.cache\nnosetests.xml\ncoverage.xml\n*,cover\n\n# Translations\n*.mo\n*.pot\n\n# Django stuff:\n*.log\n\n# Sphinx documentation\ndocs/_build/\n\n# PyBuilder\ntarget/\n\n.pytest_cache\n.vscode\n\n**/.ipynb_checkpoints/\n\n# Added by cargo\n/target\nCargo.lock\n\ncotengra/_version.py\n\nexperiments\n\n# pixi environments\n.pixi/*\n!.pixi/config.toml"
},
{
"path": ".readthedocs.yml",
"content": "version: 2\n\nsphinx:\n configuration: docs/conf.py\n\nbuild:\n os: \"ubuntu-24.04\"\n tools:\n python: \"latest\"\n jobs:\n create_environment:\n - asdf plugin add pixi\n - asdf install pixi latest\n - asdf global pixi latest\n install:\n - pixi install -e docs\n build:\n html:\n - pixi run readthedocs\n\nformats: []\n"
},
{
"path": "LICENSE.md",
"content": "\n Apache License\n Version 2.0, January 2004\n http://www.apache.org/licenses/\n\n TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION\n\n 1. Definitions.\n\n \"License\" shall mean the terms and conditions for use, reproduction,\n and distribution as defined by Sections 1 through 9 of this document.\n\n \"Licensor\" shall mean the copyright owner or entity authorized by\n the copyright owner that is granting the License.\n\n \"Legal Entity\" shall mean the union of the acting entity and all\n other entities that control, are controlled by, or are under common\n control with that entity. For the purposes of this definition,\n \"control\" means (i) the power, direct or indirect, to cause the\n direction or management of such entity, whether by contract or\n otherwise, or (ii) ownership of fifty percent (50%) or more of the\n outstanding shares, or (iii) beneficial ownership of such entity.\n\n \"You\" (or \"Your\") shall mean an individual or Legal Entity\n exercising permissions granted by this License.\n\n \"Source\" form shall mean the preferred form for making modifications,\n including but not limited to software source code, documentation\n source, and configuration files.\n\n \"Object\" form shall mean any form resulting from mechanical\n transformation or translation of a Source form, including but\n not limited to compiled object code, generated documentation,\n and conversions to other media types.\n\n \"Work\" shall mean the work of authorship, whether in Source or\n Object form, made available under the License, as indicated by a\n copyright notice that is included in or attached to the work\n (an example is provided in the Appendix below).\n\n \"Derivative Works\" shall mean any work, whether in Source or Object\n form, that is based on (or derived from) the Work and for which the\n editorial revisions, annotations, elaborations, or other modifications\n represent, as a whole, an original work of authorship. For the purposes\n of this License, Derivative Works shall not include works that remain\n separable from, or merely link (or bind by name) to the interfaces of,\n the Work and Derivative Works thereof.\n\n \"Contribution\" shall mean any work of authorship, including\n the original version of the Work and any modifications or additions\n to that Work or Derivative Works thereof, that is intentionally\n submitted to Licensor for inclusion in the Work by the copyright owner\n or by an individual or Legal Entity authorized to submit on behalf of\n the copyright owner. For the purposes of this definition, \"submitted\"\n means any form of electronic, verbal, or written communication sent\n to the Licensor or its representatives, including but not limited to\n communication on electronic mailing lists, source code control systems,\n and issue tracking systems that are managed by, or on behalf of, the\n Licensor for the purpose of discussing and improving the Work, but\n excluding communication that is conspicuously marked or otherwise\n designated in writing by the copyright owner as \"Not a Contribution.\"\n\n \"Contributor\" shall mean Licensor and any individual or Legal Entity\n on behalf of whom a Contribution has been received by Licensor and\n subsequently incorporated within the Work.\n\n 2. Grant of Copyright License. Subject to the terms and conditions of\n this License, each Contributor hereby grants to You a perpetual,\n worldwide, non-exclusive, no-charge, royalty-free, irrevocable\n copyright license to reproduce, prepare Derivative Works of,\n publicly display, publicly perform, sublicense, and distribute the\n Work and such Derivative Works in Source or Object form.\n\n 3. Grant of Patent License. Subject to the terms and conditions of\n this License, each Contributor hereby grants to You a perpetual,\n worldwide, non-exclusive, no-charge, royalty-free, irrevocable\n (except as stated in this section) patent license to make, have made,\n use, offer to sell, sell, import, and otherwise transfer the Work,\n where such license applies only to those patent claims licensable\n by such Contributor that are necessarily infringed by their\n Contribution(s) alone or by combination of their Contribution(s)\n with the Work to which such Contribution(s) was submitted. If You\n institute patent litigation against any entity (including a\n cross-claim or counterclaim in a lawsuit) alleging that the Work\n or a Contribution incorporated within the Work constitutes direct\n or contributory patent infringement, then any patent licenses\n granted to You under this License for that Work shall terminate\n as of the date such litigation is filed.\n\n 4. Redistribution. You may reproduce and distribute copies of the\n Work or Derivative Works thereof in any medium, with or without\n modifications, and in Source or Object form, provided that You\n meet the following conditions:\n\n (a) You must give any other recipients of the Work or\n Derivative Works a copy of this License; and\n\n (b) You must cause any modified files to carry prominent notices\n stating that You changed the files; and\n\n (c) You must retain, in the Source form of any Derivative Works\n that You distribute, all copyright, patent, trademark, and\n attribution notices from the Source form of the Work,\n excluding those notices that do not pertain to any part of\n the Derivative Works; and\n\n (d) If the Work includes a \"NOTICE\" text file as part of its\n distribution, then any Derivative Works that You distribute must\n include a readable copy of the attribution notices contained\n within such NOTICE file, excluding those notices that do not\n pertain to any part of the Derivative Works, in at least one\n of the following places: within a NOTICE text file distributed\n as part of the Derivative Works; within the Source form or\n documentation, if provided along with the Derivative Works; or,\n within a display generated by the Derivative Works, if and\n wherever such third-party notices normally appear. The contents\n of the NOTICE file are for informational purposes only and\n do not modify the License. You may add Your own attribution\n notices within Derivative Works that You distribute, alongside\n or as an addendum to the NOTICE text from the Work, provided\n that such additional attribution notices cannot be construed\n as modifying the License.\n\n You may add Your own copyright statement to Your modifications and\n may provide additional or different license terms and conditions\n for use, reproduction, or distribution of Your modifications, or\n for any such Derivative Works as a whole, provided Your use,\n reproduction, and distribution of the Work otherwise complies with\n the conditions stated in this License.\n\n 5. Submission of Contributions. Unless You explicitly state otherwise,\n any Contribution intentionally submitted for inclusion in the Work\n by You to the Licensor shall be under the terms and conditions of\n this License, without any additional terms or conditions.\n Notwithstanding the above, nothing herein shall supersede or modify\n the terms of any separate license agreement you may have executed\n with Licensor regarding such Contributions.\n\n 6. Trademarks. This License does not grant permission to use the trade\n names, trademarks, service marks, or product names of the Licensor,\n except as required for reasonable and customary use in describing the\n origin of the Work and reproducing the content of the NOTICE file.\n\n 7. Disclaimer of Warranty. Unless required by applicable law or\n agreed to in writing, Licensor provides the Work (and each\n Contributor provides its Contributions) on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or\n implied, including, without limitation, any warranties or conditions\n of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A\n PARTICULAR PURPOSE. You are solely responsible for determining the\n appropriateness of using or redistributing the Work and assume any\n risks associated with Your exercise of permissions under this License.\n\n 8. Limitation of Liability. In no event and under no legal theory,\n whether in tort (including negligence), contract, or otherwise,\n unless required by applicable law (such as deliberate and grossly\n negligent acts) or agreed to in writing, shall any Contributor be\n liable to You for damages, including any direct, indirect, special,\n incidental, or consequential damages of any character arising as a\n result of this License or out of the use or inability to use the\n Work (including but not limited to damages for loss of goodwill,\n work stoppage, computer failure or malfunction, or any and all\n other commercial damages or losses), even if such Contributor\n has been advised of the possibility of such damages.\n\n 9. Accepting Warranty or Additional Liability. While redistributing\n the Work or Derivative Works thereof, You may choose to offer,\n and charge a fee for, acceptance of support, warranty, indemnity,\n or other liability obligations and/or rights consistent with this\n License. However, in accepting such obligations, You may act only\n on Your own behalf and on Your sole responsibility, not on behalf\n of any other Contributor, and only if You agree to indemnify,\n defend, and hold each Contributor harmless for any liability\n incurred by, or claims asserted against, such Contributor by reason\n of your accepting any such warranty or additional liability.\n\n END OF TERMS AND CONDITIONS\n\n APPENDIX: How to apply the Apache License to your work.\n\n To apply the Apache License to your work, attach the following\n boilerplate notice, with the fields enclosed by brackets \"{}\"\n replaced with your own identifying information. (Don't include\n the brackets!) The text should be enclosed in the appropriate\n comment syntax for the file format. We also recommend that a\n file or class name and description of purpose be included on the\n same \"printed page\" as the copyright notice for easier\n identification within third-party archives.\n\n Copyright {yyyy} {name of copyright owner}\n\n Licensed under the Apache License, Version 2.0 (the \"License\");\n you may not use this file except in compliance with the License.\n You may obtain a copy of the License at\n\n http://www.apache.org/licenses/LICENSE-2.0\n\n Unless required by applicable law or agreed to in writing, software\n distributed under the License is distributed on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n See the License for the specific language governing permissions and\n limitations under the License.\n"
},
{
"path": "MANIFEST.in",
"content": "include cotengra/kahypar_profiles/*.ini\ninclude cotengra/kahypar_profiles/old/*.ini\ninclude LICENSE.md\ninclude README.md\ngraft tests\n"
},
{
"path": "README.md",
"content": "
\n\n[](https://github.com/jcmgray/cotengra/actions/workflows/test.yml)\n[](https://codecov.io/gh/jcmgray/cotengra)\n[](https://cotengra.readthedocs.io)\n[](https://pypi.org/project/cotengra/)\n[](https://anaconda.org/conda-forge/cotengra)\n[](https://pixi.sh)\n\n`cotengra` is a python library for contracting tensor networks or einsum\nexpressions involving large numbers of tensors - the main docs can be found\nat [cotengra.readthedocs.io](https://cotengra.readthedocs.io/).\nSome of the key feautures of `cotengra` include:\n\n* drop-in ``einsum`` and ``ncon`` replacement\n* an explicit **contraction tree** object that can be flexibly built, modified and visualized\n* a **'hyper optimizer'** that samples trees while tuning the generating meta-paremeters\n* **dynamic slicing** for massive memory savings and parallelism\n* **simulated annealing** as an alternative optimizing and slicing strategy\n* support for **hyper** edge tensor networks and thus arbitrary einsum equations\n* **paths** that can be supplied to [`numpy.einsum`](https://numpy.org/doc/stable/reference/generated/numpy.einsum.html), [`opt_einsum`](https://dgasmith.github.io/opt_einsum/), [`quimb`](https://quimb.readthedocs.io/en/latest/) among others\n* **performing contractions** with tensors from many libraries via [`autoray`](https://github.com/jcmgray/autoray),\n even if they don't provide `einsum` or `tensordot` but do have (batch) matrix\n multiplication\n\n
\n"
},
{
"path": "cotengra/__init__.py",
"content": "\"\"\"Hyper optimized contraction trees for large tensor networks and einsums.\"\"\"\n\nfrom importlib.metadata import PackageNotFoundError as _PackageNotFoundError\nfrom importlib.metadata import version as _version\n\ntry:\n __version__ = _version(\"cotengra\")\nexcept _PackageNotFoundError:\n try:\n # fallback for source trees where hatch-vcs has generated _version.py.\n from ._version import version as __version__\n except ImportError:\n __version__ = \"0.0.0+unknown\"\n\nimport functools\nimport warnings\n\nfrom . import utils\nfrom .core import (\n ContractionTree,\n ContractionTreeCompressed,\n)\nfrom .core_multi import (\n ContractionTreeMulti,\n)\nfrom .hypergraph import (\n HyperGraph,\n get_hypergraph,\n)\nfrom .hyperoptimizers import (\n hyper_cmaes,\n hyper_es,\n hyper_neldermead,\n hyper_nevergrad,\n hyper_optuna,\n hyper_random,\n hyper_sbplx,\n hyper_skopt,\n)\nfrom .hyperoptimizers.hyper import (\n HyperCompressedOptimizer,\n HyperMultiOptimizer,\n HyperOptimizer,\n ReusableHyperCompressedOptimizer,\n ReusableHyperOptimizer,\n get_hyper_space,\n list_hyper_functions,\n)\nfrom .interface import (\n array_contract,\n array_contract_expression,\n array_contract_path,\n array_contract_tree,\n einsum,\n einsum_expression,\n einsum_tree,\n ncon,\n register_preset,\n)\nfrom .oe import PathOptimizer\nfrom .pathfinders import (\n path_basic,\n path_compressed_greedy,\n path_greedy,\n path_igraph,\n path_kahypar,\n path_labels,\n)\nfrom .pathfinders.path_basic import (\n GreedyOptimizer,\n OptimalOptimizer,\n RandomGreedyOptimizer,\n ReusableRandomGreedyOptimizer,\n edge_path_to_linear,\n edge_path_to_ssa,\n linear_to_ssa,\n ssa_to_linear,\n)\nfrom .pathfinders.path_flowcutter import (\n FlowCutterOptimizer,\n optimize_flowcutter,\n)\nfrom .pathfinders.path_quickbb import QuickBBOptimizer, optimize_quickbb\nfrom .pathfinders.path_random import RandomOptimizer\nfrom .plot import (\n plot_contractions,\n plot_contractions_alt,\n plot_scatter,\n plot_scatter_alt,\n plot_slicings,\n plot_slicings_alt,\n plot_tree,\n plot_tree_ring,\n plot_tree_span,\n plot_tree_tent,\n plot_trials,\n plot_trials_alt,\n)\nfrom .presets import (\n AutoHQOptimizer,\n AutoOptimizer,\n auto_hq_optimize,\n auto_optimize,\n greedy_optimize,\n optimal_optimize,\n optimal_outer_optimize,\n)\nfrom .reusable import (\n hash_contraction,\n)\nfrom .slicer import SliceFinder\nfrom .utils import (\n get_symbol,\n get_symbol_map,\n)\n\nUniformOptimizer = functools.partial(HyperOptimizer, optlib=\"random\")\n\"\"\"Does no gaussian process tuning by default, just randomly samples - requires\nno optimization library.\n\"\"\"\n\ncontract_expression = einsum_expression\n\"\"\"Alias for :func:`cotengra.einsum_expression`.\"\"\"\n\ncontract = einsum\n\"\"\"Alias for :func:`cotengra.einsum`.\"\"\"\n\n\n__all__ = (\n \"array_contract_expression\",\n \"array_contract_path\",\n \"array_contract_tree\",\n \"array_contract\",\n \"auto_hq_optimize\",\n \"auto_optimize\",\n \"AutoHQOptimizer\",\n \"AutoOptimizer\",\n \"contract_expression\",\n \"contract\",\n \"ContractionTree\",\n \"ContractionTreeCompressed\",\n \"ContractionTreeMulti\",\n \"edge_path_to_linear\",\n \"edge_path_to_ssa\",\n \"einsum_expression\",\n \"einsum_tree\",\n \"einsum\",\n \"FlowCutterOptimizer\",\n \"get_hyper_space\",\n \"get_hypergraph\",\n \"get_symbol_map\",\n \"get_symbol\",\n \"greedy_optimize\",\n \"GreedyOptimizer\",\n \"hash_contraction\",\n \"hyper_cmaes\",\n \"hyper_nevergrad\",\n \"hyper_neldermead\",\n \"hyper_optimize\",\n \"hyper_optuna\",\n \"hyper_random\",\n \"hyper_es\",\n \"hyper_skopt\",\n \"hyper_sbplx\",\n \"HyperCompressedOptimizer\",\n \"HyperGraph\",\n \"HyperMultiOptimizer\",\n \"HyperOptimizer\",\n \"linear_to_ssa\",\n \"list_hyper_functions\",\n \"ncon\",\n \"optimal_optimize\",\n \"optimal_outer_optimize\",\n \"OptimalOptimizer\",\n \"optimize_flowcutter\",\n \"optimize_quickbb\",\n \"path_basic\",\n \"path_compressed_greedy\",\n \"path_greedy\",\n \"path_igraph\",\n \"path_kahypar\",\n \"path_labels\",\n \"PathOptimizer\",\n \"plot_contractions_alt\",\n \"plot_contractions\",\n \"plot_scatter_alt\",\n \"plot_scatter\",\n \"plot_slicings_alt\",\n \"plot_slicings\",\n \"plot_tree_ring\",\n \"plot_tree_span\",\n \"plot_tree_tent\",\n \"plot_tree\",\n \"plot_trials_alt\",\n \"plot_trials\",\n \"QuasiRandOptimizer\",\n \"QuickBBOptimizer\",\n \"RandomGreedyOptimizer\",\n \"RandomOptimizer\",\n \"register_preset\",\n \"ReusableHyperCompressedOptimizer\",\n \"ReusableHyperOptimizer\",\n \"ReusableRandomGreedyOptimizer\",\n \"SliceFinder\",\n \"ssa_to_linear\",\n \"UniformOptimizer\",\n \"utils\",\n)\n\n\n# add some presets\n\n\ndef hyper_optimize(\n inputs,\n output,\n size_dict,\n memory_limit=None,\n get=\"path\",\n **opts,\n):\n if memory_limit is not None:\n warnings.warn(\n \"`memory_limit` is not supported in hyper_optimize, ignoring.\"\n )\n\n optimizer = HyperOptimizer(**opts)\n if get == \"path\":\n return optimizer(inputs, output, size_dict)\n elif get == \"tree\":\n return optimizer.search(inputs, output, size_dict)\n else:\n raise ValueError(f\"Unknown get option {get}\")\n\n\ndef hyper_compressed_optimize(\n inputs,\n output,\n size_dict,\n get=\"path\",\n **opts,\n):\n optimizer = HyperCompressedOptimizer(**opts)\n\n if get == \"path\":\n return optimizer(inputs, output, size_dict)\n elif get == \"tree\":\n return optimizer.search(inputs, output, size_dict)\n else:\n raise ValueError(f\"Unknown get option {get}\")\n\n\ndef random_greedy_optimize(\n inputs, output, size_dict, memory_limit=None, **opts\n):\n if memory_limit is not None:\n warnings.warn(\n \"`memory_limit` is not supported in \"\n \"random_greedy_optimize, ignoring.\"\n )\n\n optimizer = RandomGreedyOptimizer(**opts)\n return optimizer(inputs, output, size_dict)\n\n\ntry:\n register_preset(\n \"hyper\",\n hyper_optimize,\n optimizer_tree=functools.partial(hyper_optimize, get=\"tree\"),\n )\n register_preset(\n \"hyper-256\",\n functools.partial(hyper_optimize, max_repeats=256),\n optimizer_tree=functools.partial(\n hyper_optimize, max_repeats=256, get=\"tree\"\n ),\n )\n register_preset(\n \"hyper-greedy\",\n functools.partial(hyper_optimize, methods=[\"greedy\"]),\n optimizer_tree=functools.partial(\n hyper_optimize, methods=[\"greedy\"], get=\"tree\"\n ),\n )\n register_preset(\n \"hyper-labels\",\n functools.partial(hyper_optimize, methods=[\"labels\"]),\n optimizer_tree=functools.partial(\n hyper_optimize, methods=[\"labels\"], get=\"tree\"\n ),\n )\n register_preset(\n \"hyper-kahypar\",\n functools.partial(hyper_optimize, methods=[\"kahypar\"]),\n optimizer_tree=functools.partial(\n hyper_optimize, methods=[\"kahypar\"], get=\"tree\"\n ),\n )\n register_preset(\n \"hyper-balanced\",\n functools.partial(\n hyper_optimize, methods=[\"kahypar-balanced\"], max_repeats=16\n ),\n optimizer_tree=functools.partial(\n hyper_optimize,\n methods=[\"kahypar-balanced\"],\n max_repeats=16,\n get=\"tree\",\n ),\n )\n register_preset(\n \"hyper-compressed\",\n hyper_compressed_optimize,\n optimizer_tree=functools.partial(\n hyper_compressed_optimize,\n get=\"tree\",\n ),\n compressed=True,\n )\n register_preset(\n \"hyper-spinglass\",\n functools.partial(hyper_optimize, methods=[\"spinglass\"]),\n )\n register_preset(\n \"hyper-betweenness\",\n functools.partial(hyper_optimize, methods=[\"betweenness\"]),\n )\n register_preset(\n \"random-greedy\",\n random_greedy_optimize,\n )\n register_preset(\n \"random-greedy-128\",\n functools.partial(random_greedy_optimize, max_repeats=128),\n )\n register_preset(\n \"flowcutter-2\",\n functools.partial(optimize_flowcutter, max_time=2),\n )\n register_preset(\n \"flowcutter-10\",\n functools.partial(optimize_flowcutter, max_time=10),\n )\n register_preset(\n \"flowcutter-60\",\n functools.partial(optimize_flowcutter, max_time=60),\n )\n register_preset(\n \"quickbb-2\",\n functools.partial(optimize_quickbb, max_time=2),\n )\n register_preset(\n \"quickbb-10\",\n functools.partial(optimize_quickbb, max_time=10),\n )\n register_preset(\n \"quickbb-60\",\n functools.partial(optimize_quickbb, max_time=60),\n )\n register_preset(\n \"greedy-compressed\",\n path_compressed_greedy.greedy_compressed,\n path_compressed_greedy.trial_greedy_compressed,\n compressed=True,\n )\n register_preset(\n \"greedy-span\",\n path_compressed_greedy.greedy_span,\n path_compressed_greedy.trial_greedy_span,\n compressed=True,\n )\nexcept KeyError:\n # KeyError: if reloading cotengra e.g. library entries already registered\n pass\n"
},
{
"path": "cotengra/contract.py",
"content": "\"\"\"Functionality relating to actually contracting.\"\"\"\n\nimport contextlib\nimport functools\nimport itertools\nimport operator\n\nfrom autoray import do, get_namespace, infer_backend_multi, shape\n\nDEFAULT_IMPLEMENTATION = \"auto\"\n\n\ndef set_default_implementation(impl):\n global DEFAULT_IMPLEMENTATION\n DEFAULT_IMPLEMENTATION = impl\n\n\ndef get_default_implementation():\n return DEFAULT_IMPLEMENTATION\n\n\n@contextlib.contextmanager\ndef default_implementation(impl):\n \"\"\"Context manager for temporarily setting the default implementation.\"\"\"\n global DEFAULT_IMPLEMENTATION\n old_impl = DEFAULT_IMPLEMENTATION\n DEFAULT_IMPLEMENTATION = impl\n try:\n yield\n finally:\n DEFAULT_IMPLEMENTATION = old_impl\n\n\n@functools.lru_cache(2**12)\ndef _sanitize_equation(eq):\n \"\"\"Get the input and output indices of an equation, computing the output\n implicitly as the sorted sequence of every index that appears exactly once\n if it is not provided.\n \"\"\"\n # remove spaces\n eq = eq.replace(\" \", \"\")\n\n if \"...\" in eq:\n raise NotImplementedError(\"Ellipsis not supported.\")\n\n if \"->\" not in eq:\n lhs = eq\n tmp_subscripts = lhs.replace(\",\", \"\")\n out = \"\".join(\n # sorted sequence of indices\n s\n for s in sorted(set(tmp_subscripts))\n # that appear exactly once\n if tmp_subscripts.count(s) == 1\n )\n else:\n lhs, out = eq.split(\"->\")\n return lhs, out\n\n\n@functools.lru_cache(2**12)\ndef _parse_einsum_single(eq, shape):\n \"\"\"Cached parsing of a single term einsum equation into the necessary\n sequence of arguments for axes diagonals, sums, and transposes.\n \"\"\"\n lhs, out = _sanitize_equation(eq)\n\n # parse each index\n need_to_diag = []\n need_to_sum = []\n seen = set()\n for ix in lhs:\n if ix in need_to_diag:\n continue\n if ix in seen:\n need_to_diag.append(ix)\n continue\n seen.add(ix)\n if ix not in out:\n need_to_sum.append(ix)\n\n # first handle diagonal reductions\n if need_to_diag:\n diag_sels = []\n sizes = dict(zip(lhs, shape))\n while need_to_diag:\n ixd = need_to_diag.pop()\n dinds = tuple(range(sizes[ixd]))\n\n # construct advanced indexing object\n selector = tuple(dinds if ix == ixd else slice(None) for ix in lhs)\n diag_sels.append(selector)\n\n # after taking the diagonal what are new indices?\n ixd_contig = ixd * lhs.count(ixd)\n if ixd_contig in lhs:\n # contig axes, new axis is at same position\n lhs = lhs.replace(ixd_contig, ixd)\n else:\n # non-contig, new axis is at beginning\n lhs = ixd + lhs.replace(ixd, \"\")\n else:\n diag_sels = None\n\n # then sum reductions\n if need_to_sum:\n sum_axes = tuple(map(lhs.index, need_to_sum))\n for ix in need_to_sum:\n lhs = lhs.replace(ix, \"\")\n else:\n sum_axes = None\n\n # then transposition\n if lhs == out:\n perm = None\n else:\n perm = tuple(lhs.index(ix) for ix in out)\n\n return diag_sels, sum_axes, perm\n\n\ndef _parse_eq_to_pure_multiplication(a_term, shape_a, b_term, shape_b, out):\n \"\"\"If there are no contracted indices, then we can directly transpose and\n insert singleton dimensions into ``a`` and ``b`` such that (broadcast)\n elementwise multiplication performs the einsum.\n\n No need to cache this as it is within the cached\n ``_parse_eq_to_batch_matmul``.\n\n \"\"\"\n desired_a = \"\"\n desired_b = \"\"\n new_shape_a = []\n new_shape_b = []\n for ix in out:\n if ix in a_term:\n desired_a += ix\n new_shape_a.append(shape_a[a_term.index(ix)])\n else:\n new_shape_a.append(1)\n if ix in b_term:\n desired_b += ix\n new_shape_b.append(shape_b[b_term.index(ix)])\n else:\n new_shape_b.append(1)\n\n if desired_a != a_term:\n eq_a = f\"{a_term}->{desired_a}\"\n else:\n eq_a = None\n if desired_b != b_term:\n eq_b = f\"{b_term}->{desired_b}\"\n else:\n eq_b = None\n\n return (\n eq_a,\n eq_b,\n new_shape_a,\n new_shape_b,\n None, # new_shape_ab, not needed since not fusing\n None, # perm_ab, not needed as we transpose a and b first\n True, # pure_multiplication=True\n )\n\n\n@functools.lru_cache(2**12)\ndef _parse_eq_to_batch_matmul(eq, shape_a, shape_b):\n \"\"\"Cached parsing of a two term einsum equation into the necessary\n sequence of arguments for contracttion via batched matrix multiplication.\n The steps we need to specify are:\n\n 1. Remove repeated and trivial indices from the left and right terms,\n and transpose them, done as a single einsum.\n 2. Fuse the remaining indices so we have two 3D tensors.\n 3. Perform the batched matrix multiplication.\n 4. Unfuse the output to get the desired final index order.\n\n \"\"\"\n lhs, out = eq.split(\"->\")\n a_term, b_term = lhs.split(\",\")\n\n if len(a_term) != len(shape_a):\n raise ValueError(f\"Term '{a_term}' does not match shape {shape_a}.\")\n if len(b_term) != len(shape_b):\n raise ValueError(f\"Term '{b_term}' does not match shape {shape_b}.\")\n\n sizes = {}\n singletons = set()\n\n # parse left term to unique indices with size > 1\n left = {}\n for ix, d in zip(a_term, shape_a):\n if d == 1:\n # everything (including broadcasting) works nicely if simply ignore\n # such dimensions, but we do need to track if they appear in output\n # and thus should be reintroduced later\n singletons.add(ix)\n continue\n if sizes.setdefault(ix, d) != d:\n # set and check size\n raise ValueError(\n f\"Index {ix} has mismatched sizes {sizes[ix]} and {d}.\"\n )\n left[ix] = True\n\n # parse right term to unique indices with size > 1\n right = {}\n for ix, d in zip(b_term, shape_b):\n # broadcast indices (size 1 on one input and size != 1\n # on the other) should not be treated as singletons\n if d == 1:\n if ix not in left:\n singletons.add(ix)\n continue\n singletons.discard(ix)\n\n if sizes.setdefault(ix, d) != d:\n # set and check size\n raise ValueError(\n f\"Index {ix} has mismatched sizes {sizes[ix]} and {d}.\"\n )\n right[ix] = True\n\n # now we classify the unique size > 1 indices only\n bat_inds = [] # appears on A, B, O\n con_inds = [] # appears on A, B, .\n a_keep = [] # appears on A, ., O\n b_keep = [] # appears on ., B, O\n # other indices (appearing on A or B only) will\n # be summed or traced out prior to the matmul\n for ix in left:\n if right.pop(ix, False):\n if ix in out:\n bat_inds.append(ix)\n else:\n con_inds.append(ix)\n elif ix in out:\n a_keep.append(ix)\n # now only indices unique to right remain\n for ix in right:\n if ix in out:\n b_keep.append(ix)\n\n if not con_inds:\n # contraction is pure multiplication, prepare inputs differently\n return _parse_eq_to_pure_multiplication(\n a_term, shape_a, b_term, shape_b, out\n )\n\n # only need the size one indices that appear in the output\n singletons = [ix for ix in out if ix in singletons]\n\n # take diagonal, remove any trivial axes and transpose left\n desired_a = \"\".join((*bat_inds, *a_keep, *con_inds))\n if a_term != desired_a:\n if set(a_term) == set(desired_a):\n # only need to transpose, don't invoke einsum\n eq_a = tuple(a_term.index(ix) for ix in desired_a)\n else:\n eq_a = f\"{a_term}->{desired_a}\"\n else:\n eq_a = None\n\n # take diagonal, remove any trivial axes and transpose right\n desired_b = \"\".join((*bat_inds, *con_inds, *b_keep))\n if b_term != desired_b:\n if set(b_term) == set(desired_b):\n # only need to transpose, don't invoke einsum\n eq_b = tuple(b_term.index(ix) for ix in desired_b)\n else:\n eq_b = f\"{b_term}->{desired_b}\"\n else:\n eq_b = None\n\n # then we want to reshape\n if bat_inds:\n lgroups = (bat_inds, a_keep, con_inds)\n rgroups = (bat_inds, con_inds, b_keep)\n ogroups = (bat_inds, a_keep, b_keep)\n else:\n # avoid size 1 batch dimension if no batch indices\n lgroups = (a_keep, con_inds)\n rgroups = (con_inds, b_keep)\n ogroups = (a_keep, b_keep)\n\n if any(len(group) != 1 for group in lgroups):\n # need to fuse 'kept' and contracted indices\n # (though could allow batch indices to be broadcast)\n new_shape_a = tuple(\n functools.reduce(operator.mul, (sizes[ix] for ix in ix_group), 1)\n for ix_group in lgroups\n )\n else:\n new_shape_a = None\n\n if any(len(group) != 1 for group in rgroups):\n # need to fuse 'kept' and contracted indices\n # (though could allow batch indices to be broadcast)\n new_shape_b = tuple(\n functools.reduce(operator.mul, (sizes[ix] for ix in ix_group), 1)\n for ix_group in rgroups\n )\n else:\n new_shape_b = None\n\n if any(len(group) != 1 for group in ogroups) or singletons:\n new_shape_ab = (1,) * len(singletons) + tuple(\n sizes[ix] for ix_group in ogroups for ix in ix_group\n )\n else:\n new_shape_ab = None\n\n # then we might need to permute the matmul produced output:\n out_produced = \"\".join((*singletons, *bat_inds, *a_keep, *b_keep))\n if out_produced != out:\n perm_ab = tuple(out_produced.index(ix) for ix in out)\n else:\n perm_ab = None\n\n return (\n eq_a,\n eq_b,\n new_shape_a,\n new_shape_b,\n new_shape_ab,\n perm_ab,\n False, # pure_multiplication=False\n )\n\n\ndef _einsum_single(eq, x, backend=None):\n \"\"\"Einsum on a single tensor, via three steps: diagonal selection\n (via advanced indexing), axes summations, transposition. The logic for each\n is cached based on the equation and array shape, and each step is only\n performed if necessary.\n \"\"\"\n try:\n return do(\"einsum\", eq, x, like=backend)\n except ImportError:\n pass\n\n diag_sels, sum_axes, perm = _parse_einsum_single(eq, shape(x))\n\n if diag_sels is not None:\n # diagonal reduction via advanced indexing\n # e.g ababbac->abc\n for selector in diag_sels:\n x = x[selector]\n\n if sum_axes is not None:\n # trivial removal of axes via summation\n # e.g. abc->c\n x = do(\"sum\", x, sum_axes, like=backend)\n\n if perm is not None:\n # transpose to desired output\n # e.g. abc->cba\n x = do(\"transpose\", x, perm, like=backend)\n\n return x\n\n\ndef _do_contraction_via_bmm(\n a,\n b,\n eq_a,\n eq_b,\n new_shape_a,\n new_shape_b,\n new_shape_ab,\n perm_ab,\n pure_multiplication,\n backend,\n):\n # prepare left\n if eq_a is not None:\n if isinstance(eq_a, tuple):\n # only transpose\n a = do(\"transpose\", a, eq_a, like=backend)\n else:\n # diagonals, sums, and tranpose\n a = _einsum_single(eq_a, a)\n if new_shape_a is not None:\n a = do(\"reshape\", a, new_shape_a, like=backend)\n\n # prepare right\n if eq_b is not None:\n if isinstance(eq_b, tuple):\n # only transpose\n b = do(\"transpose\", b, eq_b, like=backend)\n else:\n # diagonals, sums, and tranpose\n b = _einsum_single(eq_b, b)\n if new_shape_b is not None:\n b = do(\"reshape\", b, new_shape_b, like=backend)\n\n if pure_multiplication:\n # no contracted indices\n return do(\"multiply\", a, b)\n\n # do the contraction!\n ab = do(\"matmul\", a, b, like=backend)\n\n # prepare the output\n if new_shape_ab is not None:\n ab = do(\"reshape\", ab, new_shape_ab, like=backend)\n if perm_ab is not None:\n ab = do(\"transpose\", ab, perm_ab, like=backend)\n\n return ab\n\n\ndef einsum(eq, a, b=None, *, backend=None):\n \"\"\"Perform arbitrary single and pairwise einsums using only `matmul`,\n `transpose`, `reshape` and `sum`. The logic for each is cached based on\n the equation and array shape, and each step is only performed if necessary.\n\n Parameters\n ----------\n eq : str\n The einsum equation.\n a : array_like\n The first array to contract.\n b : array_like, optional\n The second array to contract.\n backend : str, optional\n The backend to use for array operations. If ``None``, dispatch\n automatically based on ``a`` and ``b``.\n\n Returns\n -------\n array_like\n \"\"\"\n if b is None:\n return _einsum_single(eq, a, backend=backend)\n\n (\n eq_a,\n eq_b,\n new_shape_a,\n new_shape_b,\n new_shape_ab,\n perm_ab,\n pure_multiplication,\n ) = _parse_eq_to_batch_matmul(eq, shape(a), shape(b))\n\n return _do_contraction_via_bmm(\n a,\n b,\n eq_a,\n eq_b,\n new_shape_a,\n new_shape_b,\n new_shape_ab,\n perm_ab,\n pure_multiplication,\n backend,\n )\n\n\ndef gen_nice_inds():\n \"\"\"Generate the indices from [a-z, A-Z, reasonable unicode...].\"\"\"\n for i in range(26):\n yield chr(ord(\"a\") + i)\n for i in range(26):\n yield chr(ord(\"A\") + i)\n for i in itertools.count(192):\n yield chr(i)\n\n\n@functools.lru_cache(2**12)\ndef _parse_tensordot_axes_to_matmul(axes, shape_a, shape_b):\n \"\"\"Parse a tensordot specification into the necessary sequence of arguments\n for contracttion via matrix multiplication. This just converts ``axes``\n into an ``einsum`` eq string then calls ``_parse_eq_to_batch_matmul``.\n \"\"\"\n ndim_a = len(shape_a)\n ndim_b = len(shape_b)\n\n if isinstance(axes, int):\n axes_a = tuple(range(ndim_a - axes, ndim_a))\n axes_b = tuple(range(axes))\n else:\n axes_a, axes_b = axes\n\n num_con = len(axes_a)\n if num_con != len(axes_b):\n raise ValueError(\n f\"Axes should have the same length, got {axes_a} and {axes_b}.\"\n )\n\n possible_inds = gen_nice_inds()\n inds_a = [next(possible_inds) for _ in range(ndim_a)]\n inds_b = []\n inds_out = inds_a.copy()\n\n for axb in range(ndim_b):\n if axb not in axes_b:\n # right uncontracted index\n ind = next(possible_inds)\n inds_out.append(ind)\n else:\n # contracted index\n axa = axes_a[axes_b.index(axb)]\n # check that the shapes match\n if shape_a[axa] != shape_b[axb]:\n raise ValueError(\n f\"Dimension mismatch between axes {axa} of {shape_a} and \"\n f\"{axb} of {shape_b}: {shape_a[axa]} != {shape_b[axb]}.\"\n )\n ind = inds_a[axa]\n inds_out.remove(ind)\n inds_b.append(ind)\n\n eq = f\"{''.join(inds_a)},{''.join(inds_b)}->{''.join(inds_out)}\"\n\n return _parse_eq_to_batch_matmul(eq, shape_a, shape_b)\n\n\ndef tensordot(a, b, axes=2, *, backend=None):\n \"\"\"Perform a tensordot using only `matmul`, `transpose`, `reshape`. The\n logic for each is cached based on the equation and array shape, and each\n step is only performed if necessary.\n\n Parameters\n ----------\n a, b : array_like\n The arrays to contract.\n axes : int or tuple of (sequence[int], sequence[int])\n The number of axes to contract, or the axes to contract. If an int,\n the last ``axes`` axes of ``a`` and the first ``axes`` axes of ``b``\n are contracted. If a tuple, the axes to contract for ``a`` and ``b``\n respectively.\n backend : str or None, optional\n The backend to use for array operations. If ``None``, dispatch\n automatically based on ``a`` and ``b``.\n\n Returns\n -------\n array_like\n \"\"\"\n try:\n # ensure hashable\n axes = tuple(map(int, axes[0])), tuple(map(int, axes[1]))\n except IndexError:\n axes = int(axes)\n\n (\n eq_a,\n eq_b,\n new_shape_a,\n new_shape_b,\n new_shape_ab,\n perm_ab,\n pure_multiplication,\n ) = _parse_tensordot_axes_to_matmul(axes, shape(a), shape(b))\n\n return _do_contraction_via_bmm(\n a,\n b,\n eq_a,\n eq_b,\n new_shape_a,\n new_shape_b,\n new_shape_ab,\n perm_ab,\n pure_multiplication,\n backend,\n )\n\n\ndef extract_contractions(\n tree,\n order=None,\n prefer_einsum=False,\n):\n \"\"\"Extract just the information needed to perform the contraction.\n\n Parameters\n ----------\n order : str or callable, optional\n Supplied to :meth:`ContractionTree.traverse`.\n prefer_einsum : bool, optional\n Prefer to use ``einsum`` for pairwise contractions, even if\n ``tensordot`` can perform the contraction.\n\n Returns\n -------\n contractions : tuple\n A tuple of tuples, each containing the information needed to\n perform a pairwise contraction. Each tuple contains:\n\n - ``p``: the parent node,\n - ``l``: the left child node,\n - ``r``: the right child node,\n - ``tdot``: whether to use ``tensordot`` or ``einsum``,\n - ``arg``: the argument to pass to ``tensordot`` or ``einsum``\n i.e. ``axes`` or ``eq``,\n - ``perm``: the permutation required after the contraction, if\n any (only applies to tensordot).\n\n If both ``l`` and ``r`` are ``None``, the the operation is a single\n term simplification performed with ``einsum``.\n \"\"\"\n if tree.N == 1:\n # trivial 'contraction', single input maps directly to output\n # possibly with reductions/transpose, setting l but not r flags this\n pi = 1\n li = 0\n ri = None\n tdot = False\n arg = tree.get_eq_sliced()\n perm = None\n return [(pi, li, ri, tdot, arg, perm)]\n\n contractions = []\n\n # for compactness we convert nodes to ssa indices\n ssas = {leaf: i for i, leaf in enumerate(tree.gen_leaves())}\n ssa = len(ssas)\n\n # pairwise contractions\n for p, l, r in tree.traverse(order=order):\n li = ssas.pop(l)\n ri = ssas.pop(r)\n pi = ssas[p] = ssa\n ssa += 1\n\n if prefer_einsum or not tree.get_can_dot(p):\n tdot = False\n arg = tree.get_einsum_eq(p)\n perm = None\n else:\n tdot = True\n arg = tree.get_tensordot_axes(p)\n perm = tree.get_tensordot_perm(p)\n\n contractions.append((pi, li, ri, tdot, arg, perm))\n\n if tree.preprocessing:\n # inplace single term simplifications\n # n.b. these are populated lazily when the other information is\n # computed above, so we do it after\n pre_contractions = (\n (i, None, None, False, eq, None)\n for i, eq in tree.preprocessing.items()\n )\n return (*pre_contractions, *contractions)\n\n return tuple(contractions)\n\n\nclass Contractor:\n \"\"\"Default cotengra network contractor.\n\n Parameters\n ----------\n contractions : tuple[tuple]\n The sequence of contractions to perform. Each contraction should be a\n tuple containing:\n\n - ``p``: the parent node,\n - ``l``: the left child node,\n - ``r``: the right child node,\n - ``tdot``: whether to use ``tensordot`` or ``einsum``,\n - ``arg``: the argument to pass to ``tensordot`` or ``einsum``\n i.e. ``axes`` or ``eq``,\n - ``perm``: the permutation required after the contraction, if\n any (only applies to tensordot).\n\n e.g. built by calling ``extract_contractions(tree)``.\n\n strip_exponent : bool, optional\n If ``True``, eagerly strip the exponent (in log10) from\n intermediate tensors to control numerical problems from leaving the\n range of the datatype. This method then returns the scaled\n 'mantissa' output array and the exponent separately.\n check_zero : bool, optional\n If ``True``, when ``strip_exponent=True``, explicitly check for\n zero-valued intermediates that would otherwise produce ``nan``,\n instead terminating early if encounteredand returning\n ``(0.0, 0.0)``.\n backend : str, optional\n What library to use for ``tensordot``, ``einsum`` and\n ``transpose``, it will be automatically inferred from the input\n arrays if not given.\n progbar : bool, optional\n Whether to show a progress bar.\n \"\"\"\n\n __slots__ = (\n \"contractions\",\n \"strip_exponent\",\n \"check_zero\",\n \"implementation\",\n \"backend\",\n \"progbar\",\n \"__weakref__\",\n )\n\n def __init__(\n self,\n contractions,\n strip_exponent=False,\n check_zero=False,\n implementation=\"auto\",\n backend=None,\n progbar=False,\n ):\n self.contractions = contractions\n self.strip_exponent = strip_exponent\n self.check_zero = check_zero\n self.implementation = implementation\n self.backend = backend\n self.progbar = progbar\n\n def __call__(self, *arrays, **kwargs):\n \"\"\"Contract ``arrays`` using operations listed in ``contractions``.\n\n Parameters\n ----------\n arrays : sequence of array-like\n The arrays to contract.\n kwargs : dict\n Override the default settings for this contraction only.\n\n Returns\n -------\n output : array\n The contracted output, it will be scaled if ``strip_exponent==True``.\n exponent : float\n The exponent of the output in base 10, returned only if\n ``strip_exponent==True``.\n \"\"\"\n backend = kwargs.pop(\"backend\", self.backend)\n progbar = kwargs.pop(\"progbar\", self.progbar)\n check_zero = kwargs.pop(\"check_zero\", self.check_zero)\n strip_exponent = kwargs.pop(\"strip_exponent\", self.strip_exponent)\n implementation = kwargs.pop(\"implementation\", self.implementation)\n if kwargs:\n raise TypeError(f\"Unknown keyword arguments: {kwargs}.\")\n\n if backend is None:\n backend = infer_backend_multi(*arrays)\n xp = get_namespace(backend)\n\n if implementation == \"auto\":\n if backend == \"numpy\":\n # by default only replace numpy's einsum/tensordot\n implementation = \"cotengra\"\n else:\n implementation = \"autoray\"\n\n if implementation == \"cotengra\":\n _einsum, _tensordot = einsum, tensordot\n elif implementation == \"pytblis\":\n import pytblis\n\n _einsum = pytblis.einsum\n _tensordot = pytblis.tensordot\n elif implementation == \"autoray\":\n try:\n _einsum = xp.einsum\n except ImportError:\n # fallback to cotengra (matmul) implementation\n _einsum = einsum\n\n try:\n _tensordot = xp.tensordot\n except ImportError:\n # fallback to cotengra (matmul) implementation\n _tensordot = tensordot\n else:\n # manually supplied\n _einsum, _tensordot = implementation\n\n # temporary storage for intermediates\n N = len(arrays)\n temps = dict(enumerate(arrays))\n\n exponent = 0.0 if (strip_exponent is not False) else None\n\n if progbar:\n import tqdm\n\n contractions = tqdm.tqdm(self.contractions, total=N - 1)\n else:\n contractions = self.contractions\n\n for pi, li, ri, tdot, arg, perm in contractions:\n if ri is None:\n if li is None:\n # single term simplification, perform inplace with einsum\n temps[pi] = _einsum(arg, temps[pi])\n continue\n else:\n # trivial 'contraction', single input maps directly to\n # output, possibly with reductions/transpose via einsum\n p_array = _einsum(arg, temps[li])\n if strip_exponent:\n return p_array, 0.0\n return p_array\n\n # get input arrays for this contraction\n l_array = temps.pop(li)\n r_array = temps.pop(ri)\n\n if tdot:\n p_array = _tensordot(l_array, r_array, arg)\n if perm:\n p_array = xp.transpose(p_array, perm)\n else:\n p_array = _einsum(arg, l_array, r_array)\n\n if exponent is not None:\n factor = xp.max(xp.abs(p_array))\n\n if check_zero and float(factor) == 0.0:\n return 0.0, float(\"-inf\")\n exponent = exponent + xp.log10(factor)\n\n if backend == \"tensorflow\":\n factor = xp.astype(factor, p_array.dtype)\n # TODO:\n # currently special case tensorflow\n # autoray needs fix for autojit and astype to use generally\n\n p_array = p_array / factor\n\n # insert the new intermediate array\n temps[pi] = p_array\n\n if exponent is not None:\n return p_array, exponent\n\n return p_array\n\n\nclass CuQuantumContractor:\n def __init__(\n self,\n tree,\n handle_slicing=False,\n autotune=False,\n **kwargs,\n ):\n if kwargs.pop(\"strip_exponent\", None):\n raise ValueError(\n \"strip_exponent=True not supported with cuQuantum\"\n )\n\n if tree.has_preprocessing():\n raise ValueError(\"Preprocessing not supported with cuQuantum yet.\")\n\n if kwargs.pop(\"progbar\", None):\n import warnings\n\n warnings.warn(\"Progress bar not supported with cuQuantum yet.\")\n\n if handle_slicing:\n self.eq = tree.get_eq()\n self.shapes = tree.get_shapes()\n else:\n self.eq = tree.get_eq_sliced()\n self.shapes = tree.get_shapes_sliced()\n\n if tree.is_complete():\n kwargs.setdefault(\"optimize\", {})\n kwargs[\"optimize\"].setdefault(\"path\", tree.get_path())\n\n if handle_slicing and tree.sliced_inds:\n kwargs[\"optimize\"].setdefault(\n \"slicing\",\n [(ix, tree.size_dict[ix] - 1) for ix in tree.sliced_inds],\n )\n\n self.kwargs = kwargs\n self.autotune = 3 if autotune is True else autotune\n self.handle = None\n self.network = None\n\n def setup(self, *arrays):\n import cuquantum\n\n if hasattr(cuquantum, \"bindings\"):\n # cuquantum-python >= 25.03\n from cuquantum.tensornet import Network\n else:\n # for cuquantum < 25.03\n from cuquantum import Network\n\n self.network = Network(\n self.eq,\n *arrays,\n )\n self.network.contract_path(**self.kwargs)\n if self.autotune:\n self.network.autotune(iterations=self.autotune)\n\n def __call__(\n self,\n *arrays,\n check_zero=False,\n backend=None,\n progbar=False,\n ):\n # can't handle these yet\n assert not check_zero\n assert not progbar\n assert backend is None\n\n if self.network is None:\n self.setup(*arrays)\n else:\n self.network.reset_operands(*arrays)\n\n return self.network.contract()\n\n def __del__(self):\n if self.network is not None:\n self.network.free()\n\n\ndef make_contractor(\n tree,\n order=None,\n prefer_einsum=False,\n strip_exponent=False,\n check_zero=False,\n implementation=None,\n autojit=False,\n progbar=False,\n):\n \"\"\"Get a reusable function which performs the contraction corresponding\n to ``tree``. The various options provide defaults that can also be overrode\n when calling the standard contractor.\n\n Parameters\n ----------\n tree : ContractionTree\n The contraction tree.\n order : str or callable, optional\n Supplied to :meth:`ContractionTree.traverse`, the order in which\n to perform the pairwise contractions given by the tree.\n prefer_einsum : bool, optional\n Prefer to use ``einsum`` for pairwise contractions, even if\n ``tensordot`` can perform the contraction.\n strip_exponent : bool, optional\n If ``True``, the function will strip the exponent from the output\n array and return it separately.\n check_zero : bool, optional\n If ``True``, when ``strip_exponent=True``, explicitly check for\n zero-valued intermediates that would otherwise produce ``nan``,\n instead terminating early if encountered and returning\n ``(0.0, 0.0)``.\n implementation : str or tuple[callable, callable], optional\n What library to use to actually perform the contractions. Options are\n\n - \"auto\": let cotengra choose\n - \"autoray\": dispatch with autoray, using the ``tensordot`` and\n ``einsum`` implementation of the backend\n - \"cotengra\": use the ``tensordot`` and ``einsum`` implementation of\n cotengra, which is based on batch matrix multiplication. This is\n faster for some backends like numpy, and also enables libraries\n which don't yet provide ``tensordot`` and ``einsum`` to be used.\n - \"cuquantum\": use the cuquantum library to perform the whole\n contraction (not just individual contractions).\n - tuple[callable, callable]: manually supply the ``tensordot`` and\n ``einsum`` implementations to use.\n\n autojit : bool, optional\n If ``True``, use :func:`autoray.autojit` to compile the contraction\n function.\n progbar : bool, optional\n Whether to show progress through the contraction by default.\n\n Returns\n -------\n fn : callable\n The contraction function, with signature ``fn(*arrays)``.\n \"\"\"\n if implementation is None:\n implementation = get_default_implementation()\n\n if implementation == \"cuquantum\":\n fn = CuQuantumContractor(\n tree,\n strip_exponent=strip_exponent,\n check_zero=check_zero,\n progbar=progbar,\n )\n else:\n fn = Contractor(\n contractions=extract_contractions(tree, order, prefer_einsum),\n strip_exponent=strip_exponent,\n check_zero=check_zero,\n implementation=implementation,\n progbar=progbar,\n )\n if autojit:\n from autoray import autojit as _autojit\n\n fn = _autojit(fn)\n\n return fn\n"
},
{
"path": "cotengra/core.py",
"content": "\"\"\"Core contraction tree data structure and methods.\"\"\"\n\nimport collections\nimport functools\nimport itertools\nimport math\nimport warnings\nfrom dataclasses import dataclass\nfrom typing import Optional\n\nfrom autoray import do, get_namespace, infer_backend\n\nfrom .contract import make_contractor\nfrom .hypergraph import get_hypergraph\nfrom .nodeops import get_nodeops\nfrom .parallel import (\n can_scatter,\n maybe_leave_pool,\n maybe_rejoin_pool,\n parse_parallel_arg,\n scatter,\n submit,\n)\nfrom .pathfinders.path_simulated_annealing import (\n parallel_temper_tree,\n simulated_anneal_tree,\n)\nfrom .plot import (\n plot_contractions,\n plot_contractions_alt,\n plot_hypergraph,\n plot_tree_circuit,\n plot_tree_flat,\n plot_tree_ring,\n plot_tree_rubberband,\n plot_tree_span,\n plot_tree_tent,\n)\nfrom .scoring import (\n DEFAULT_COMBO_FACTOR,\n CompressedStatsTracker,\n get_score_fn,\n)\nfrom .utils import (\n MaxCounter,\n compute_size_by_dict,\n deprecated,\n get_rng,\n get_symbol,\n groupby,\n inputs_output_to_eq,\n interleave,\n oset,\n prod,\n unique,\n)\n\n\ndef cached_node_property(name):\n \"\"\"Decorator for caching information about nodes.\"\"\"\n\n def wrapper(meth):\n @functools.wraps(meth)\n def getter(self, node):\n try:\n return self.info[node][name]\n except KeyError:\n self.info[node][name] = value = meth(self, node)\n return value\n\n return getter\n\n return wrapper\n\n\ndef legs_union(legs_seq):\n \"\"\"Combine a sequence of legs into a single set of legs, summing their\n appearances.\n \"\"\"\n new_legs, *rem_legs = legs_seq\n new_legs = new_legs.copy()\n for legs in rem_legs:\n for ix, ix_count in legs.items():\n new_legs[ix] = new_legs.get(ix, 0) + ix_count\n return new_legs\n\n\ndef legs_without(legs, ind):\n \"\"\"Discard ``ind`` from legs to create a new set of legs.\"\"\"\n new_legs = legs.copy()\n new_legs.pop(ind, None)\n return new_legs\n\n\ndef get_with_default(k, obj, default):\n return obj.get(k, default)\n\n\n@dataclass(order=True, frozen=True)\nclass SliceInfo:\n inner: bool\n ind: str\n size: int\n project: Optional[int]\n\n @property\n def sliced_range(self):\n if self.project is None:\n return range(self.size)\n else:\n return [self.project]\n\n\ndef get_slice_strides(sliced_inds):\n \"\"\"Compute the 'strides' given the (ordered) dictionary of sliced indices.\"\"\"\n slice_infos = list(sliced_inds.values())\n nsliced = len(slice_infos)\n strides = [1] * nsliced\n # backwards cumulative product\n for i in range(nsliced - 2, -1, -1):\n strides[i] = strides[i + 1] * slice_infos[i + 1].size\n return strides\n\n\nclass AdderWithMaybeExponentStripped:\n \"\"\"Object that ddds two arrays, or tuples of (array, exponent) together in\n a stable and branchless way. It also internally caches the backend on the\n first call.\n \"\"\"\n\n __slots__ = (\"backend\", \"namespace\", \"need_to_cast\")\n\n def __init__(self):\n self.backend = None\n self.namespace = None\n self.need_to_cast = False\n\n def __call__(self, x, y):\n xistup = isinstance(x, tuple)\n yistup = isinstance(y, tuple)\n if not (xistup or yistup):\n # simple sum without exponent\n return x + y\n\n if xistup:\n xm, xe = x\n else:\n xm = x\n xe = 0.0\n\n if yistup:\n ym, ye = y\n else:\n ym = y\n ye = 0.0\n\n if self.backend is None:\n self.backend = infer_backend(xm)\n self.namespace = get_namespace(self.backend)\n self.need_to_cast = self.backend == \"tensorflow\"\n\n # perform branchless for jit etc.\n e = max(xe, ye)\n\n if self.need_to_cast:\n xcoeff = self.namespace.astype(10.0 ** (xe - e), xm.dtype)\n ycoeff = self.namespace.astype(10.0 ** (ye - e), ym.dtype)\n m = xm * xcoeff + ym * ycoeff\n else:\n m = xm * 10 ** (xe - e) + ym * 10 ** (ye - e)\n\n return (m, e)\n\n\nclass ContractionTree:\n \"\"\"Binary tree representing a tensor network contraction.\n\n Parameters\n ----------\n inputs : sequence of str\n The list of input tensor's indices.\n output : str\n The output indices.\n size_dict : dict[str, int]\n The size of each index.\n track_childless : bool, optional\n Whether to dynamically keep track of which nodes are childless. Useful\n if you are 'divisively' building the tree.\n track_flops : bool, optional\n Whether to dynamically keep track of the total number of flops. If\n ``False`` You can still compute this once the tree is complete.\n track_write : bool, optional\n Whether to dynamically keep track of the total number of elements\n written. If ``False`` You can still compute this once the tree is\n complete.\n track_size : bool, optional\n Whether to dynamically keep track of the largest tensor so far. If\n ``False`` You can still compute this once the tree is complete.\n objective : str or Objective, optional\n An default objective function to use for further optimization and\n scoring, for example reconfiguring or computing the combo cost. If not\n supplied the default is to create a flops objective when needed.\n\n Attributes\n ----------\n children : dict[node, tuple[node]]\n Mapping of each node to two children.\n info : dict[node, dict]\n Information about the tree nodes. The key is the set of inputs (a\n set of inputs indices) the node contains. Or in other words, the\n subgraph of the node. The value is a dictionary to cache information\n about effective 'leg' indices, size, flops of formation etc.\n \"\"\"\n\n def __init__(\n self,\n inputs,\n output,\n size_dict,\n track_childless=False,\n track_flops=False,\n track_write=False,\n track_size=False,\n objective=None,\n nodeops=\"auto\",\n ):\n self.inputs = inputs\n self.output = output\n\n if isinstance(self.inputs[0], set) or isinstance(self.output, set):\n warnings.warn(\n \"The inputs or output of this tree are not ordered.\"\n \"Costs will be accurate but actually contracting requires \"\n \"ordered indices corresponding to array axes.\"\n )\n\n if not isinstance(next(iter(size_dict.values()), 1), int):\n # make sure we are working with python integers to avoid overflow\n # comparison errors with inf etc.\n self.size_dict = {k: int(v) for k, v in size_dict.items()}\n else:\n self.size_dict = size_dict\n\n self.N = len(self.inputs)\n\n # the index representation for each input is an ordered mapping of\n # each index to the number of times it has appeared on children. By\n # also tracking the total number of appearances one can efficiently\n # and locally compute which indices should be kept or contracted\n self.appearances = {}\n for term in self.inputs:\n for ix in term:\n self.appearances[ix] = self.appearances.get(ix, 0) + 1\n # adding output appearances ensures these are never contracted away,\n # N.B. if after this step every appearance count is exactly 2,\n # then there are no 'hyper' indices in the contraction\n for ix in self.output:\n self.appearances[ix] = self.appearances.get(ix, 0) + 1\n\n # this stores potentialy preprocessing steps that are not part of the\n # main contraction tree, but assumed to have been applied, for example\n # tracing or summing over indices that appear only once\n self.preprocessing = {}\n\n # mapping of parents to children - the core binary tree object\n self.children = {}\n\n # information about all the nodes\n self.nodeops = get_nodeops(nodeops, self.N)\n self.info = {}\n\n # add constant nodes: the leaves\n for leaf in self.gen_leaves():\n self._add_node(leaf, extent=1) # leaf extent is always 1\n\n # and the root or top node\n self.root = self.nodeops.node_supremum(self.N)\n self._add_node(self.root, extent=self.N) # root extent is always N\n\n if self.N == 1:\n # trivial 'contraction', single input maps directly to output,\n self.children[self.root] = (leaf,)\n\n # whether to keep track of dangling nodes/subgraphs\n self.track_childless = track_childless\n if self.track_childless:\n # the set of dangling nodes\n self.childless = oset([self.root])\n\n # running largest_intermediate and total flops\n self._track_flops = track_flops\n if track_flops:\n self._flops = 0\n\n self._track_write = track_write\n if track_write:\n self._write = 0\n\n self._track_size = track_size\n if track_size:\n self._sizes = MaxCounter()\n\n # container for caching subtree reconfiguration condidates\n self.already_optimized = dict()\n\n # info relating to slicing (base constructor is always unsliced)\n self.multiplicity = 1\n self.sliced_inds = {}\n self.sliced_inputs = frozenset()\n\n # cache for compiled contraction cores\n self.contraction_cores = {}\n\n # a default objective function useful for\n # further optimization and scoring\n self._default_objective = objective\n\n def set_state_from(self, other):\n \"\"\"Set the internal state of this tree to that of ``other``.\"\"\"\n # immutable or never mutated properties\n for attr in (\n \"appearances\",\n \"inputs\",\n \"multiplicity\",\n \"N\",\n \"output\",\n \"root\",\n \"size_dict\",\n \"sliced_inputs\",\n \"_default_objective\",\n ):\n setattr(self, attr, getattr(other, attr))\n\n # mutable properties\n for attr in (\n \"children\",\n \"contraction_cores\",\n \"nodeops\",\n \"sliced_inds\",\n \"preprocessing\",\n ):\n setattr(self, attr, getattr(other, attr).copy())\n\n # dicts of mutable\n for attr in (\"info\", \"already_optimized\"):\n setattr(\n self,\n attr,\n {k: v.copy() for k, v in getattr(other, attr).items()},\n )\n\n self.track_childless = other.track_childless\n if other.track_childless:\n self.childless = other.childless.copy()\n\n self._track_flops = other._track_flops\n if other._track_flops:\n self._flops = other._flops\n\n self._track_write = other._track_write\n if other._track_write:\n self._write = other._write\n\n self._track_size = other._track_size\n if other._track_size:\n self._sizes = other._sizes.copy()\n\n def copy(self):\n \"\"\"Create a copy of this ``ContractionTree``.\"\"\"\n tree = object.__new__(self.__class__)\n tree.set_state_from(self)\n return tree\n\n def set_default_objective(self, objective):\n \"\"\"Set the objective function for this tree.\"\"\"\n self._default_objective = get_score_fn(objective)\n\n def get_default_objective(self):\n \"\"\"Get the objective function for this tree.\"\"\"\n if self._default_objective is None:\n self._default_objective = get_score_fn(\"flops\")\n return self._default_objective\n\n def get_default_combo_factor(self):\n \"\"\"Get the default combo factor for this tree.\"\"\"\n objective = self.get_default_objective()\n try:\n return objective.factor\n except AttributeError:\n return DEFAULT_COMBO_FACTOR\n\n def get_score(self, objective=None):\n \"\"\"Score this tree using the default objective function.\"\"\"\n from .scoring import get_score_fn\n\n if objective is None:\n objective = self.get_default_objective()\n\n objective = get_score_fn(objective)\n\n return objective({\"tree\": self})\n\n @property\n def nslices(self):\n \"\"\"Simple alias for how many independent contractions this tree\n represents overall.\n \"\"\"\n return self.multiplicity\n\n @property\n def nchunks(self):\n \"\"\"The number of 'chunks' - determined by the number of sliced output\n indices.\n \"\"\"\n return prod(\n si.size for si in self.sliced_inds.values() if not si.inner\n )\n\n def input_to_node(self, i):\n \"\"\"Create a node from a single input index, i.e. the subgraph that\n only contains the input tensor ``i``.\n\n Parameters\n ----------\n i : int\n The input index.\n\n Returns\n -------\n node : node_type\n \"\"\"\n return self.nodeops.node_from_single(i)\n\n def node_to_input(self, node):\n \"\"\"Assuming ``node`` has one element, i.e. is a leaf, return the\n corresponding input index.\n\n Parameters\n ----------\n node : node_type\n The node to convert.\n\n Returns\n -------\n i : int\n \"\"\"\n return self.nodeops.node_get_single_el(node)\n\n def node_to_terms(self, node):\n \"\"\"Turn a node into the corresponding terms a sequence of leaf legs,\n corresponding to input indices.\n \"\"\"\n return (\n self.get_legs(self.input_to_node(i))\n for i in self.get_subgraph(node)\n )\n\n def gen_leaves(self):\n \"\"\"Generate the nodes representing leaves of the contraction tree, i.e.\n of size 1 each corresponding to a single input tensor.\n \"\"\"\n return map(self.input_to_node, range(self.N))\n\n def get_incomplete_nodes(self):\n \"\"\"Get the set of current nodes that have no children and the set of\n nodes that have no parents. These are the 'childless' and 'parentless'\n nodes respectively, that need to be contracted to complete the tree.\n The parentless nodes are grouped into the childless nodes that contain\n them as subgraphs.\n\n Returns\n -------\n groups : dict[node_type, list[node_type]]\n A mapping of childless nodes to the list of parentless nodes are\n beneath them.\n\n See Also\n --------\n autocomplete\n \"\"\"\n childless = dict.fromkeys(\n node\n for node in self.info\n # start wth all but leaves\n if not self.is_leaf(node)\n )\n parentless = dict.fromkeys(\n node\n for node in self.info\n # start with all but root\n if not self.is_root(node)\n )\n for p, (l, r) in self.children.items():\n parentless.pop(l)\n parentless.pop(r)\n childless.pop(p)\n\n groups = {node: [] for node in childless}\n for node in parentless:\n # get the smallest node that contains this node\n ancestor = min(\n (\n possible_parent\n for possible_parent in childless\n if set(self.get_subgraph(node)).issubset(\n self.get_subgraph(possible_parent)\n )\n # XXX: for non-ssa node types could do:\n # if self.is_descendant(node, possible_parent)\n ),\n key=self.get_extent,\n )\n groups[ancestor].append(node)\n\n return groups\n\n def autocomplete(self, **contract_opts):\n \"\"\"Contract all remaining node groups (as computed by\n ``tree.get_incomplete_nodes``) in the tree to complete it.\n\n Parameters\n ----------\n contract_opts\n Options to pass to ``tree.contract_nodes``.\n\n See Also\n --------\n get_incomplete_nodes, contract_nodes\n \"\"\"\n groups = self.get_incomplete_nodes()\n for grandparent, parentless_subnodes in groups.items():\n self.contract_nodes(\n parentless_subnodes, grandparent=grandparent, **contract_opts\n )\n\n @classmethod\n def from_path(\n cls,\n inputs,\n output,\n size_dict,\n *,\n path=None,\n ssa_path=None,\n edge_path=None,\n optimize=\"auto\",\n autocomplete=\"auto\",\n check=False,\n **kwargs,\n ):\n \"\"\"Create a (completed) ``ContractionTree`` from the usual inputs plus\n a standard contraction path or 'ssa_path' - you need to supply one.\n\n Parameters\n ----------\n inputs : Sequence[Sequence[str]]\n The input indices of each tensor, as single unicode characters.\n output : Sequence[str]\n The output indices.\n size_dict : dict[str, int]\n The size of each index.\n path : Sequence[Sequence[int]], optional\n The contraction path, a sequence of pairs of tensor ids to\n contract. The ids are linear indices into the list of temporary\n tensors, which are recycled as each contraction pops a pair and\n appends the result. One of ``path``, ``ssa_path`` or ``edge_path``\n must be supplied.\n ssa_path : Sequence[Sequence[int]], optional\n The contraction path, a sequence of pairs of indices to contract.\n The indices are single use, as if the result of each contraction is\n appended to the end of the list of temporary tensors without\n popping. One of ``path``, ``ssa_path`` or ``edge_path`` must be\n supplied.\n edge_path : Sequence[str], optional\n The contraction path, a sequence of indices to contract in order.\n One of ``path``, ``ssa_path`` or ``edge_path`` must be supplied.\n optimize : str, optional\n If a contraction within the path contains 3 or more tensors, how to\n optimize this subcontraction into a binary tree.\n autocomplete : \"auto\" or bool, optional\n Whether to automatically complete the path, i.e. contract all\n remaining nodes. If \"auto\" then a warning is issued if the path is\n not complete.\n check : bool, optional\n Whether to perform some basic checks while creating the contraction\n nodes.\n\n Returns\n -------\n ContractionTree\n \"\"\"\n if (path is None) + (ssa_path is None) + (edge_path is None) != 2:\n raise ValueError(\n \"Exactly one of ``path`` or ``ssa_path`` must be supplied.\"\n )\n\n contract_opts = {\"optimize\": optimize, \"check\": check}\n\n if edge_path is not None:\n from .pathfinders.path_basic import edge_path_to_ssa\n\n ssa_path = edge_path_to_ssa(edge_path, inputs)\n\n if ssa_path is not None:\n path = ssa_path\n\n tree = cls(inputs, output, size_dict, **kwargs)\n\n if ssa_path is not None:\n # ssa path ('single static assignment' ids)\n nodes = dict(enumerate(tree.gen_leaves()))\n ssa = len(nodes)\n for p in path:\n merge = [nodes.pop(i) for i in p]\n nodes[ssa] = tree.contract_nodes(merge, **contract_opts)\n ssa += 1\n nodes = nodes.values()\n else:\n # regular path ('recycled' ids)\n nodes = list(tree.gen_leaves())\n for p in path:\n merge = [nodes.pop(i) for i in sorted(p, reverse=True)]\n nodes.append(tree.contract_nodes(merge, **contract_opts))\n\n if len(nodes) > 1 and autocomplete:\n if autocomplete == \"auto\":\n # warn that we are completing\n warnings.warn(\n \"Path was not complete - contracting all remaining. \"\n \"You can silence this warning with `autocomplete=True`.\"\n \"Or produce an incomplete tree with `autocomplete=False`.\"\n )\n\n tree.contract_nodes(nodes, grandparent=tree.root, **contract_opts)\n\n return tree\n\n @classmethod\n def from_info(cls, info, **kwargs):\n \"\"\"Create a ``ContractionTree`` from an ``opt_einsum.PathInfo`` object.\"\"\"\n return cls.from_path(\n inputs=info.input_subscripts.split(\",\"),\n output=info.output_subscript,\n size_dict=info.size_dict,\n path=info.path,\n **kwargs,\n )\n\n @classmethod\n def from_eq(cls, eq, size_dict, **kwargs):\n \"\"\"Create a empty ``ContractionTree`` directly from an equation and set\n of shapes.\n\n Parameters\n ----------\n eq : str\n The einsum string equation.\n size_dict : dict[str, int]\n The size of each index.\n \"\"\"\n lhs, output = eq.split(\"->\")\n inputs = lhs.split(\",\")\n return cls(inputs, output, size_dict, **kwargs)\n\n def get_eq(self):\n \"\"\"Get the einsum equation corresponding to this tree. Note that this\n is the total (or original) equation, so includes indices which have\n been sliced.\n\n Returns\n -------\n eq : str\n \"\"\"\n return inputs_output_to_eq(self.inputs, self.output)\n\n def get_shapes(self):\n \"\"\"Get the shapes of the input tensors corresponding to this tree.\n\n Returns\n -------\n shapes : tuple[tuple[int]]\n \"\"\"\n return tuple(\n tuple(self.size_dict[ix] for ix in term) for term in self.inputs\n )\n\n def get_inputs_sliced(self):\n \"\"\"Get the input indices corresponding to a single slice of this tree,\n i.e. with sliced indices removed.\n\n Returns\n -------\n inputs : tuple[tuple[str]]\n \"\"\"\n return tuple(\n tuple(ix for ix in term if ix not in self.sliced_inds)\n for term in self.inputs\n )\n\n def get_output_sliced(self):\n \"\"\"Get the output indices corresponding to a single slice of this tree,\n i.e. with sliced indices removed.\n\n Returns\n -------\n output : tuple[str]\n \"\"\"\n return tuple(ix for ix in self.output if ix not in self.sliced_inds)\n\n def get_eq_sliced(self):\n \"\"\"Get the einsum equation corresponding to a single slice of this\n tree, i.e. with sliced indices removed.\n\n Returns\n -------\n eq : str\n \"\"\"\n return inputs_output_to_eq(\n self.get_inputs_sliced(), self.get_output_sliced()\n )\n\n def get_shapes_sliced(self):\n \"\"\"Get the shapes of the input tensors corresponding to a single slice\n of this tree, i.e. with sliced indices removed.\n\n Returns\n -------\n shapes : tuple[tuple[int]]\n \"\"\"\n return tuple(\n tuple(\n self.size_dict[ix] for ix in term if ix not in self.sliced_inds\n )\n for term in self.inputs\n )\n\n @classmethod\n def from_edge_path(\n cls,\n edge_path,\n inputs,\n output,\n size_dict,\n optimize=\"auto\",\n autocomplete=\"auto\",\n check=False,\n **kwargs,\n ):\n \"\"\"Create a ``ContractionTree`` from an edge elimination ordering.\"\"\"\n warnings.warn(\n \"ContractionTree.from_edge_path(edge_path, ...) is deprecated. Use\"\n \" ContractionTree.from_path(edge_path=edge_path, ...) instead.\"\n )\n return cls.from_path(\n inputs,\n output,\n size_dict,\n edge_path=edge_path,\n optimize=optimize,\n autocomplete=autocomplete,\n check=check,\n **kwargs,\n )\n\n def _add_node(self, node, check=False, **kwargs):\n \"\"\"Add a node to this tree, specified either directly as a existing\n node type, or as a subgraph (i.e. a sequence of input positions) which\n is then converted to a node with the corresponding extent and subgraph\n information.\n\n Note if \"ssa\" nodes are used, then adding two equivalent subgraphs\n will result in *two* new nodes, since the node labels do not\n themselves encode the subgraph information.\n\n Parameters\n ----------\n node : node_type or Sequence[int]\n The node to add, either directly as a node type, or as a subgraph\n specified by the sequence of input positions it contains.\n check : bool, optional\n Whether to perform some basic checks on the node and tree state\n before adding the node.\n kwargs : dict, optional\n Additional information to cache about this node, for example its\n 'extent' or 'subgraph'. If it is being specified as a sequence of\n input positions, these two will be injected automatically.\n\n Returns\n -------\n node : node_type\n The node that was added, which may be different from the input if\n the input was specified as a sequence of input positions.\n \"\"\"\n # first we possibly convert from subgraph spec to node\n if not isinstance(node, self.nodeops.node_type):\n # assume node *has* been specified as sequence of input positions\n subgraph = tuple(node)\n\n if len(subgraph) == 1:\n # leaf node, for ssa we don't want to generate a new node\n (i,) = subgraph\n node = self.nodeops.node_from_single(i)\n elif len(subgraph) == self.N:\n # root node, for ssa we don't want to generate a new node\n node = self.root\n else:\n # intermediate, assume we can generate new node identifier\n node = self.nodeops.new_node_for_seq(subgraph)\n kwargs.setdefault(\"extent\", len(subgraph))\n kwargs.setdefault(\"subgraph\", subgraph)\n\n if check:\n if len(self.info) > 2 * self.N - 1:\n raise ValueError(\"There are too many children already.\")\n if len(self.children) > self.N - 1:\n raise ValueError(\"There are too many branches already.\")\n if not self.nodeops.is_valid_node(node):\n raise ValueError(\"{} is not a valid node.\".format(node))\n\n try:\n d = self.info[node]\n except KeyError:\n d = self.info[node] = {}\n\n if kwargs:\n d.update(kwargs)\n\n return node\n\n def _remove_node(self, node):\n \"\"\"Remove ``node`` from this tree and update the flops and maximum size\n if tracking them respectively, as well as input pre-processing.\n \"\"\"\n node_extent = self.get_extent(node)\n\n if node_extent == 1:\n # leaf nodes should always exist\n self.info[node].clear()\n self.info[node][\"extent\"] = 1 # leaf extent is always 1\n # input: remove any associated preprocessing\n self.preprocessing.pop(self.nodeops.node_get_single_el(node), None)\n else:\n # only non-leaf nodes contribute to size, flops and write\n if self._track_size:\n self._sizes.discard(self.get_size(node))\n\n if self._track_flops:\n self._flops -= self.get_flops(node)\n\n if self._track_write:\n self._write -= self.get_size(node)\n\n del self.children[node]\n if node_extent == self.N:\n # root node should always exist\n self.info[node].clear()\n self.info[node][\"extent\"] = self.N # root extent is always N\n else:\n del self.info[node]\n\n def compute_leaf_legs(self, i):\n \"\"\"Compute the effective 'outer' indices for the ith input tensor. This\n is not always simply the ith input indices, due to A) potential slicing\n and B) potential preprocessing.\n \"\"\"\n # indices of input tensor (after slicing which is done immediately)\n if self.sliced_inds:\n term = tuple(\n ix for ix in self.inputs[i] if ix not in self.sliced_inds\n )\n else:\n term = self.inputs[i]\n\n legs = {}\n for ix in term:\n legs[ix] = legs.get(ix, 0) + 1\n\n # check for single term simplifications, these are treated as a simple\n # preprocessing step that only is taken into account during actual\n # contraction, and are not represented in the binary tree\n # N.B. need to compute simplifiability *after* slicing\n is_simplifiable = (\n # repeated indices (diag or traces)\n (len(term) != len(legs))\n or\n # reduced indices (are summed immediately)\n any(\n ix_count == self.appearances[ix]\n for ix, ix_count in legs.items()\n )\n )\n\n if is_simplifiable:\n # compute the simplified legs -> the new effective input legs\n legs = {\n ix: ix_count\n for ix, ix_count in legs.items()\n if ix_count != self.appearances[ix]\n }\n # add a preprocessing step to the list of contractions\n eq = inputs_output_to_eq((term,), legs, canonicalize=True)\n self.preprocessing[i] = eq\n\n return legs\n\n def has_preprocessing(self):\n # touch all inputs legs, since preprocessing is lazily computed\n for node in self.gen_leaves():\n self.get_legs(node)\n return bool(self.preprocessing)\n\n def has_hyper_indices(self):\n \"\"\"Check if there are any 'hyper' indices in the contraction, i.e.\n indices that don't appear exactly twice, when considering the inputs\n and output.\n \"\"\"\n return any(ix_count != 2 for ix_count in self.appearances.values())\n\n @cached_node_property(\"extent\")\n def get_extent(self, node):\n \"\"\"Get the number of input tensors contained in the subgraph\n represented by ``node``.\n\n Parameters\n ----------\n node : node_type\n The node to compute the extent of.\n\n Returns\n -------\n extent : int\n \"\"\"\n if node in self.children:\n l, r = self.children[node]\n return self.get_extent(l) + self.get_extent(r)\n else:\n return self.nodeops.node_size(node)\n\n @cached_node_property(\"subgraph\")\n def get_subgraph(self, node) -> tuple[int, ...]:\n \"\"\"Get the sequence of input tensors contained in subgraph represented\n by ``node``.\n\n Parameters\n ----------\n node : node_type\n The node to compute the subgraph of.\n\n Returns\n -------\n subgraph : tuple[int]\n The input tensor indices contained in this subgraph.\n \"\"\"\n node_extent = self.get_extent(node)\n if node_extent == 1:\n return (self.nodeops.node_get_single_el(node),)\n elif node_extent == self.N:\n return tuple(range(self.N))\n else:\n try:\n left, right = self.children[node]\n return self.get_subgraph(left) + self.get_subgraph(right)\n except KeyError:\n # this should only happen if directly creating\n # incomplete nodes e.g. not in a bottom up fashion\n # ssa nodes e.g. will not support this operation\n return tuple(node)\n\n @cached_node_property(\"legs\")\n def get_legs(self, node):\n \"\"\"Get the effective 'outer' indices for the collection of tensors\n in ``node``.\n \"\"\"\n # should this comparison be with self.N for efficiency?\n if node == self.root:\n # root legs are output, after slicing\n # n.b. the index counts are irrelevant for the output\n return {ix: 0 for ix in self.output if ix not in self.sliced_inds}\n\n node_extent = self.get_extent(node)\n\n if node_extent == 1:\n # leaf legs are inputs\n return self.compute_leaf_legs(\n self.nodeops.node_get_single_el(node)\n )\n\n try:\n involved = self.get_involved(node)\n except KeyError:\n # this should only happen if directly creating\n # incomplete nodes e.g. not in a bottom up fashion\n involved = legs_union(self.node_to_terms(node))\n\n return {\n ix: ix_count\n for ix, ix_count in involved.items()\n if ix_count < self.appearances[ix]\n }\n\n @cached_node_property(\"involved\")\n def get_involved(self, node):\n \"\"\"Get all the indices involved in the formation of subgraph ``node``.\"\"\"\n if self.is_leaf(node):\n return {}\n sub_legs = map(self.get_legs, self.children[node])\n return legs_union(sub_legs)\n\n @cached_node_property(\"size\")\n def get_size(self, node):\n \"\"\"Get the tensor size of ``node``.\"\"\"\n return compute_size_by_dict(self.get_legs(node), self.size_dict)\n\n @cached_node_property(\"flops\")\n def get_flops(self, node):\n \"\"\"Get the FLOPs for the pairwise contraction that will create\n ``node``.\n \"\"\"\n if self.is_leaf(node):\n return 0\n involved = self.get_involved(node)\n return compute_size_by_dict(involved, self.size_dict)\n\n @cached_node_property(\"can_dot\")\n def get_can_dot(self, node):\n \"\"\"Get whether this contraction can be performed as a dot product (i.e.\n with ``tensordot``), or else requires ``einsum``, as it has indices\n that don't appear exactly twice in either the inputs or the output.\n \"\"\"\n l, r = self.children[node]\n sp, sl, sr = map(self.get_legs, (node, l, r))\n return set(sp) == set(sl).symmetric_difference(sr)\n\n @cached_node_property(\"inds\")\n def get_inds(self, node):\n \"\"\"Get the indices of this node - an ordered string version of\n ``get_legs`` that starts with ``tree.inputs`` and maintains the order\n they appear in each contraction 'ABC,abc->ABCabc', to match tensordot.\n \"\"\"\n # NB: self.inputs and self.output contain the full (unsliced) indices\n # thus we filter even the input legs and output legs\n\n if self.get_extent(node) in (1, self.N):\n return \"\".join(self.get_legs(node))\n\n legs = self.get_legs(node)\n l_inds, r_inds = map(self.get_inds, self.children[node])\n # the filter here takes care of contracted indices\n return \"\".join(\n unique(filter(legs.__contains__, itertools.chain(l_inds, r_inds)))\n )\n\n @cached_node_property(\"tensordot_axes\")\n def get_tensordot_axes(self, node):\n \"\"\"Get the ``axes`` arg for a tensordot ocontraction that produces\n ``node``. The pairs are sorted in order of appearance on the left\n input.\n \"\"\"\n l_inds, r_inds = map(self.get_inds, self.children[node])\n l_axes, r_axes = [], []\n for i, ind in enumerate(l_inds):\n j = r_inds.find(ind)\n if j != -1:\n l_axes.append(i)\n r_axes.append(j)\n return tuple(l_axes), tuple(r_axes)\n\n @cached_node_property(\"tensordot_perm\")\n def get_tensordot_perm(self, node):\n \"\"\"Get the permutation required, if any, to bring the tensordot output\n of this nodes contraction into line with ``self.get_inds(node)``.\n \"\"\"\n l_inds, r_inds = map(self.get_inds, self.children[node])\n # the target output inds\n p_inds = self.get_inds(node)\n # the tensordot output inds\n td_inds = \"\".join(sorted(p_inds, key=f\"{l_inds}{r_inds}\".find))\n if td_inds == p_inds:\n return None\n return tuple(map(td_inds.find, p_inds))\n\n @cached_node_property(\"einsum_eq\")\n def get_einsum_eq(self, node):\n \"\"\"Get the einsum string describing the contraction that produces\n ``node``, unlike ``get_inds`` the characters are mapped into [a-zA-Z],\n for compatibility with ``numpy.einsum`` for example.\n \"\"\"\n l, r = self.children[node]\n l_inds, r_inds, p_inds = map(self.get_inds, (l, r, node))\n # we need to map any extended unicode characters into ascii\n char_mapping = {\n ord(ix): get_symbol(i)\n for i, ix in enumerate(unique(itertools.chain(l_inds, r_inds)))\n }\n return f\"{l_inds},{r_inds}->{p_inds}\".translate(char_mapping)\n\n def is_leaf(self, node):\n \"\"\"Check if ``node`` is a leaf node in this tree.\n\n Parameters\n ----------\n node : node_type\n The node to check.\n\n Returns\n -------\n is_leaf : bool\n \"\"\"\n return self.nodeops.is_leaf(node)\n\n def is_root(self, node):\n \"\"\"Check if ``node`` is the root node in this tree.\n\n Parameters\n ----------\n node : node_type\n The node to check.\n\n Returns\n -------\n is_root : bool\n \"\"\"\n return self.nodeops.is_supremum(node, self.N)\n\n def is_descendant(self, node, ancestor):\n \"\"\"Check if ``node`` is a descendant of ``ancestor`` in this tree.\n\n Parameters\n ----------\n node : node_type\n The node to check.\n ancestor : node_type\n The potential ancestor node.\n\n Returns\n -------\n is_descendant : bool\n \"\"\"\n return self.nodeops.node_issubset(node, ancestor)\n\n def get_peak_size(self, node):\n \"\"\"Get the peak size for all but only the contractions required to\n produce ``node``. The value for the root note will be the peak size of\n the entire contraction.\n\n Parameters\n ----------\n node : node_type\n The node to compute the peak size of.\n\n Returns\n -------\n peak_size : int\n \"\"\"\n if self.is_leaf(node):\n # leaf node is input\n return self.get_size(node)\n\n l, r = self.children[node]\n # peak either occured while:\n # 1. we were forming left intermediate\n peakleft = self.get_peak_size(l)\n # 2. we were forming right intermediate (whilst holding left)\n peakright = self.get_size(l) + self.get_peak_size(r)\n # 3. or we were performing this contraction, including output\n peakthis = self.get_size(l) + self.get_size(r) + self.get_size(node)\n\n return max(peakleft, peakright, peakthis)\n\n def reorder_for_peak_size(self):\n \"\"\"This reorders the depth first traversal of the tree to minimize\n the peak size of the contraction.\n \"\"\"\n changed = False\n for p, l, r in self.traverse():\n sl = self.get_size(l)\n pl = self.get_peak_size(l)\n sr = self.get_size(r)\n pr = self.get_peak_size(r)\n # peak if hold left while we form right\n plr = max(pl, sl + pr)\n # peak if hold right while we form left\n prl = max(pr, sr + pl)\n if prl < plr:\n self.children[p] = r, l\n changed = True\n return changed\n\n def get_centrality(self, node):\n try:\n return self.info[node][\"centrality\"]\n except KeyError:\n self.compute_centralities()\n return self.info[node][\"centrality\"]\n\n def total_flops(self, dtype=None, log=None):\n \"\"\"Sum the flops contribution from every node in the tree.\n\n Parameters\n ----------\n dtype : {'float', 'complex', None}, optional\n Scale the answer depending on the assumed data type.\n \"\"\"\n if self._track_flops:\n C = self.multiplicity * self._flops\n\n else:\n self._flops = 0\n for node, _, _ in self.traverse():\n self._flops += self.get_flops(node)\n\n self._track_flops = True\n C = self.multiplicity * self._flops\n\n if dtype is None:\n pass\n elif \"float\" in dtype:\n C *= 2\n elif \"complex\" in dtype:\n C *= 4\n else:\n raise ValueError(f\"Unknown dtype {dtype}\")\n\n if log is not None:\n C = math.log(max(C, 1), log)\n\n return C\n\n def total_write(self):\n \"\"\"Sum the total amount of memory that will be created and operated on.\"\"\"\n if not self._track_write:\n self._write = 0\n for node, _, _ in self.traverse():\n self._write += self.get_size(node)\n\n self._track_write = True\n\n return self.multiplicity * self._write\n\n def combo_cost(self, factor=DEFAULT_COMBO_FACTOR, combine=sum, log=None):\n t = 0\n for p in self.children:\n f = self.get_flops(p)\n w = self.get_size(p)\n t += combine((f, factor * w))\n\n t *= self.multiplicity\n\n if log is not None:\n t = math.log(t, log)\n\n return t\n\n total_cost = combo_cost\n\n def max_size(self, log=None):\n \"\"\"The size of the largest intermediate tensor.\"\"\"\n if self.N == 1:\n return self.get_size(self.root)\n\n if not self._track_size:\n self._sizes = MaxCounter()\n for node, _, _ in self.traverse():\n self._sizes.add(self.get_size(node))\n self._track_size = True\n\n size = self._sizes.max()\n\n if log is not None:\n size = math.log(size, log)\n\n return size\n\n def max_contraction_size(self, log=None):\n \"\"\"The maximum size of a single contraction in the tree. This includes\n the size of the two input tensors and the output tensor, and can be a\n more practical measure of the peak memory required.\n\n Parameters\n ----------\n log : float, optional\n If provided, return the log of the size to this base.\n\n Returns\n -------\n size : int or float\n The maximum size of a single contraction in the tree, or its log.\n \"\"\"\n Y = max(\n self.get_size(p) + self.get_size(l) + self.get_size(r)\n for p, (l, r) in self.children.items()\n )\n\n if log is not None:\n Y = math.log(Y, log)\n\n return Y\n\n def peak_size(self, order=None, log=None):\n \"\"\"Get the peak concurrent size of tensors needed - this depends on the\n traversal order, i.e. the exact contraction path, not just the\n contraction tree.\n \"\"\"\n tot_size = sum(self.get_size(node) for node in self.gen_leaves())\n peak = tot_size\n for p, l, r in self.traverse(order=order):\n tot_size += self.get_size(p)\n # measure peak assuming we need both inputs and output\n peak = max(peak, tot_size)\n tot_size -= self.get_size(l)\n tot_size -= self.get_size(r)\n\n if log is not None:\n peak = math.log(peak, log)\n\n return peak\n\n def contract_stats(self, force=False):\n \"\"\"Simulteneously compute the total flops, write and size of the\n contraction tree. This is more efficient than calling each of the\n individual methods separately. Once computed, each quantity is then\n automatically tracked.\n\n Returns\n -------\n stats : dict[str, int]\n The total flops, write and size.\n \"\"\"\n if force or not (\n self._track_flops and self._track_write and self._track_size\n ):\n self._flops = self._write = 0\n self._sizes = MaxCounter()\n\n for node, _, _ in self.traverse():\n self._flops += self.get_flops(node)\n node_size = self.get_size(node)\n self._write += node_size\n self._sizes.add(node_size)\n\n self._track_flops = self._track_write = self._track_size = True\n\n return {\n \"flops\": max(self.multiplicity * self._flops, 1),\n \"write\": max(self.multiplicity * self._write, 1),\n \"size\": max(self._sizes.max(), 1),\n }\n\n def arithmetic_intensity(self):\n \"\"\"The ratio of total flops to total write - the higher the better for\n extracting good computational performance.\n \"\"\"\n return self.total_flops(dtype=None) / self.total_write()\n\n def contraction_scaling(self):\n \"\"\"This is computed simply as the maximum number of indices involved\n in any single contraction, which will match the scaling assuming that\n all dimensions are equal.\n \"\"\"\n return max(len(self.get_involved(node)) for node in self.info)\n\n def contraction_cost(self, log=None):\n \"\"\"Get the total number of scalar operations ~ time complexity.\"\"\"\n return self.total_flops(dtype=None, log=log)\n\n def naive_cost(self, log=None):\n \"\"\"Get the naive cost of performing this contraction as a single\n einsum summation, without any intermediate contractions. This is given\n the as product of the size of all indices.\n\n Parameters\n ----------\n log : float, optional\n If provided, return log of the cost to this base.\n \"\"\"\n if log is None:\n return prod(self.size_dict[ix] for ix in self.appearances)\n else:\n return sum(\n math.log(self.size_dict[ix], log) for ix in self.appearances\n )\n\n def speedup(self, log=None):\n \"\"\"Speedup compared to naive summation.\n\n Parameters\n ----------\n log : float, optional\n If provided, return log of the speedup to this base.\n \"\"\"\n if log is None:\n return self.naive_cost() / self.contraction_cost()\n else:\n logc = self.contraction_cost(log=log)\n logn = self.naive_cost(log=log)\n return logn - logc\n\n def contraction_width(self, log=2):\n \"\"\"Get log2 of the size of the largest tensor.\"\"\"\n return self.max_size(log=log)\n\n def compressed_contract_stats(\n self,\n chi=None,\n order=\"surface_order\",\n compress_late=None,\n ):\n if chi is None:\n chi = self.get_default_chi()\n\n if compress_late is None:\n compress_late = self.get_default_compress_late()\n\n hg = self.get_hypergraph(accel=\"auto\")\n\n # conversion between tree nodes <-> hypergraph nodes during contraction\n tree_map = dict(zip(self.gen_leaves(), range(hg.get_num_nodes())))\n\n tracker = CompressedStatsTracker(hg, chi)\n\n for p, l, r in self.traverse(order):\n li = tree_map[l]\n ri = tree_map[r]\n\n tracker.update_pre_step()\n\n if compress_late:\n tracker.update_pre_compress(hg, li, ri)\n # compress just before we contract tensors\n hg.compress(chi=chi, edges=hg.get_node(li))\n hg.compress(chi=chi, edges=hg.get_node(ri))\n tracker.update_post_compress(hg, li, ri)\n\n tracker.update_pre_contract(hg, li, ri)\n pi = tree_map[p] = hg.contract(li, ri)\n tracker.update_post_contract(hg, pi)\n\n if not compress_late:\n # compress as soon as we can after contracting tensors\n tracker.update_pre_compress(hg, pi)\n hg.compress(chi=chi, edges=hg.get_node(pi))\n tracker.update_post_compress(hg, pi)\n\n tracker.update_post_step()\n\n return tracker\n\n def total_flops_compressed(\n self,\n chi=None,\n order=\"surface_order\",\n compress_late=None,\n dtype=None,\n log=None,\n ):\n \"\"\"Estimate the total flops for a compressed contraction of this tree\n with maximum bond size ``chi``. This includes basic estimates of the\n ops to perform contractions, QRs and SVDs.\n \"\"\"\n if dtype is not None:\n raise ValueError(\n \"Can only estimate cost in terms of \"\n \"number of abstract scalar ops.\"\n )\n\n F = self.compressed_contract_stats(\n chi=chi,\n order=order,\n compress_late=compress_late,\n ).flops\n\n if log is not None:\n F = math.log(F, log)\n\n return F\n\n contraction_cost_compressed = total_flops_compressed\n\n def total_write_compressed(\n self,\n chi=None,\n order=\"surface_order\",\n compress_late=None,\n log=None,\n ):\n \"\"\"Compute the total size of all intermediate tensors when a\n compressed contraction is performed with maximum bond size ``chi``,\n ordered by ``order``. This is relevant maybe for time complexity and\n e.g. autodiff space complexity (since every intermediate is kept).\n \"\"\"\n W = self.compressed_contract_stats(\n chi=chi,\n order=order,\n compress_late=compress_late,\n ).write\n\n if log is not None:\n W = math.log(W, log)\n\n return W\n\n def combo_cost_compressed(\n self,\n chi=None,\n order=\"surface_order\",\n compress_late=None,\n factor=None,\n log=None,\n ):\n if factor is None:\n factor = self.get_default_combo_factor()\n\n C = self.total_flops_compressed(\n chi=chi, order=order, compress_late=compress_late\n ) + factor * self.total_write_compressed(\n chi=chi, order=order, compress_late=compress_late\n )\n\n if log is not None:\n C = math.log(C, log)\n\n return C\n\n total_cost_compressed = combo_cost_compressed\n\n def max_size_compressed(\n self, chi=None, order=\"surface_order\", compress_late=None, log=None\n ):\n \"\"\"Compute the maximum sized tensor produced when a compressed\n contraction is performed with maximum bond size ``chi``, ordered by\n ``order``. This is close to the ideal space complexity if only\n tensors that are being directly operated on are kept in memory.\n \"\"\"\n S = self.compressed_contract_stats(\n chi=chi,\n order=order,\n compress_late=compress_late,\n ).max_size\n\n if log is not None:\n S = math.log(S, log)\n\n return S\n\n def peak_size_compressed(\n self,\n chi=None,\n order=\"surface_order\",\n compress_late=None,\n accel=\"auto\",\n log=None,\n ):\n \"\"\"Compute the peak size of combined intermediate tensors when a\n compressed contraction is performed with maximum bond size ``chi``,\n ordered by ``order``. This is the practical space complexity if one is\n not swapping intermediates in and out of memory.\n \"\"\"\n P = self.compressed_contract_stats(\n chi=chi,\n order=order,\n compress_late=compress_late,\n ).peak_size\n\n if log is not None:\n P = math.log(P, log)\n\n return P\n\n def contraction_width_compressed(\n self, chi=None, order=\"surface_order\", compress_late=None, log=2\n ):\n \"\"\"Compute log2 of the maximum sized tensor produced when a compressed\n contraction is performed with maximum bond size ``chi``, ordered by\n ``order``.\n \"\"\"\n return self.max_size_compressed(chi, order, compress_late, log=log)\n\n def _update_tracked(self, node):\n if self._track_flops:\n self._flops += self.get_flops(node)\n if self._track_write:\n self._write += self.get_size(node)\n if self._track_size:\n self._sizes.add(self.get_size(node))\n\n def contract_nodes_pair(\n self,\n x,\n y,\n legs=None,\n cost=None,\n size=None,\n parent=None,\n check=False,\n ):\n \"\"\"Contract node ``x`` with node ``y`` in the tree to create a new\n parent node, which is returned.\n\n Parameters\n ----------\n x : node_type\n The first node to contract.\n y : node_type\n The second node to contract.\n legs : dict[str, int], optional\n The effective 'legs' of the new node if already known. If not\n given, this is computed from the inputs of ``x`` and ``y``.\n cost : int, optional\n The cost of the contraction if already known. If not given, this is\n computed from the inputs of ``x`` and ``y``.\n size : int, optional\n The size of the new node if already known. If not given, this is\n computed from the inputs of ``x`` and ``y``.\n check : bool, optional\n Whether to check the inputs are valid.\n\n Returns\n -------\n parent : node_type\n The new parent node of ``x`` and ``y``.\n \"\"\"\n self._add_node(x, check=check)\n self._add_node(y, check=check)\n nx, ny = self.get_extent(x), self.get_extent(y)\n\n if parent is None:\n if nx + ny == self.N:\n parent = self.root\n else:\n parent = self.nodeops.new_node_for_union([x, y])\n self._add_node(parent, check=check)\n\n # enforce left ordering of 'heaviest' subtrees\n if nx == ny:\n # deterministically break ties\n sortx = self.nodeops.node_tie_breaker(x)\n sorty = self.nodeops.node_tie_breaker(y)\n else:\n sortx = nx\n sorty = ny\n\n if sortx > sorty:\n lr = (x, y)\n else:\n lr = (y, x)\n\n self.children[parent] = lr\n\n if self.track_childless:\n self.childless.discard(parent)\n if x not in self.children and nx > 1:\n self.childless.add(x)\n if y not in self.children and ny > 1:\n self.childless.add(y)\n\n # pre-computed information\n if legs is not None:\n self.info[parent][\"legs\"] = legs\n if cost is not None:\n self.info[parent][\"flops\"] = cost\n if size is not None:\n self.info[parent][\"size\"] = size\n\n self._update_tracked(parent)\n\n return parent\n\n def contract_nodes(\n self,\n nodes,\n optimize=\"auto\",\n grandparent=None,\n check=False,\n extra_opts=None,\n ):\n \"\"\"Contract an arbitrary number of ``nodes`` in the tree to build up a\n subtree. The root of this subtree (a new intermediate) is returned.\n \"\"\"\n # possibly convert from subgraph spec to node types\n nodes = tuple(self._add_node(node, check=check) for node in nodes)\n\n if len(nodes) == 1:\n return nodes[0]\n\n if len(nodes) == 2:\n return self.contract_nodes_pair(\n *nodes, parent=grandparent, check=check\n )\n\n from .interface import find_path\n\n # create the bottom and top nodes\n if grandparent is None:\n if sum(map(self.get_extent, nodes)) == self.N:\n # don't generate new node if root\n grandparent = self.root\n else:\n # assume we can generate new node\n grandparent = self.nodeops.new_node_for_union(nodes)\n\n self._add_node(grandparent, check=check)\n\n # if more than two nodes need to find the path to fill in between\n # \\\n # GN <- 'grandparent'\n # / \\\n # ?????????\n # ????????????? <- to be filled with 'temp nodes'\n # / \\ / / \\\n # N0 N1 N2 N3 N4 <- ``nodes``, or, subgraphs\n # / \\ / / \\\n legs_inputs = tuple(map(self.get_legs, nodes))\n path_inputs = tuple(map(tuple, legs_inputs))\n\n try:\n # output legs of the grandparent (after slicing)\n # we dont' use get_legs since we can do the shortcut below\n grand_legs = self.info[grandparent][\"legs\"]\n except KeyError:\n # compute legs directly from children\n if grandparent == self.root:\n # special case, need output ordering and sliced indices\n grand_legs = self.get_legs(grandparent)\n else:\n involved = legs_union(legs_inputs)\n grand_legs = {\n ix: ix_count\n for ix, ix_count in involved.items()\n if ix_count < self.appearances[ix]\n }\n self.info[grandparent][\"legs\"] = grand_legs\n\n path_output = tuple(grand_legs)\n\n path = find_path(\n path_inputs,\n path_output,\n self.size_dict,\n optimize=optimize,\n **(extra_opts or {}),\n )\n\n # now we have path create the nodes in between\n temp_nodes = list(nodes)\n for p in path[:-1]:\n to_contract = [temp_nodes.pop(i) for i in sorted(p, reverse=True)]\n temp_nodes.append(self.contract_nodes(to_contract, check=check))\n\n # want to explicitly specify the grandparent node:\n # so do the final pairwise contraction separately\n self.contract_nodes(temp_nodes, grandparent=grandparent, check=check)\n\n return grandparent\n\n def is_complete(self):\n \"\"\"Check every node has two children, unless it is a leaf.\"\"\"\n if self.N == 1:\n return True\n\n too_many_nodes = len(self.info) > 2 * self.N - 1\n too_many_branches = len(self.children) > self.N - 1\n\n if too_many_nodes or too_many_branches:\n raise ValueError(\"Contraction tree seems to be over complete!\")\n\n queue = [self.root]\n while queue:\n node = queue.pop()\n if self.is_leaf(node):\n continue\n try:\n queue.extend(self.children[node])\n except KeyError:\n return False\n\n return True\n\n def get_default_order(self):\n return \"dfs\"\n\n def _traverse_dfs(self):\n \"\"\"Traverse the tree in a depth first, non-recursive, order.\"\"\"\n ready = set(self.gen_leaves())\n queue = [self.root]\n\n while queue:\n node = queue[-1]\n l, r = self.children[node]\n\n # both node's children are ready -> we can yield this contraction\n if (l in ready) and (r in ready):\n ready.add(queue.pop())\n yield node, l, r\n continue\n\n if r not in ready:\n queue.append(r)\n if l not in ready:\n queue.append(l)\n\n def _traverse_ordered(self, order):\n \"\"\"Traverse the tree in the order that minimizes ``order(node)``, but\n still constrained to produce children before parents.\n \"\"\"\n from bisect import bisect\n\n if order == \"surface_order\":\n order = self.surface_order\n\n seen = set()\n queue = [self.root]\n scores = [order(self.root)]\n\n while len(seen) != len(self.children):\n i = 0\n while i < len(queue):\n node = queue[i]\n if node not in seen:\n for child in self.children[node]:\n if self.get_extent(child) > 1:\n # insert child into queue by score + before parent\n score = order(child)\n ci = bisect(scores[:i], score)\n scores.insert(ci, score)\n queue.insert(ci, child)\n # parent moves extra place to right\n i += 1\n seen.add(node)\n i += 1\n\n for node in queue:\n yield (node, *self.children[node])\n\n def traverse(self, order=None):\n \"\"\"Generate, in order, all the node merges in this tree. Non-recursive!\n This ensures children are always visited before their parent.\n\n Parameters\n ----------\n order : None, \"dfs\", or callable, optional\n How to order the contractions within the tree. If a callable is\n given (which should take a node as its argument), try to contract\n nodes that minimize this function first.\n\n Returns\n -------\n generator[tuple[node]]\n The bottom up ordered sequence of tree merges, each a\n tuple of ``(parent, left_child, right_child)``.\n\n See Also\n --------\n descend\n \"\"\"\n if self.N == 1:\n return\n\n if order is None:\n order = self.get_default_order()\n\n if order == \"dfs\":\n yield from self._traverse_dfs()\n else:\n yield from self._traverse_ordered(order=order)\n\n def descend(self, mode=\"dfs\"):\n \"\"\"Generate, from root to leaves, all the node merges in this tree.\n Non-recursive! This ensures parents are visited before their children.\n\n Parameters\n ----------\n mode : {'dfs', bfs}, optional\n How expand from a parent.\n\n Returns\n -------\n generator[tuple[node]\n The top down ordered sequence of tree merges, each a\n tuple of ``(parent, left_child, right_child)``.\n\n See Also\n --------\n traverse\n \"\"\"\n queue = [self.root]\n while queue:\n if mode == \"dfs\":\n parent = queue.pop(-1)\n elif mode == \"bfs\":\n parent = queue.pop(0)\n l, r = self.children[parent]\n yield parent, l, r\n if self.get_extent(l) > 1:\n queue.append(l)\n if self.get_extent(r) > 1:\n queue.append(r)\n\n def get_subtree(self, node, size, search=\"bfs\", seed=None):\n \"\"\"Get a subtree spanning down from ``node`` which will have ``size``\n leaves (themselves not necessarily leaves of the actual tree).\n\n Parameters\n ----------\n node : node_type\n The node of the tree to start with.\n size : int\n How many subtree leaves to aim for.\n search : {'bfs', 'dfs', 'random'}, optional\n How to build the tree:\n\n - 'bfs': breadth first expansion\n - 'dfs': depth first expansion (largest nodes first)\n - 'random': random expansion\n\n seed : None, int or random.Random, optional\n Random number generator seed, if ``search`` is 'random'.\n\n Returns\n -------\n sub_leaves : tuple[node_type]\n Nodes which are subtree leaves.\n branches : tuple[node_type]\n Nodes which are between the subtree leaves and root.\n \"\"\"\n # nodes which are subtree leaves\n branches = []\n\n # actual tree leaves - can't expand\n real_leaves = []\n\n # nodes to expand\n queue = [node]\n\n if search == \"random\":\n rng = get_rng(seed)\n else:\n rng = None\n if search == \"bfs\":\n i = 0\n elif search == \"dfs\":\n i = -1\n\n while (len(queue) + len(real_leaves) < size) and queue:\n if rng is not None:\n i = rng.randint(0, len(queue) - 1)\n\n p = queue.pop(i)\n if self.is_leaf(p):\n real_leaves.append(p)\n continue\n\n # the left child is always >= in weight that right child\n # if we append it last then ``.pop(-1)`` above perform the\n # depth first search sorting by node subgraph size\n l, r = self.children[p]\n\n queue.append(r)\n queue.append(l)\n branches.append(p)\n\n # nodes at the bottom of the subtree\n sub_leaves = queue + real_leaves\n\n return tuple(sub_leaves), tuple(branches)\n\n def remove_ind(self, ind, project=None, inplace=False):\n \"\"\"Remove (i.e. by default slice) index ``ind`` from this contraction\n tree, taking care to update all relevant information about each node.\n \"\"\"\n tree = self if inplace else self.copy()\n\n if ind in tree.sliced_inds:\n raise ValueError(f\"Index {ind} already sliced.\")\n\n # make sure all flops and size information has been populated\n tree.contract_stats()\n\n d = tree.size_dict[ind]\n if project is None:\n # we are slicing the index\n si = SliceInfo(ind not in tree.output, ind, d, None)\n tree.multiplicity = tree.multiplicity * d\n else:\n si = SliceInfo(ind not in tree.output, ind, 1, project)\n\n # update the ordered slice information dictionary, but maintain the\n # order such that output sliced indices always appear first ->\n # enforced by the dataclass SliceInfo ordering\n tree.sliced_inds = {\n si.ind: si for si in sorted((*tree.sliced_inds.values(), si))\n }\n\n for node, node_info in tree.info.items():\n if self.is_leaf(node):\n # handle leaves separately\n i = self.nodeops.node_get_single_el(node)\n term = tree.inputs[i]\n if ind in term:\n # n.b. leaves don't contribute to size, flops or write\n # simply recalculate all information, incl. preprocessing\n tree._remove_node(node)\n tree.sliced_inputs = tree.sliced_inputs | frozenset([i])\n else:\n involved = tree.get_involved(node)\n if ind not in involved:\n # if ind doesn't feature in this node (contraction)\n # -> nothing to do\n continue\n\n # else update all the relevant information about this node\n # -> flops changes for all involved indices\n node_info[\"involved\"] = legs_without(involved, ind)\n old_flops = tree.get_flops(node)\n new_flops = old_flops // d\n node_info[\"flops\"] = new_flops\n tree._flops += new_flops - old_flops\n\n # -> size and write only changes for node legs (output) indices\n legs = tree.get_legs(node)\n if ind in legs:\n node_info[\"legs\"] = legs_without(legs, ind)\n old_size = tree.get_size(node)\n tree._sizes.discard(old_size)\n new_size = old_size // d\n tree._sizes.add(new_size)\n node_info[\"size\"] = new_size\n tree._write += new_size - old_size\n\n # delete info we can't change\n for k in (\n \"inds\",\n \"einsum_eq\",\n \"can_dot\",\n \"tensordot_axes\",\n \"tensordot_perm\",\n ):\n tree.info[node].pop(k, None)\n\n tree.already_optimized.clear()\n tree.contraction_cores.clear()\n\n return tree\n\n remove_ind_ = functools.partialmethod(remove_ind, inplace=True)\n\n def restore_ind(self, ind, inplace=False):\n \"\"\"Restore (unslice or un-project) index ``ind`` to this contraction\n tree, taking care to update all relevant information about each node.\n\n Parameters\n ----------\n ind : str\n The index to restore.\n inplace : bool, optional\n Whether to perform the restoration inplace or not.\n\n Returns\n -------\n ContractionTree\n \"\"\"\n tree = self if inplace else self.copy()\n\n # pop sliced index info\n si = tree.sliced_inds.pop(ind)\n\n # make sure all flops and size information has been populated\n tree.contract_stats()\n tree.multiplicity //= si.size\n\n # handle inputs\n for i, term in enumerate(tree.inputs):\n # this is the original term with all indices\n if ind in term:\n tree._remove_node(self.input_to_node(i))\n if all(ix not in tree.sliced_inds for ix in term):\n # mark this input as not sliced\n tree.sliced_inputs = tree.sliced_inputs - frozenset([i])\n\n # delete and re-add dependent intermediates\n for p, l, r in tree.traverse():\n if ind in tree.get_legs(l) or ind in tree.get_legs(r):\n tree._remove_node(p)\n tree.contract_nodes_pair(l, r, parent=p)\n\n # reset caches\n tree.already_optimized.clear()\n tree.contraction_cores.clear()\n\n return tree\n\n restore_ind_ = functools.partialmethod(restore_ind, inplace=True)\n\n def unslice_rand(self, seed=None, inplace=False):\n \"\"\"Unslice (restore) a random index from this contraction tree.\n\n Parameters\n ----------\n seed : None, int or random.Random, optional\n Random number generator seed.\n inplace : bool, optional\n Whether to perform the unslicing inplace or not.\n\n Returns\n -------\n ContractionTree\n \"\"\"\n rng = get_rng(seed)\n ix = rng.choice(tuple(self.sliced_inds))\n return self.restore_ind(ix, inplace=inplace)\n\n unslice_rand_ = functools.partialmethod(unslice_rand, inplace=True)\n\n def unslice_all(self, inplace=False):\n \"\"\"Unslice (restore) all sliced indices from this contraction tree.\n\n Parameters\n ----------\n inplace : bool, optional\n Whether to perform the unslicing inplace or not.\n\n Returns\n -------\n ContractionTree\n \"\"\"\n tree = self if inplace else self.copy()\n\n for ind in tuple(tree.sliced_inds):\n tree.restore_ind_(ind)\n\n return tree\n\n unslice_all_ = functools.partialmethod(unslice_all, inplace=True)\n\n def calc_subtree_candidates(self, pwr=2, what=\"flops\"):\n # get all intermediate nodes\n candidates = list(self.children)\n\n if what == \"size\":\n weights = [self.get_size(x) for x in candidates]\n\n elif what == \"flops\":\n weights = [self.get_flops(x) for x in candidates]\n\n if pwr == \"log\":\n weights = [math.log2(max(2, w)) for w in weights]\n else:\n max_weight = max(weights)\n # can be bigger than numpy int/float allows\n weights = [float(w / max_weight) ** (1 / pwr) for w in weights]\n\n # sort by descending score\n candidates, weights = zip(\n *sorted(zip(candidates, weights), key=lambda x: -x[1])\n )\n\n return list(candidates), list(weights)\n\n def _subtree_remove_and_optimize(\n self,\n sub_root,\n sub_leaves,\n sub_branches,\n already_optimized,\n node_cost,\n minimize,\n opt,\n pbar,\n ):\n current_cost = node_cost(self, sub_root)\n for node in sub_branches:\n # these are the intermediates *between* leaves and sub-root\n if minimize == \"size\":\n current_cost = max(current_cost, node_cost(self, node))\n else:\n current_cost += node_cost(self, node)\n self._remove_node(node)\n\n # make the optimizer more efficient by supplying accurate cap\n opt.cost_cap = max(2, current_cost)\n\n # and reoptimize the leaves\n self.contract_nodes(sub_leaves, optimize=opt, grandparent=sub_root)\n already_optimized.add(sub_leaves)\n\n if pbar is not None:\n pbar.update()\n pbar.set_description(_describe_tree(self), refresh=False)\n\n def _subtree_reconfigure_descend(\n self,\n subtree_size,\n subtree_search,\n maxiter,\n seed,\n minimize,\n opt,\n already_optimized,\n node_cost,\n pbar,\n ):\n candidates = [self.root]\n any_modified = False\n\n def _possibly_add_children(sub_root, any_modified):\n if self.get_extent(sub_root) > subtree_size:\n # possibly extend with node children, if not close to bottom\n lnode, rnode = self.children[sub_root]\n if self.get_extent(lnode) >= 2:\n candidates.append(lnode)\n if self.get_extent(rnode) >= 2:\n candidates.append(rnode)\n\n if len(candidates) == 0:\n # exhausted queue\n if any_modified:\n # but have made *any* changes -> go again from top\n candidates.append(self.root)\n any_modified = False\n\n return any_modified\n\n r = 0\n while candidates and r < maxiter:\n sub_root = candidates.pop(0)\n\n # get a subtree to possibly reconfigure\n sub_leaves, sub_branches = self.get_subtree(\n sub_root, size=subtree_size, search=subtree_search, seed=seed\n )\n\n # check if its already been optimized\n sub_leaves = frozenset(sub_leaves)\n if sub_leaves in already_optimized:\n any_modified = _possibly_add_children(sub_root, any_modified)\n continue\n\n # else remove the branches, keeping track of current cost\n self._subtree_remove_and_optimize(\n sub_root,\n sub_leaves,\n sub_branches,\n already_optimized,\n node_cost,\n minimize,\n opt,\n pbar,\n )\n any_modified = _possibly_add_children(sub_root, True)\n r += 1\n\n def _subtree_reconfigure_rand_select(\n self,\n subtree_size,\n subtree_search,\n weight_what,\n weight_pwr,\n select,\n maxiter,\n seed,\n minimize,\n opt,\n already_optimized,\n node_cost,\n pbar,\n ):\n if select == \"random\":\n rng = get_rng(seed)\n else:\n rng = None\n if select == \"max\":\n i = 0\n elif select == \"min\":\n i = -1\n\n candidates, weights = self.calc_subtree_candidates(\n pwr=weight_pwr, what=weight_what\n )\n\n r = 0\n while candidates and r < maxiter:\n if rng is not None:\n (i,) = rng.choices(range(len(candidates)), weights=weights)\n\n weights.pop(i)\n sub_root = candidates.pop(i)\n\n # get a subtree to possibly reconfigure\n sub_leaves, sub_branches = self.get_subtree(\n sub_root, size=subtree_size, search=subtree_search, seed=seed\n )\n\n # check if its already been optimized\n sub_leaves = frozenset(sub_leaves)\n if sub_leaves in already_optimized:\n continue\n\n # else remove the branches, keeping track of current cost\n self._subtree_remove_and_optimize(\n sub_root,\n sub_leaves,\n sub_branches,\n already_optimized,\n node_cost,\n minimize,\n opt,\n pbar,\n )\n\n # if we have reconfigured simply re-add all candidates\n candidates, weights = self.calc_subtree_candidates(\n pwr=weight_pwr, what=weight_what\n )\n\n r += 1\n\n def subtree_reconfigure(\n self,\n subtree_size=8,\n subtree_search=\"bfs\",\n weight_what=\"flops\",\n weight_pwr=2,\n select=\"max\",\n maxiter=500,\n seed=None,\n minimize=None,\n optimize=None,\n inplace=False,\n progbar=False,\n ):\n \"\"\"Reconfigure subtrees of this tree with locally optimal paths.\n\n Parameters\n ----------\n subtree_size : int, optional\n The size of subtree to consider. Cost is exponential in this.\n subtree_search : {'bfs', 'dfs', 'random'}, optional\n How to build the subtrees:\n\n - 'bfs': breadth-first-search creating balanced subtrees\n - 'dfs': depth-first-search creating imbalanced subtrees\n - 'random': random subtree building\n\n weight_what : {'flops', 'size'}, optional\n When assessing nodes to build and optimize subtrees from whether to\n score them by the (local) contraction cost, or tensor size.\n weight_pwr : int, optional\n When assessing nodes to build and optimize subtrees from, how to\n scale their score into a probability: ``score**(1 / weight_pwr)``.\n The larger this is the more explorative the algorithm is when\n ``select='random'``.\n select : {'descend', 'max', 'min', 'random'}, optional\n What order to select node subtrees to optimize:\n\n - 'descend': start from the root and then descend into children. In\n this case the weights and weight_pwr are ignored since this is a\n deterministic order.\n - 'max': choose the highest score first\n - 'min': choose the lowest score first\n - 'random': choose randomly weighted on score - see ``weight_pwr``.\n\n maxiter : int, optional\n How many subtree optimizations to perform, the algorithm can\n terminate before this if all subtrees have been optimized.\n seed : int, optional\n A random seed (seeds python system random module).\n minimize : {'flops', 'size'}, optional\n Whether to minimize with respect to contraction flops or size.\n inplace : bool, optional\n Whether to perform the reconfiguration inplace or not.\n progbar : bool, optional\n Whether to show live progress of the reconfiguration.\n\n Returns\n -------\n ContractionTree\n \"\"\"\n tree = self if inplace else self.copy()\n tree.reset_contraction_indices()\n\n # ensure these have been computed and thus are being tracked\n tree.contract_stats()\n\n if minimize is None:\n minimize = self.get_default_objective()\n scorer = get_score_fn(minimize)\n node_cost = getattr(scorer, \"cost_local_tree_node\", lambda _: 2)\n\n if optimize is None:\n from .pathfinders.path_basic import OptimalOptimizer\n\n minimize = scorer.get_dynamic_programming_minimize()\n opt = OptimalOptimizer(minimize=minimize)\n else:\n opt = optimize\n\n # different caches as we might want to reconfigure one before other\n tree.already_optimized.setdefault(minimize, set())\n already_optimized = tree.already_optimized[minimize]\n\n if progbar:\n import tqdm\n\n pbar = tqdm.tqdm()\n pbar.set_description(_describe_tree(tree), refresh=False)\n else:\n pbar = None\n\n try:\n reconf_kwargs = {\n \"subtree_size\": subtree_size,\n \"subtree_search\": subtree_search,\n \"maxiter\": maxiter,\n \"seed\": seed,\n \"minimize\": minimize,\n \"opt\": opt,\n \"already_optimized\": already_optimized,\n \"node_cost\": node_cost,\n \"pbar\": pbar,\n }\n\n if select == \"descend\":\n tree._subtree_reconfigure_descend(**reconf_kwargs)\n else:\n reconf_kwargs[\"weight_what\"] = weight_what\n reconf_kwargs[\"weight_pwr\"] = weight_pwr\n reconf_kwargs[\"select\"] = select\n tree._subtree_reconfigure_rand_select(**reconf_kwargs)\n\n finally:\n if progbar:\n pbar.close()\n\n return tree\n\n subtree_reconfigure_ = functools.partialmethod(\n subtree_reconfigure, inplace=True\n )\n\n def subtree_reconfigure_forest(\n self,\n num_trees=8,\n num_restarts=10,\n restart_fraction=0.5,\n subtree_maxiter=100,\n subtree_size=10,\n subtree_search=(\"random\", \"bfs\"),\n subtree_select=(\"random\",),\n subtree_weight_what=(\"flops\", \"size\"),\n subtree_weight_pwr=(2,),\n parallel=\"auto\",\n parallel_maxiter_steps=4,\n minimize=None,\n seed=None,\n progbar=False,\n inplace=False,\n ):\n \"\"\"'Forested' version of ``subtree_reconfigure`` which is more\n explorative and can be parallelized. It stochastically generates\n a 'forest' reconfigured trees, then only keeps some fraction of these\n to generate the next forest.\n\n Parameters\n ----------\n num_trees : int, optional\n The number of trees to reconfigure at each stage.\n num_restarts : int, optional\n The number of times to halt, prune and then restart the\n tree reconfigurations.\n restart_fraction : float, optional\n The fraction of trees to keep at each stage and generate the next\n forest from.\n subtree_maxiter : int, optional\n Number of subtree reconfigurations per step.\n ``num_restarts * subtree_maxiter`` is the max number of total\n subtree reconfigurations for the final tree produced.\n subtree_size : int, optional\n The size of subtrees to search for and reconfigure.\n subtree_search : tuple[{'random', 'bfs', 'dfs'}], optional\n Tuple of options for the ``search`` kwarg of\n :meth:`ContractionTree.subtree_reconfigure` to randomly sample.\n subtree_select : tuple[{'random', 'max', 'min'}], optional\n Tuple of options for the ``select`` kwarg of\n :meth:`ContractionTree.subtree_reconfigure` to randomly sample.\n subtree_weight_what : tuple[{'flops', 'size'}], optional\n Tuple of options for the ``weight_what`` kwarg of\n :meth:`ContractionTree.subtree_reconfigure` to randomly sample.\n subtree_weight_pwr : tuple[int], optional\n Tuple of options for the ``weight_pwr`` kwarg of\n :meth:`ContractionTree.subtree_reconfigure` to randomly sample.\n parallel : 'auto', False, True, int, or distributed.Client\n Whether to parallelize the search.\n parallel_maxiter_steps : int, optional\n If parallelizing, how many steps to break each reconfiguration into\n in order to evenly saturate many processes.\n minimize : {'flops', 'size', ..., Objective}, optional\n Whether to minimize the total flops or maximum size of the\n contraction tree.\n seed : None, int or random.Random, optional\n A random seed to use.\n progbar : bool, optional\n Whether to show live progress.\n inplace : bool, optional\n Whether to perform the subtree reconfiguration inplace.\n\n Returns\n -------\n ContractionTree\n \"\"\"\n tree = self if inplace else self.copy()\n tree.reset_contraction_indices()\n\n # candidate trees\n num_keep = max(1, int(num_trees * restart_fraction))\n\n # how to rank the trees\n if minimize is None:\n minimize = self.get_default_objective()\n score = get_score_fn(minimize)\n\n rng = get_rng(seed)\n\n # set up the initial 'forest' and parallel machinery\n pool = parse_parallel_arg(parallel)\n is_scatter_pool = can_scatter(pool)\n if is_scatter_pool:\n is_worker = maybe_leave_pool(pool)\n # store the trees as futures for the entire process\n forest = [scatter(pool, tree)]\n maxiter = subtree_maxiter // parallel_maxiter_steps\n else:\n forest = [tree]\n maxiter = subtree_maxiter\n\n if progbar:\n import tqdm\n\n pbar = tqdm.tqdm(total=num_restarts)\n pbar.set_description(_describe_tree(tree), refresh=False)\n\n try:\n for _ in range(num_restarts):\n # on the next round take only the best trees\n forest = itertools.cycle(forest[:num_keep])\n\n # select some random configurations\n saplings = [\n {\n \"tree\": next(forest),\n \"maxiter\": maxiter,\n \"minimize\": minimize,\n \"subtree_size\": subtree_size,\n \"subtree_search\": rng.choice(subtree_search),\n \"select\": rng.choice(subtree_select),\n \"weight_pwr\": rng.choice(subtree_weight_pwr),\n \"weight_what\": rng.choice(subtree_weight_what),\n }\n for _ in range(num_trees)\n ]\n\n if pool is None:\n forest = [_reconfigure_tree(**s) for s in saplings]\n res = [{\"tree\": t, **_get_tree_info(t)} for t in forest]\n elif not is_scatter_pool:\n forest_futures = [\n submit(pool, _reconfigure_tree, **s) for s in saplings\n ]\n forest = [f.result() for f in forest_futures]\n res = [{\"tree\": t, **_get_tree_info(t)} for t in forest]\n else:\n # submit in smaller steps to saturate processes\n for _ in range(parallel_maxiter_steps):\n for s in saplings:\n s[\"tree\"] = submit(pool, _reconfigure_tree, **s)\n\n # compute scores remotely then gather\n forest_futures = [s[\"tree\"] for s in saplings]\n res_futures = [\n submit(pool, _get_tree_info, t) for t in forest_futures\n ]\n res = [\n {\"tree\": tree_future, **res_future.result()}\n for tree_future, res_future in zip(\n forest_futures, res_futures\n )\n ]\n\n # update the order of the new forest\n res.sort(key=score)\n forest = [r[\"tree\"] for r in res]\n\n if progbar:\n pbar.update()\n if pool is None:\n d = _describe_tree(forest[0])\n else:\n d = submit(pool, _describe_tree, forest[0]).result()\n pbar.set_description(d, refresh=False)\n\n finally:\n if progbar:\n pbar.close()\n\n if is_scatter_pool:\n tree.set_state_from(forest[0].result())\n maybe_rejoin_pool(is_worker, pool)\n else:\n tree.set_state_from(forest[0])\n\n return tree\n\n subtree_reconfigure_forest_ = functools.partialmethod(\n subtree_reconfigure_forest, inplace=True\n )\n\n simulated_anneal = simulated_anneal_tree\n simulated_anneal_ = functools.partialmethod(simulated_anneal, inplace=True)\n parallel_temper = parallel_temper_tree\n parallel_temper_ = functools.partialmethod(parallel_temper, inplace=True)\n\n def slice(\n self,\n target_size=None,\n target_overhead=None,\n target_slices=None,\n temperature=0.01,\n minimize=None,\n allow_outer=True,\n max_repeats=16,\n reslice=False,\n seed=None,\n inplace=False,\n ):\n \"\"\"Slice this tree (turn some indices into indices which are explicitly\n summed over rather than being part of contractions). The indices are\n stored in ``tree.sliced_inds``, and the contraction width updated to\n take account of the slicing. Calling ``tree.contract(arrays)`` moreover\n which automatically perform the slicing and summation.\n\n Parameters\n ----------\n target_size : int, optional\n The target number of entries in the largest tensor of the sliced\n contraction. The search algorithm will terminate after this is\n reached.\n target_slices : int, optional\n The target or minimum number of 'slices' to consider - individual\n contractions after slicing indices. The search algorithm will\n terminate after this is breached. This is on top of the current\n number of slices.\n target_overhead : float, optional\n The target increase in total number of floating point operations.\n For example, a value of ``2.0`` will terminate the search just\n before the cost of computing all the slices individually breaches\n twice that of computing the original contraction all at once.\n temperature : float, optional\n How much to randomize the repeated search.\n minimize : {'flops', 'size', ..., Objective}, optional\n Which metric to score the overhead increase against.\n allow_outer : bool, optional\n Whether to allow slicing of outer indices.\n max_repeats : int, optional\n How many times to repeat the search with a slight randomization.\n reslice : bool, optional\n Whether to reslice the tree, i.e. first remove all currently\n sliced indices and start the search again. Generally any 'good'\n sliced indices will be easily found again.\n seed : None, int or random.Random, optional\n A random seed or generator to use for the search.\n inplace : bool, optional\n Whether the remove the indices from this tree inplace or not.\n\n Returns\n -------\n ContractionTree\n\n See Also\n --------\n SliceFinder, ContractionTree.slice_and_reconfigure\n \"\"\"\n from .slicer import SliceFinder\n\n if minimize is None:\n minimize = self.get_default_objective()\n\n tree = self if inplace else self.copy()\n\n if reslice:\n if target_slices is not None:\n target_slices *= tree.nslices\n tree.unslice_all_()\n\n sf = SliceFinder(\n tree,\n target_size=target_size,\n target_overhead=target_overhead,\n target_slices=target_slices,\n temperature=temperature,\n minimize=minimize,\n allow_outer=allow_outer,\n seed=seed,\n )\n\n ix_sl, _ = sf.search(max_repeats)\n for ix in ix_sl:\n tree.remove_ind_(ix)\n\n return tree\n\n slice_ = functools.partialmethod(slice, inplace=True)\n\n def slice_and_reconfigure(\n self,\n target_size,\n step_size=2,\n temperature=0.01,\n minimize=None,\n allow_outer=True,\n max_repeats=16,\n reslice=False,\n reconf_opts=None,\n progbar=False,\n inplace=False,\n ):\n \"\"\"Interleave slicing (removing indices into an exterior sum) with\n subtree reconfiguration to minimize the overhead induced by this\n slicing.\n\n Parameters\n ----------\n target_size : int\n Slice the tree until the maximum intermediate size is this or\n smaller.\n step_size : int, optional\n The minimum size reduction to try and achieve before switching to a\n round of subtree reconfiguration.\n temperature : float, optional\n The temperature to supply to ``SliceFinder`` for searching for\n indices.\n minimize : {'flops', 'size', ..., Objective}, optional\n The metric to minimize when slicing and reconfiguring subtrees.\n max_repeats : int, optional\n The number of slicing attempts to perform per search.\n progbar : bool, optional\n Whether to show live progress.\n inplace : bool, optional\n Whether to perform the slicing and reconfiguration inplace.\n reconf_opts : None or dict, optional\n Supplied to\n :meth:`ContractionTree.subtree_reconfigure` or\n :meth:`ContractionTree.subtree_reconfigure_forest`, depending on\n `'forested'` key value.\n \"\"\"\n tree = self if inplace else self.copy()\n\n reconf_opts = {} if reconf_opts is None else dict(reconf_opts)\n\n if minimize is None:\n minimize = self.get_default_objective()\n minimize = get_score_fn(minimize)\n\n reconf_opts.setdefault(\"minimize\", minimize)\n forested_reconf = reconf_opts.pop(\"forested\", False)\n\n if progbar:\n import tqdm\n\n pbar = tqdm.tqdm()\n pbar.set_description(_describe_tree(tree), refresh=False)\n\n try:\n while tree.max_size() > target_size:\n tree.slice_(\n temperature=temperature,\n target_slices=step_size,\n minimize=minimize,\n allow_outer=allow_outer,\n max_repeats=max_repeats,\n reslice=reslice,\n )\n if forested_reconf:\n tree.subtree_reconfigure_forest_(**reconf_opts)\n else:\n tree.subtree_reconfigure_(**reconf_opts)\n\n if progbar:\n pbar.update()\n pbar.set_description(_describe_tree(tree), refresh=False)\n finally:\n if progbar:\n pbar.close()\n\n return tree\n\n slice_and_reconfigure_ = functools.partialmethod(\n slice_and_reconfigure, inplace=True\n )\n\n def slice_and_reconfigure_forest(\n self,\n target_size,\n step_size=2,\n num_trees=8,\n restart_fraction=0.5,\n temperature=0.02,\n max_repeats=32,\n reslice=False,\n minimize=None,\n allow_outer=True,\n parallel=\"auto\",\n progbar=False,\n inplace=False,\n reconf_opts=None,\n ):\n \"\"\"'Forested' version of :meth:`ContractionTree.slice_and_reconfigure`.\n This maintains a 'forest' of trees with different slicing and subtree\n reconfiguration attempts, pruning the worst at each step and generating\n a new forest from the best.\n\n Parameters\n ----------\n target_size : int\n Slice the tree until the maximum intermediate size is this or\n smaller.\n step_size : int, optional\n The minimum size reduction to try and achieve before switching to a\n round of subtree reconfiguration.\n num_restarts : int, optional\n The number of times to halt, prune and then restart the\n tree reconfigurations.\n restart_fraction : float, optional\n The fraction of trees to keep at each stage and generate the next\n forest from.\n temperature : float, optional\n The temperature at which to randomize the sliced index search.\n max_repeats : int, optional\n The number of slicing attempts to perform per search.\n parallel : 'auto', False, True, int, or distributed.Client\n Whether to parallelize the search.\n progbar : bool, optional\n Whether to show live progress.\n inplace : bool, optional\n Whether to perform the slicing and reconfiguration inplace.\n reconf_opts : None or dict, optional\n Supplied to\n :meth:`ContractionTree.slice_and_reconfigure`.\n\n Returns\n -------\n ContractionTree\n \"\"\"\n tree = self if inplace else self.copy()\n tree.reset_contraction_indices()\n\n # candidate trees\n num_keep = max(1, int(num_trees * restart_fraction))\n\n # how to rank the trees\n if minimize is None:\n minimize = self.get_default_objective()\n score = get_score_fn(minimize)\n\n # set up the initial 'forest' and parallel machinery\n pool = parse_parallel_arg(parallel)\n is_scatter_pool = can_scatter(pool)\n if is_scatter_pool:\n is_worker = maybe_leave_pool(pool)\n # store the trees as futures for the entire process\n forest = [scatter(pool, tree)]\n else:\n forest = [tree]\n\n if progbar:\n import tqdm\n\n pbar = tqdm.tqdm()\n pbar.set_description(_describe_tree(tree), refresh=False)\n\n next_size = tree.max_size()\n\n try:\n while True:\n next_size //= step_size\n\n # on the next round take only the best trees\n forest = itertools.cycle(forest[:num_keep])\n\n saplings = [\n {\n \"tree\": next(forest),\n \"target_size\": next_size,\n \"step_size\": step_size,\n \"temperature\": temperature,\n \"max_repeats\": max_repeats,\n \"reconf_opts\": reconf_opts,\n \"allow_outer\": allow_outer,\n \"reslice\": reslice,\n }\n for _ in range(num_trees)\n ]\n\n if pool is None:\n forest = [\n _slice_and_reconfigure_tree(**s) for s in saplings\n ]\n res = [{\"tree\": t, **_get_tree_info(t)} for t in forest]\n\n elif not is_scatter_pool:\n # simple pool with no pass by reference\n forest_futures = [\n submit(pool, _slice_and_reconfigure_tree, **s)\n for s in saplings\n ]\n forest = [f.result() for f in forest_futures]\n res = [{\"tree\": t, **_get_tree_info(t)} for t in forest]\n\n else:\n forest_futures = [\n submit(pool, _slice_and_reconfigure_tree, **s)\n for s in saplings\n ]\n\n # compute scores remotely then gather\n res_futures = [\n submit(pool, _get_tree_info, t) for t in forest_futures\n ]\n res = [\n {\"tree\": tree_future, **res_future.result()}\n for tree_future, res_future in zip(\n forest_futures, res_futures\n )\n ]\n\n # we want to sort by flops, but also favour sampling as\n # many different sliced index combos as possible\n # ~ [1, 1, 1, 2, 2, 3] -> [1, 2, 3, 1, 2, 1]\n res.sort(key=score)\n res = list(\n interleave(\n groupby(lambda r: r[\"sliced_ind_set\"], res).values()\n )\n )\n\n # update the order of the new forest\n forest = [r[\"tree\"] for r in res]\n\n if progbar:\n pbar.update()\n if pool is None:\n d = _describe_tree(forest[0])\n else:\n d = submit(pool, _describe_tree, forest[0]).result()\n pbar.set_description(d, refresh=False)\n\n if res[0][\"size\"] <= target_size:\n break\n\n finally:\n if progbar:\n pbar.close()\n\n if is_scatter_pool:\n tree.set_state_from(forest[0].result())\n maybe_rejoin_pool(is_worker, pool)\n else:\n tree.set_state_from(forest[0])\n\n return tree\n\n slice_and_reconfigure_forest_ = functools.partialmethod(\n slice_and_reconfigure_forest, inplace=True\n )\n\n def compressed_reconfigure(\n self,\n minimize=None,\n order_only=False,\n max_nodes=\"auto\",\n max_time=None,\n local_score=None,\n exploration_power=0,\n best_score=None,\n progbar=False,\n inplace=False,\n ):\n \"\"\"Reconfigure this tree according to ``peak_size_compressed``.\n\n Parameters\n ----------\n chi : int\n The maximum bond dimension to consider.\n order_only : bool, optional\n Whether to only consider the ordering of the current tree\n contractions, or all possible contractions, starting with the\n current.\n max_nodes : int, optional\n Set the maximum number of contraction steps to consider.\n max_time : float, optional\n Set the maximum time to spend on the search.\n local_score : callable, optional\n A function that assigns a score to a potential contraction, with a\n lower score giving more priority to explore that contraction\n earlier. It should have signature::\n\n local_score(step, new_score, dsize, new_size)\n\n where ``step`` is the number of steps so far, ``new_score`` is the\n score of the contraction so far, ``dsize`` is the change in memory\n by the current step, and ``new_size`` is the new memory size after\n contraction.\n exploration_power : float, optional\n If not ``0.0``, the inverse power to which the step is raised in\n the default local score function. Higher values favor exploring\n more promising branches early on - at the cost of increased memory.\n Ignored if ``local_score`` is supplied.\n best_score : float, optional\n Manually specify an upper bound for best score found so far.\n progbar : bool, optional\n If ``True``, display a progress bar.\n inplace : bool, optional\n Whether to perform the reconfiguration inplace on this tree.\n\n Returns\n -------\n ContractionTree\n \"\"\"\n from .experimental.path_compressed_branchbound import (\n CompressedExhaustive,\n )\n\n if minimize is None:\n minimize = self.get_default_objective()\n\n if max_nodes == \"auto\":\n if max_time is None:\n max_nodes = max(10_000, self.N**2)\n else:\n max_nodes = float(\"inf\")\n\n opt = CompressedExhaustive(\n minimize=minimize,\n local_score=local_score,\n max_nodes=max_nodes,\n max_time=max_time,\n exploration_power=exploration_power,\n best_score=best_score,\n progbar=progbar,\n )\n opt.setup(self.inputs, self.output, self.size_dict)\n opt.explore_path(self.get_path_surface(), restrict=order_only)\n\n # rtree = opt.search(self.inputs, self.output, self.size_dict)\n\n opt.run(self.inputs, self.output, self.size_dict)\n ssa_path = opt.ssa_path\n # ssa_path = opt(self.inputs, self.output, self.size_dict)\n rtree = self.__class__.from_path(\n self.inputs,\n self.output,\n self.size_dict,\n ssa_path=ssa_path,\n objective=minimize,\n )\n if inplace:\n self.set_state_from(rtree)\n rtree = self\n\n rtree.reset_contraction_indices()\n return rtree\n\n compressed_reconfigure_ = functools.partialmethod(\n compressed_reconfigure, inplace=True\n )\n\n def windowed_reconfigure(\n self,\n minimize=None,\n order_only=False,\n window_size=20,\n max_iterations=100,\n max_window_tries=1000,\n score_temperature=0.0,\n queue_temperature=1.0,\n scorer=None,\n queue_scorer=None,\n seed=None,\n inplace=False,\n progbar=False,\n **kwargs,\n ):\n from .pathfinders.path_compressed import WindowedOptimizer\n\n if minimize is None:\n minimize = self.get_default_objective()\n\n wo = WindowedOptimizer(\n self.inputs,\n self.output,\n self.size_dict,\n minimize=minimize,\n ssa_path=self.get_ssa_path(),\n seed=seed,\n )\n\n wo.refine(\n window_size=window_size,\n max_iterations=max_iterations,\n order_only=order_only,\n max_window_tries=max_window_tries,\n score_temperature=score_temperature,\n queue_temperature=queue_temperature,\n scorer=scorer,\n queue_scorer=queue_scorer,\n progbar=progbar,\n **kwargs,\n )\n ssa_path = wo.get_ssa_path()\n\n rtree = self.__class__.from_path(\n self.inputs,\n self.output,\n self.size_dict,\n ssa_path=ssa_path,\n objective=minimize,\n )\n\n if inplace:\n self.set_state_from(rtree)\n rtree = self\n\n rtree.reset_contraction_indices()\n return rtree\n\n windowed_reconfigure_ = functools.partialmethod(\n windowed_reconfigure, inplace=True\n )\n\n def flat_tree(self, order=None):\n \"\"\"Create a nested tuple representation of the contraction tree like::\n\n ((0, (1, 2)), ((3, 4), ((5, (6, 7)), (8, 9))))\n\n Such that the contraction will progress like::\n\n ((0, (1, 2)), ((3, 4), ((5, (6, 7)), (8, 9))))\n ((0, 12), (34, ((5, 67), 89)))\n (012, (34, (567, 89)))\n (012, (34, 56789))\n (012, 3456789)\n 0123456789\n\n Where each integer represents a leaf (i.e. single element node).\n \"\"\"\n tups = dict(zip(self.gen_leaves(), range(self.N)))\n\n for parent, l, r in self.traverse(order=order):\n tups[parent] = tups[l], tups[r]\n\n return tups[self.root]\n\n def get_leaves_ordered(self):\n \"\"\"Return the list of leaves as ordered by the contraction tree.\n\n Returns\n -------\n tuple[node_type]\n \"\"\"\n if not self.is_complete():\n raise ValueError(\"Can't order the leaves until tree is complete.\")\n\n return tuple(\n node\n for node in itertools.chain.from_iterable(self.traverse())\n if self.is_leaf(node)\n )\n\n def get_path(self, order=None):\n \"\"\"Generate a standard path (with linear recycled ids) from the\n contraction tree.\n\n Parameters\n ----------\n order : None, \"dfs\", or callable, optional\n How to order the contractions within the tree. If a callable is\n given (which should take a node as its argument), try to contract\n nodes that minimize this function first.\n\n Returns\n -------\n path: tuple[tuple[int, int]]\n \"\"\"\n from bisect import bisect_left\n\n ssa = self.N\n ssas = list(range(ssa))\n node_to_ssa = dict(zip(self.gen_leaves(), ssas))\n path = []\n\n for parent, left, right in self.traverse(order=order):\n # map nodes to ssas\n lssa = node_to_ssa[left]\n rssa = node_to_ssa[right]\n # map ssas to linear indices, using bisection\n i, j = sorted((bisect_left(ssas, lssa), bisect_left(ssas, rssa)))\n # 'contract' nodes\n ssas.pop(j)\n ssas.pop(i)\n path.append((i, j))\n ssas.append(ssa)\n # update mapping\n node_to_ssa[parent] = ssa\n ssa += 1\n\n return tuple(path)\n\n path = deprecated(get_path, \"path\", \"get_path\")\n\n def get_numpy_path(self, order=None):\n \"\"\"Generate a path compatible with the `optimize` kwarg of\n `numpy.einsum`.\n \"\"\"\n return [\"einsum_path\", *self.get_path(order=order)]\n\n def get_ssa_path(self, order=None):\n \"\"\"Generate a single static assignment path from the contraction tree.\n\n Parameters\n ----------\n order : None, \"dfs\", or callable, optional\n How to order the contractions within the tree. If a callable is\n given (which should take a node as its argument), try to contract\n nodes that minimize this function first.\n\n Returns\n -------\n ssa_path: tuple[tuple[int, int]]\n \"\"\"\n ssa_path = []\n pos = dict(zip(self.gen_leaves(), range(self.N)))\n\n for parent, l, r in self.traverse(order=order):\n i, j = sorted((pos[l], pos[r]))\n ssa_path.append((i, j))\n pos[parent] = len(ssa_path) + self.N - 1\n\n return tuple(ssa_path)\n\n ssa_path = deprecated(get_ssa_path, \"ssa_path\", \"get_ssa_path\")\n\n def surface_order(self, node):\n return (self.get_extent(node), self.get_centrality(node))\n\n def set_surface_order_from_path(self, ssa_path):\n\n # first get dict from contractions to parents (don't usually store)\n parent_map = {}\n for p, l, r in self.traverse():\n parent_map[frozenset([l, r])] = p\n\n # then traverse up in given ssa_path order,\n # assigning parent node ordering 'score' incrementally\n parent_scores = {}\n node_map = {i: n for i, n in enumerate(self.gen_leaves())}\n for j, p in enumerate(ssa_path):\n lr = frozenset(node_map[i] for i in p)\n p = parent_map[lr]\n parent_scores[p] = j\n node_map[self.N + j] = p\n\n self.surface_order = functools.partial(\n get_with_default, obj=parent_scores, default=float(\"inf\")\n )\n\n def get_path_surface(self):\n return self.get_path(order=self.surface_order)\n\n path_surface = deprecated(\n get_path_surface, \"path_surface\", \"get_path_surface\"\n )\n\n def get_ssa_path_surface(self):\n return self.get_ssa_path(order=self.surface_order)\n\n ssa_path_surface = deprecated(\n get_ssa_path_surface, \"ssa_path_surface\", \"get_ssa_path_surface\"\n )\n\n def get_spans(self):\n \"\"\"Get all (which could mean none) potential embeddings of this\n contraction tree into a spanning tree of the original graph.\n\n Returns\n -------\n tuple[dict[frozenset[int], frozenset[int]]]\n \"\"\"\n ind_to_term = collections.defaultdict(set)\n for i, term in enumerate(self.inputs):\n for ix in term:\n ind_to_term[ix].add(i)\n\n def boundary_pairs(node):\n \"\"\"Get nodes along the boundary of the bipartition represented by\n ``node``.\n \"\"\"\n pairs = set()\n involved = self.get_involved(node)\n legs = self.get_legs(node)\n removed = [ix for ix in involved if ix not in legs]\n for ix in removed:\n # for every index across the contraction\n l1, l2 = ind_to_term[ix]\n\n # can either span from left to right or right to left\n pairs.add((l1, l2))\n pairs.add((l2, l1))\n\n return pairs\n\n # first span choice is any nodes across the top level bipart\n candidates = [\n {\n # which intermedate nodes map to which leaf nodes\n \"map\": {self.root: self.input_to_node(l2)},\n # the leaf nodes in the spanning tree\n \"spine\": {l1, l2},\n }\n for l1, l2 in boundary_pairs(self.root)\n ]\n\n for _, l, r in self.descend():\n for child in (r, l):\n # for each current candidate check all the possible extensions\n for _ in range(len(candidates)):\n cand = candidates.pop(0)\n\n # don't need to do anything for\n if self.is_leaf(child):\n candidates.append(\n {\n \"map\": {child: child, **cand[\"map\"]},\n \"spine\": cand[\"spine\"].copy(),\n }\n )\n\n for l1, l2 in boundary_pairs(child):\n if (l1 in cand[\"spine\"]) or (l2 not in cand[\"spine\"]):\n # pair does not merge inwards into spine\n continue\n\n # valid extension of spanning tree\n candidates.append(\n {\n \"map\": {\n child: self.input_to_node(l2),\n **cand[\"map\"],\n },\n \"spine\": cand[\"spine\"] | {l1, l2},\n }\n )\n\n return tuple(c[\"map\"] for c in candidates)\n\n def compute_centralities(self, combine=\"mean\"):\n \"\"\"Compute a centrality for every node in this contraction tree.\"\"\"\n hg = self.get_hypergraph(accel=\"auto\")\n cents = hg.simple_centrality()\n\n for i, leaf in enumerate(self.gen_leaves()):\n self.info[leaf][\"centrality\"] = cents[i]\n\n combine = {\n \"mean\": lambda x, y: (x + y) / 2,\n \"sum\": lambda x, y: x + y,\n \"max\": max,\n \"min\": min,\n }.get(combine, combine)\n\n for p, l, r in self.traverse(\"dfs\"):\n self.info[p][\"centrality\"] = combine(\n self.info[l][\"centrality\"], self.info[r][\"centrality\"]\n )\n\n def get_hypergraph(self, accel=False):\n \"\"\"Get a hypergraph representing the uncontracted network (i.e. the\n leaves).\n \"\"\"\n return get_hypergraph(self.inputs, self.output, self.size_dict, accel)\n\n def reset_contraction_indices(self):\n \"\"\"Reset all information regarding a) the explicit contraction indices\n ordering and b) cached contraction expressions. This should probably be\n called any time structural changes are made to the tree, e.g.\n reconfiguration.\n \"\"\"\n # delete all derived information\n # (note legs, involved, etc. are order invariant so we can keep those)\n for node in self.children:\n for k in (\n \"inds\",\n \"einsum_eq\",\n \"can_dot\",\n \"tensordot_axes\",\n \"tensordot_perm\",\n ):\n self.info[node].pop(k, None)\n\n # invalidate any compiled contractions\n self.contraction_cores.clear()\n\n def sort_contraction_indices(\n self,\n priority=\"flops\",\n make_output_contig=True,\n make_contracted_contig=True,\n reset=True,\n ):\n \"\"\"Set explicit orders for the contraction indices of this self to\n optimize for one of two things: contiguity in contracted ('k') indices,\n or contiguity of left and right output ('m' and 'n') indices.\n\n Parameters\n ----------\n priority : {'flops', 'size', 'root', 'leaves'}, optional\n Which order to process the intermediate nodes in. Later nodes\n re-sort previous nodes so are more likely to keep their ordering.\n E.g. for 'flops' the mostly costly contracton will be process last\n and thus will be guaranteed to have its indices exactly sorted.\n make_output_contig : bool, optional\n When processing a pairwise contraction, sort the parent contraction\n indices so that the order of indices is the order they appear\n from left to right in the two child (input) tensors.\n make_contracted_contig : bool, optional\n When processing a pairwise contraction, sort the child (input)\n tensor indices so that all contracted indices appear contiguously.\n reset : bool, optional\n Reset all indices to the default order before sorting.\n \"\"\"\n if reset:\n self.reset_contraction_indices()\n\n if priority == \"flops\":\n nodes = sorted(\n self.children.items(), key=lambda x: self.get_flops(x[0])\n )\n elif priority == \"size\":\n nodes = sorted(\n self.children.items(), key=lambda x: self.get_size(x[0])\n )\n elif priority == \"root\":\n nodes = ((p, (l, r)) for p, l, r in self.traverse())\n elif priority == \"leaves\":\n nodes = ((p, (l, r)) for p, l, r in self.descend())\n else:\n raise ValueError(priority)\n\n for p, (l, r) in nodes:\n p_inds, l_inds, r_inds = map(self.get_inds, (p, l, r))\n\n if make_output_contig and not self.is_root(p):\n # sort indices by whether they appear in the left or right\n # whether this happens before or after the sort below depends\n # on the order we are processing the nodes\n # (avoid root as don't want to modify output)\n\n def psort(ix):\n # group by whether in left or right input\n return (r_inds.find(ix), l_inds.find(ix))\n\n p_inds = \"\".join(sorted(p_inds, key=psort))\n self.info[p][\"inds\"] = p_inds\n\n if make_contracted_contig:\n # sort indices by:\n # 1. if they are going to be contracted\n # 2. what order they appear in the parent indices\n # (but ignore leaf indices)\n if not self.is_leaf(l):\n\n def lsort(ix):\n return (r_inds.find(ix), p_inds.find(ix))\n\n l_inds = \"\".join(sorted(self.get_legs(l), key=lsort))\n self.info[l][\"inds\"] = l_inds\n\n if not self.is_leaf(r):\n\n def rsort(ix):\n return (p_inds.find(ix), l_inds.find(ix))\n\n r_inds = \"\".join(sorted(self.get_legs(r), key=rsort))\n self.info[r][\"inds\"] = r_inds\n\n if not reset:\n # still need to invalidate any cached contraction expressions\n self.contraction_cores.clear()\n\n def print_contractions(self, sort=None, show_brackets=True):\n \"\"\"Print each pairwise contraction, with colorized indices (if\n `colorama` is installed), and other information. The color codes are:\n\n - blue: index appears on left and is kept\n - green: index appears on right and is kept\n - red: contracted index: appears on both sides and is removed\n - pink: batch index: appears on both sides and is kept\n\n Any trivial indices that appear only on one term and not in the output\n are removed and shown by the preprocessing steps.\n\n Parameters\n ----------\n sort : {'flops', 'size'}, optional\n Sort the contractions by either the number of floating point\n operations or the size of the intermediate tensor. By default the\n contraction are show in the order they are performed.\n show_brackets : bool, optional\n Whether to show the brackets around contiguous sections of the same\n type of indices.\n \"\"\"\n try:\n from colorama import Fore\n\n RESET = Fore.RESET\n GREY = Fore.WHITE\n PINK = Fore.MAGENTA\n RED = Fore.RED\n BLUE = Fore.BLUE\n GREEN = Fore.GREEN\n except ImportError:\n RESET = GREY = PINK = RED = BLUE = GREEN = \"\"\n\n entries = []\n\n if self.has_preprocessing():\n for pi, eq in self.preprocessing.items():\n # eq is with canonical indices, reinsert original inputs\n replacer = dict(zip(eq.split(\"->\")[0], self.inputs[pi]))\n eq = \"\".join(replacer.get(c, c) for c in eq)\n print(f\"{GREY}preprocess input {pi}: {RESET}{eq}\")\n print()\n\n for i, (p, l, r) in enumerate(self.traverse()):\n p_legs, l_legs, r_legs = map(self.get_legs, [p, l, r])\n p_inds, l_inds, r_inds = map(self.get_inds, [p, l, r])\n\n # print sizes and flops\n p_flops = self.get_flops(p)\n p_sz, l_sz, r_sz = (\n math.log2(self.get_size(node)) for node in [p, l, r]\n )\n # print whether tensordottable\n if self.get_can_dot(p):\n type_msg = \"tensordot\"\n perm = self.get_tensordot_perm(p)\n if perm is not None:\n # and whether indices match tensordot\n type_msg += \"+perm\"\n else:\n type_msg = \"einsum\"\n\n kpt_brck_l = \"(\" if show_brackets else \"\"\n kpt_brck_r = \")\" if show_brackets else \"\"\n con_brck_l = \"[\" if show_brackets else \"\"\n con_brck_r = \"]\" if show_brackets else \"\"\n hyp_brck_l = \"{\" if show_brackets else \"\"\n hyp_brck_r = \"}\" if show_brackets else \"\"\n\n pa = (\n \"\".join(\n PINK + f\"{hyp_brck_l}{ix}{hyp_brck_r}\"\n if (ix in l_legs) and (ix in r_legs)\n else GREEN + f\"{kpt_brck_l}{ix}{kpt_brck_r}\"\n if ix in r_legs\n else BLUE + ix\n for ix in p_inds\n )\n .replace(f\"){GREEN}(\", \"\")\n .replace(f\"}}{PINK}{{\", \"\")\n )\n la = (\n \"\".join(\n PINK + f\"{hyp_brck_l}{ix}{hyp_brck_r}\"\n if (ix in p_legs) and (ix in r_legs)\n else RED + f\"{con_brck_l}{ix}{con_brck_r}\"\n if ix in r_legs\n else BLUE + ix\n for ix in l_inds\n )\n .replace(f\"]{RED}[\", \"\")\n .replace(f\"}}{PINK}{{\", \"\")\n )\n ra = (\n \"\".join(\n PINK + f\"{hyp_brck_l}{ix}{hyp_brck_r}\"\n if (ix in p_legs) and (ix in l_legs)\n else RED + f\"{con_brck_l}{ix}{con_brck_r}\"\n if ix in l_legs\n else GREEN + ix\n for ix in r_inds\n )\n .replace(f\"]{RED}[\", \"\")\n .replace(f\"}}{PINK}{{\", \"\")\n )\n\n entries.append(\n (\n p,\n f\"{GREY}({i}) cost: {RESET}{p_flops:.1e} \"\n f\"{GREY}widths: {RESET}{l_sz:.1f},{r_sz:.1f}->{p_sz:.1f} \"\n f\"{GREY}type: {RESET}{type_msg}\\n\"\n f\"{GREY}inputs: {la},{ra}{RESET}->\\n\"\n f\"{GREY}output: {pa}\\n\",\n )\n )\n\n if sort == \"flops\":\n entries.sort(key=lambda x: self.get_flops(x[0]), reverse=True)\n if sort == \"size\":\n entries.sort(key=lambda x: self.get_size(x[0]), reverse=True)\n\n entries.append((None, f\"{RESET}\"))\n\n o = \"\\n\".join(entry for _, entry in entries)\n print(o)\n\n # --------------------- Performing the Contraction ---------------------- #\n\n def get_contractor(\n self,\n order=None,\n prefer_einsum=False,\n strip_exponent=False,\n check_zero=False,\n implementation=None,\n autojit=False,\n progbar=False,\n ):\n \"\"\"Get a reusable function which performs the contraction corresponding\n to this tree, cached.\n\n Parameters\n ----------\n tree : ContractionTree\n The contraction tree.\n order : str or callable, optional\n Supplied to :meth:`ContractionTree.traverse`, the order in which\n to perform the pairwise contractions given by the tree.\n prefer_einsum : bool, optional\n Prefer to use ``einsum`` for pairwise contractions, even if\n ``tensordot`` can perform the contraction.\n strip_exponent : bool, optional\n If ``True``, the function will eagerly strip the exponent (in\n log10) from intermediate tensors to control numerical problems from\n leaving the range of the datatype. This method then returns the\n scaled 'mantissa' output array and the exponent separately.\n check_zero : bool, optional\n If ``True``, when ``strip_exponent=True``, explicitly check for\n zero-valued intermediates that would otherwise produce ``nan``,\n instead terminating early if encountered and returning\n ``(0.0, 0.0)``.\n implementation : str or tuple[callable, callable], optional\n What library to use to actually perform the contractions. Options\n are:\n\n - None: let cotengra choose.\n - \"autoray\": dispatch with autoray, using the ``tensordot`` and\n ``einsum`` implementation of the backend.\n - \"cotengra\": use the ``tensordot`` and ``einsum`` implementation\n of cotengra, which is based on batch matrix multiplication. This\n is faster for some backends like numpy, and also enables\n libraries which don't yet provide ``tensordot`` and ``einsum`` to\n be used.\n - \"cuquantum\": use the cuquantum library to perform the whole\n contraction (not just individual contractions).\n - tuple[callable, callable]: manually supply the ``tensordot`` and\n ``einsum`` implementations to use.\n\n autojit : bool, optional\n If ``True``, use :func:`autoray.autojit` to compile the contraction\n function.\n progbar : bool, optional\n Whether to show progress through the contraction by default.\n\n Returns\n -------\n fn : callable\n The contraction function, with signature ``fn(*arrays)``.\n \"\"\"\n key = (\n autojit,\n order,\n prefer_einsum,\n strip_exponent,\n check_zero,\n implementation,\n progbar,\n )\n try:\n fn = self.contraction_cores[key]\n except KeyError:\n fn = self.contraction_cores[key] = make_contractor(\n tree=self,\n order=order,\n prefer_einsum=prefer_einsum,\n strip_exponent=strip_exponent,\n check_zero=check_zero,\n implementation=implementation,\n autojit=autojit,\n progbar=progbar,\n )\n\n return fn\n\n def contract_core(\n self,\n arrays,\n order=None,\n prefer_einsum=False,\n strip_exponent=False,\n check_zero=False,\n backend=None,\n implementation=None,\n autojit=\"auto\",\n progbar=False,\n ):\n \"\"\"Contract ``arrays`` with this tree. The order of the axes and\n output is assumed to be that of ``tree.inputs`` and ``tree.output``,\n but with sliced indices removed. This functon contracts the core tree\n and thus if indices have been sliced the arrays supplied need to be\n sliced as well.\n\n Parameters\n ----------\n arrays : sequence of array\n The arrays to contract.\n order : str or callable, optional\n Supplied to :meth:`ContractionTree.traverse`.\n prefer_einsum : bool, optional\n Prefer to use ``einsum`` for pairwise contractions, even if\n ``tensordot`` can perform the contraction.\n backend : str, optional\n What library to use for ``einsum`` and ``transpose``, will be\n automatically inferred from the arrays if not given.\n autojit : \"auto\" or bool, optional\n Whether to use ``autoray.autojit`` to jit compile the expression.\n If \"auto\", then let ``cotengra`` choose.\n progbar : bool, optional\n Show progress through the contraction.\n \"\"\"\n if autojit == \"auto\":\n # choose for the user\n autojit = backend == \"jax\"\n\n fn = self.get_contractor(\n order=order,\n prefer_einsum=prefer_einsum,\n strip_exponent=strip_exponent is not False,\n implementation=implementation,\n autojit=autojit,\n check_zero=check_zero,\n progbar=progbar,\n )\n return fn(*arrays, backend=backend)\n\n def slice_key(self, i, strides=None):\n \"\"\"Get the combination of sliced index values for overall slice ``i``.\n\n Parameters\n ----------\n i : int\n The overall slice index.\n\n Returns\n -------\n key : dict[str, int]\n The value each sliced index takes for slice ``i``.\n \"\"\"\n if strides is None:\n strides = get_slice_strides(self.sliced_inds)\n\n key = {}\n for (ind, info), stride in zip(self.sliced_inds.items(), strides):\n if info.project is None:\n key[ind] = i // stride\n i %= stride\n else:\n # size is 1 and i doesn't change\n key[ind] = info.project\n\n return key\n\n def slice_arrays(self, arrays, i):\n \"\"\"Take ``arrays`` and slice the relevant inputs according to\n ``tree.sliced_inds`` and the dynary representation of ``i``.\n \"\"\"\n temp_arrays = list(arrays)\n\n # e.g. {'a': 2, 'd': 7, 'z': 0}\n locations = self.slice_key(i)\n\n for c in self.sliced_inputs:\n # the indexing object, e.g. [:, :, 7, :, 2, :, :, 0]\n selector = tuple(\n locations.get(ix, slice(None)) for ix in self.inputs[c]\n )\n # re-insert the sliced array\n temp_arrays[c] = temp_arrays[c][selector]\n\n return temp_arrays\n\n def contract_slice(self, arrays, i, **kwargs):\n \"\"\"Get slices ``i`` of ``arrays`` and then contract them.\"\"\"\n return self.contract_core(self.slice_arrays(arrays, i), **kwargs)\n\n def gather_slices(self, slices, backend=None, progbar=False):\n \"\"\"Gather all the output contracted slices into a single full result.\n If none of the sliced indices appear in the output, then this is a\n simple sum - otherwise the slices need to be partially summed and\n partially stacked.\n \"\"\"\n if progbar:\n import tqdm\n\n slices = tqdm.tqdm(slices, total=self.multiplicity)\n\n output_pos = {\n ix: i for i, ix in enumerate(self.output) if ix in self.sliced_inds\n }\n\n add_maybe_exponent_stripped = AdderWithMaybeExponentStripped()\n\n if not output_pos:\n # we can just sum everything\n return functools.reduce(add_maybe_exponent_stripped, slices)\n\n # first we sum over non-output sliced indices\n chunks = {}\n for i, s in enumerate(slices):\n key_slice = self.slice_key(i)\n key = tuple(key_slice[ix] for ix in output_pos)\n try:\n chunks[key] = add_maybe_exponent_stripped(chunks[key], s)\n except KeyError:\n chunks[key] = s\n\n if isinstance(next(iter(chunks.values())), tuple):\n # have stripped exponents, need to scale to largest\n emax = max(v[1] for v in chunks.values())\n chunks = {\n k: mi * 10 ** (ei - emax) for k, (mi, ei) in chunks.items()\n }\n else:\n emax = None\n\n # then we stack these summed chunks over output sliced indices\n def recursively_stack_chunks(loc, remaining):\n if not remaining:\n return chunks[loc]\n arrays = [\n recursively_stack_chunks(loc + (d,), remaining[1:])\n for d in self.sliced_inds[remaining[0]].sliced_range\n ]\n axes = output_pos[remaining[0]] - len(loc)\n return do(\"stack\", arrays, axes, like=backend)\n\n result = recursively_stack_chunks((), tuple(output_pos))\n\n if emax is not None:\n # strip_exponent was True, return the exponent separately\n return result, emax\n\n return result\n\n def gen_output_chunks(\n self, arrays, with_key=False, progbar=False, **contract_opts\n ):\n \"\"\"Generate each output chunk of the contraction - i.e. take care of\n summing internally sliced indices only first. This assumes that the\n ``sliced_inds`` are sorted by whether they appear in the output or not\n (the default order). Useful for performing some kind of reduction over\n the final tensor object like ``fn(x).sum()`` without constructing the\n entire thing.\n\n Parameters\n ----------\n arrays : sequence of array\n The arrays to contract.\n with_key : bool, optional\n Whether to yield the output index configuration key along with the\n chunk.\n progbar : bool, optional\n Show progress through the contraction chunks.\n\n Yields\n ------\n chunk : array\n A chunk of the contracted result.\n key : dict[str, int]\n The value each sliced output index takes for this chunk.\n \"\"\"\n # consecutive slices of size ``stepsize`` all belong to the same output\n # block because the sliced indices are sorted output first\n stepsize = prod(\n si.size for si in self.sliced_inds.values() if si.inner\n )\n\n if progbar:\n import tqdm\n\n it = tqdm.trange(self.nslices // stepsize)\n else:\n it = range(self.nslices // stepsize)\n\n for o in it:\n chunk = self.contract_slice(arrays, o * stepsize, **contract_opts)\n\n if with_key:\n output_key = {\n ix: x\n for ix, x in self.slice_key(o * stepsize).items()\n if ix in self.output\n }\n\n for j in range(1, stepsize):\n i = o * stepsize + j\n chunk = chunk + self.contract_slice(arrays, i, **contract_opts)\n\n if with_key:\n yield chunk, output_key\n else:\n yield chunk\n\n def contract(\n self,\n arrays,\n order=None,\n prefer_einsum=False,\n strip_exponent=False,\n check_zero=False,\n backend=None,\n implementation=None,\n autojit=\"auto\",\n progbar=False,\n ):\n \"\"\"Contract ``arrays`` with this tree. This function takes *unsliced*\n arrays and handles the slicing, contractions and gathering. The order\n of the axes and output is assumed to match that of ``tree.inputs`` and\n ``tree.output``.\n\n Parameters\n ----------\n arrays : sequence of array\n The arrays to contract.\n order : str or callable, optional\n Supplied to :meth:`ContractionTree.traverse`.\n prefer_einsum : bool, optional\n Prefer to use ``einsum`` for pairwise contractions, even if\n ``tensordot`` can perform the contraction.\n strip_exponent : bool, optional\n If ``True``, eagerly strip the exponent (in log10) from\n intermediate tensors to control numerical problems from leaving the\n range of the datatype. This method then returns the scaled\n 'mantissa' output array and the exponent separately.\n check_zero : bool, optional\n If ``True``, when ``strip_exponent=True``, explicitly check for\n zero-valued intermediates that would otherwise produce ``nan``,\n instead terminating early if encountered and returning\n ``(0.0, 0.0)``.\n backend : str, optional\n What library to use for ``tensordot``, ``einsum`` and\n ``transpose``, it will be automatically inferred from the input\n arrays if not given.\n autojit : bool, optional\n Whether to use the 'autojit' feature of `autoray` to compile the\n contraction expression.\n progbar : bool, optional\n Whether to show a progress bar.\n\n Returns\n -------\n output : array\n The contracted output, it will be scaled if\n ``strip_exponent==True``.\n exponent : float\n The exponent of the output in base 10, returned only if\n ``strip_exponent==True``.\n\n See Also\n --------\n contract_core, contract_slice, slice_arrays, gather_slices\n \"\"\"\n if not self.sliced_inds:\n return self.contract_core(\n arrays,\n order=order,\n prefer_einsum=prefer_einsum,\n strip_exponent=strip_exponent,\n check_zero=check_zero,\n backend=backend,\n implementation=implementation,\n autojit=autojit,\n progbar=progbar,\n )\n\n slices = (\n self.contract_slice(\n arrays,\n i,\n order=order,\n prefer_einsum=prefer_einsum,\n strip_exponent=strip_exponent,\n check_zero=check_zero,\n backend=backend,\n implementation=implementation,\n autojit=autojit,\n )\n for i in range(self.multiplicity)\n )\n\n return self.gather_slices(slices, backend=backend, progbar=progbar)\n\n def contract_mpi(self, arrays, comm=None, root=None, **kwargs):\n \"\"\"Contract the slices of this tree and sum them in parallel -\n *assuming* we are already running under MPI.\n\n Parameters\n ----------\n arrays : sequence of array\n The input (unsliced arrays)\n comm : None or mpi4py communicator\n Defaults to ``mpi4py.MPI.COMM_WORLD`` if not given.\n root : None or int, optional\n If ``root=None``, an ``Allreduce`` will be performed such that\n every process has the resulting tensor, else if an integer e.g.\n ``root=0``, the result will be exclusively gathered to that\n process using ``Reduce``, with every other process returning\n ``None``.\n kwargs\n Supplied to :meth:`~cotengra.ContractionTree.contract_slice`.\n \"\"\"\n if not set(self.sliced_inds).isdisjoint(set(self.output)):\n raise NotImplementedError(\n \"Sliced and output indices overlap - currently only a simple \"\n \"sum of result slices is supported currently.\"\n )\n\n if comm is None:\n from mpi4py import MPI\n\n comm = MPI.COMM_WORLD\n\n if self.multiplicity < comm.size:\n raise ValueError(\n f\"Need to have more slices than MPI processes, but have \"\n f\"{self.multiplicity} and {comm.size} respectively.\"\n )\n\n # round robin compute each slice, eagerly summing\n result_i = None\n for i in range(comm.rank, self.multiplicity, comm.size):\n # note: fortran ordering is needed for the MPI reduce\n x = do(\"asfortranarray\", self.contract_slice(arrays, i, **kwargs))\n if result_i is None:\n result_i = x\n else:\n result_i += x\n\n if root is None:\n # everyone gets the summed result\n result = do(\"empty_like\", result_i)\n comm.Allreduce(result_i, result)\n return result\n\n # else we only sum reduce the result to process ``root``\n if comm.rank == root:\n result = do(\"empty_like\", result_i)\n else:\n result = None\n comm.Reduce(result_i, result, root=root)\n return result\n\n def benchmark(\n self,\n dtype=\"float64\",\n max_time=60,\n min_reps=3,\n max_reps=100,\n warmup=True,\n **contract_opts,\n ):\n \"\"\"Benchmark the contraction of this tree.\n\n Parameters\n ----------\n dtype : {\"float32\", \"float64\", \"complex64\", \"complex128\"}\n The datatype to use.\n max_time : float, optional\n The maximum time to spend benchmarking in seconds.\n min_reps : int, optional\n The minimum number of repetitions to perform, regardless of time.\n max_reps : int, optional\n The maximum number of repetitions to perform, regardless of time.\n warmup : bool or int, optional\n Whether to perform a warmup run before the benchmark. If an int,\n the number of warmup runs to perform.\n contract_opts\n Supplied to :meth:`~cotengra.ContractionTree.contract_slice`.\n\n Returns\n -------\n dict\n A dictionary of benchmarking results. The keys are:\n\n - \"time_per_slice\" : float\n The average time to contract a single slice.\n - \"est_time_total\" : float\n The estimated total time to contract all slices.\n - \"est_gigaflops\" : float\n The estimated gigaflops of the contraction.\n\n See Also\n --------\n contract_slice\n \"\"\"\n import time\n\n from .utils import make_arrays_from_inputs\n\n arrays = make_arrays_from_inputs(\n self.inputs, self.size_dict, dtype=dtype\n )\n\n for i in range(int(warmup)):\n self.contract_slice(arrays, i % self.nslices, **contract_opts)\n\n t0 = time.time()\n ti = t0\n i = 0\n while (ti - t0 < max_time) or (i < min_reps):\n self.contract_slice(arrays, i % self.nslices, **contract_opts)\n ti = time.time()\n i += 1\n if i >= max_reps:\n break\n\n time_per_slice = (ti - t0) / i\n est_time_total = time_per_slice * self.nslices\n est_gigaflops = self.total_flops(dtype=dtype) / (1e9 * est_time_total)\n\n return {\n \"time_per_slice\": time_per_slice,\n \"est_time_total\": est_time_total,\n \"est_gigaflops\": est_gigaflops,\n }\n\n plot_ring = plot_tree_ring\n plot_tent = plot_tree_tent\n plot_span = plot_tree_span\n plot_flat = plot_tree_flat\n plot_circuit = plot_tree_circuit\n plot_rubberband = plot_tree_rubberband\n plot_contractions = plot_contractions\n plot_contractions_alt = plot_contractions_alt\n\n @functools.wraps(plot_hypergraph)\n def plot_hypergraph(self, **kwargs):\n hg = self.get_hypergraph(accel=False)\n hg.plot(**kwargs)\n\n def describe(self, info=\"normal\", join=\" \"):\n \"\"\"Return a string describing the contraction tree.\"\"\"\n self.contract_stats()\n if info == \"normal\":\n return join.join(\n (\n f\"log10[FLOPs]={self.total_flops(log=10):.2f}\",\n f\"log2[SIZE]={self.max_size(log=2):.2f}\",\n )\n )\n\n elif info == \"full\":\n s = [\n f\"log10[FLOPS]={self.total_flops(log=10):.2f}\",\n f\"log10[COMBO]={self.combo_cost(log=10):.2f}\",\n f\"log2[SIZE]={self.max_size(log=2):.2f}\",\n f\"log2[PEAK]={self.peak_size(log=2):.2f}\",\n ]\n if self.sliced_inds:\n s.append(f\"NSLICES={self.multiplicity:.2f}\")\n return join.join(s)\n\n elif info == \"concise\":\n s = [\n f\"F={self.total_flops(log=10):.2f}\",\n f\"C={self.combo_cost(log=10):.2f}\",\n f\"S={self.max_size(log=2):.2f}\",\n f\"P={self.peak_size(log=2):.2f}\",\n ]\n if self.sliced_inds:\n s.append(f\"$={self.multiplicity:.2f}\")\n return join.join(s)\n\n def __repr__(self):\n if self.is_complete():\n return f\"<{self.__class__.__name__}(N={self.N})>\"\n else:\n s = \"<{}(N={}, branches={}, complete={})>\"\n return s.format(\n self.__class__.__name__,\n self.N,\n len(self.children),\n self.is_complete(),\n )\n\n def __str__(self):\n if not self.is_complete():\n return self.__repr__()\n else:\n d = self.describe(\"concise\", join=\", \")\n return f\"<{self.__class__.__name__}(N={self.N}, {d})>\"\n\n\ndef _reconfigure_tree(tree, *args, **kwargs):\n return tree.subtree_reconfigure(*args, **kwargs)\n\n\ndef _slice_and_reconfigure_tree(tree, *args, **kwargs):\n return tree.slice_and_reconfigure(*args, **kwargs)\n\n\ndef _get_tree_info(tree):\n stats = tree.contract_stats()\n stats[\"sliced_ind_set\"] = frozenset(tree.sliced_inds)\n return stats\n\n\ndef _describe_tree(tree, info=\"normal\"):\n return tree.describe(info=info)\n\n\nclass ContractionTreeCompressed(ContractionTree):\n \"\"\"A contraction tree for compressed contractions. Currently the only\n difference is that this defaults to the 'surface' traversal ordering.\n \"\"\"\n\n def set_state_from(self, other):\n super().set_state_from(other)\n self.set_surface_order_from_path(other.get_ssa_path())\n\n @classmethod\n def from_path(\n cls,\n inputs,\n output,\n size_dict,\n *,\n path=None,\n ssa_path=None,\n autocomplete=\"auto\",\n check=False,\n **kwargs,\n ):\n \"\"\"Create a (completed) ``ContractionTreeCompressed`` from the usual\n inputs plus a standard contraction path or 'ssa_path' - you need to\n supply one. This also set the default 'surface' traversal ordering to\n be the initial path.\n \"\"\"\n if int(path is None) + int(ssa_path is None) != 1:\n raise ValueError(\n \"Exactly one of ``path`` or ``ssa_path`` must be supplied.\"\n )\n\n if path is not None:\n from .pathfinders.path_basic import linear_to_ssa\n\n ssa_path = linear_to_ssa(path)\n\n tree = super().from_path(\n inputs,\n output,\n size_dict,\n ssa_path=ssa_path,\n autocomplete=False,\n check=check,\n **kwargs,\n )\n\n tree.set_surface_order_from_path(ssa_path)\n\n if (len(tree.children) < tree.N - 1) and autocomplete:\n if autocomplete == \"auto\":\n # warn that we are completing\n warnings.warn(\n \"Path was not complete - contracting all remaining. \"\n \"You can silence this warning with `autocomplete=True`.\"\n \"Or produce an incomplete tree with `autocomplete=False`.\"\n )\n\n tree.autocomplete(optimize=\"greedy-compressed\")\n\n return tree\n\n def get_default_order(self):\n return \"surface_order\"\n\n def get_default_objective(self):\n if self._default_objective is None:\n self._default_objective = get_score_fn(\"peak-compressed\")\n return self._default_objective\n\n def get_default_chi(self):\n objective = self.get_default_objective()\n try:\n chi = objective.chi\n except AttributeError:\n chi = \"auto\"\n\n if chi == \"auto\":\n chi = max(self.size_dict.values()) ** 2\n\n return chi\n\n def get_default_compress_late(self):\n objective = self.get_default_objective()\n try:\n return objective.compress_late\n except AttributeError:\n return False\n\n total_flops = ContractionTree.total_flops_compressed\n total_write = ContractionTree.total_write_compressed\n combo_cost = ContractionTree.combo_cost_compressed\n total_cost = ContractionTree.total_cost_compressed\n max_size = ContractionTree.max_size_compressed\n peak_size = ContractionTree.peak_size_compressed\n contraction_cost = ContractionTree.contraction_cost_compressed\n contraction_width = ContractionTree.contraction_width_compressed\n\n total_flops_exact = ContractionTree.total_flops\n total_write_exact = ContractionTree.total_write\n combo_cost_exact = ContractionTree.combo_cost\n total_cost_exact = ContractionTree.total_cost\n max_size_exact = ContractionTree.max_size\n peak_size_exact = ContractionTree.peak_size\n\n def get_contractor(self, *_, **__):\n raise NotImplementedError(\n \"`cotengra` doesn't implement compressed contraction itself. \"\n \"If you want to use compressed contractions, you need to use \"\n \"`quimb` and the `TensorNetwork.contract_compressed` method, \"\n \"with e.g. `optimize=tree.get_path()`.\"\n )\n\n def simulated_anneal(\n self,\n minimize=None,\n tfinal=0.0001,\n tstart=0.01,\n tsteps=50,\n numiter=50,\n seed=None,\n inplace=False,\n progbar=False,\n **kwargs,\n ):\n \"\"\"Perform simulated annealing refinement of this *compressed*\n contraction tree.\n \"\"\"\n from .pathfinders.path_compressed import WindowedOptimizer\n\n if minimize is None:\n minimize = self.get_default_objective()\n\n wo = WindowedOptimizer(\n self.inputs,\n self.output,\n self.size_dict,\n minimize=minimize,\n ssa_path=self.get_ssa_path(),\n seed=seed,\n )\n\n wo.simulated_anneal(\n tfinal=tfinal,\n tstart=tstart,\n tsteps=tsteps,\n numiter=numiter,\n progbar=progbar,\n **kwargs,\n )\n ssa_path = wo.get_ssa_path()\n\n rtree = self.__class__.from_path(\n self.inputs,\n self.output,\n self.size_dict,\n ssa_path=ssa_path,\n objective=minimize,\n )\n\n if inplace:\n self.set_state_from(rtree)\n rtree = self\n\n rtree.reset_contraction_indices()\n return rtree\n\n simulated_anneal_ = functools.partialmethod(simulated_anneal, inplace=True)\n\n\nclass PartitionTreeBuilder:\n \"\"\"Function wrapper that takes a function that partitions graphs and\n uses it to build a contraction tree. ``partition_fn`` should have\n signature:\n\n def partition_fn(inputs, output, size_dict,\n weight_nodes, weight_edges, **kwargs):\n ...\n return membership\n\n Where ``weight_nodes`` and ``weight_edges`` decsribe how to weight the\n nodes and edges of the graph respectively and ``membership`` should be a\n list of integers of length ``len(inputs)`` labelling which partition\n each input node should be put it.\n \"\"\"\n\n def __init__(self, partition_fn):\n self.partition_fn = partition_fn\n\n def build_divide(\n self,\n inputs,\n output,\n size_dict,\n random_strength=0.01,\n cutoff=10,\n parts=2,\n parts_decay=0.5,\n sub_optimize=\"greedy\",\n super_optimize=\"random-greedy-128\",\n check=False,\n seed=None,\n **partition_opts,\n ):\n tree = ContractionTree(\n inputs,\n output,\n size_dict,\n track_childless=True,\n )\n\n rng = get_rng(seed)\n rand_size_dict = jitter_dict(size_dict, random_strength, rng)\n\n dynamic_imbalance = (\"imbalance\" in partition_opts) and (\n \"imbalance_decay\" in partition_opts\n )\n if dynamic_imbalance:\n imbalance = partition_opts.pop(\"imbalance\")\n imbalance_decay = partition_opts.pop(\"imbalance_decay\")\n else:\n imbalance = imbalance_decay = None\n\n dynamic_fix = partition_opts.get(\"fix_output_nodes\", None) == \"auto\"\n\n while tree.childless:\n top_node = next(iter(tree.childless))\n subgraph = tree.get_subgraph(top_node)\n subsize = len(subgraph)\n\n # skip straight to better method\n if subsize <= cutoff:\n tree.contract_nodes(\n [tree.input_to_node(x) for x in subgraph],\n grandparent=top_node,\n optimize=sub_optimize,\n check=check,\n )\n continue\n\n # relative subgraph size\n s = subsize / tree.N\n\n # let the target number of communities depend on subgraph size\n parts_s = max(int(s**parts_decay * parts), 2)\n\n # let the imbalance either rise or fall\n if dynamic_imbalance:\n if imbalance_decay >= 0:\n imbalance_s = s**imbalance_decay * imbalance\n else:\n imbalance_s = 1 - s**-imbalance_decay * (1 - imbalance)\n partition_opts[\"imbalance\"] = imbalance_s\n\n if dynamic_fix:\n # for the top level subtree (s==1.0) we partition the outputs\n # nodes first into their own bi-partition\n parts_s = 2\n partition_opts[\"fix_output_nodes\"] = s == 1.0\n\n # partition! get community membership list e.g.\n # [0, 0, 1, 0, 1, 0, 0, 2, 2, ...]\n inputs = tuple(map(tuple, tree.node_to_terms(top_node)))\n output = tuple(tree.get_legs(top_node))\n membership = self.partition_fn(\n inputs,\n output,\n rand_size_dict,\n parts=parts_s,\n seed=rng,\n **partition_opts,\n )\n\n # divide subgraph up e.g. if we enumerate the subgraph index sets\n # (0, 1, 2, 3, 4, 5, 6, 7, 8, ...) ->\n # ({0, 1, 3, 5, 6}, {2, 4}, {7, 8})\n partitions = separate(subgraph, membership)\n\n if len(partitions) == 1:\n # no communities found - contract all remaining leaves\n tree.contract_nodes(\n tuple(map(tree.input_to_node, subgraph)),\n grandparent=top_node,\n optimize=sub_optimize,\n check=check,\n )\n continue\n\n tree.contract_nodes(\n partitions,\n grandparent=top_node,\n optimize=super_optimize,\n check=check,\n )\n\n if check:\n assert tree.is_complete()\n\n return tree\n\n def build_agglom(\n self,\n inputs,\n output,\n size_dict,\n random_strength=0.01,\n groupsize=4,\n sub_optimize=\"greedy\",\n check=False,\n seed=None,\n **partition_opts,\n ):\n tree = ContractionTree(\n inputs,\n output,\n size_dict,\n track_childless=True,\n )\n rand_size_dict = jitter_dict(size_dict, random_strength, seed)\n leaves = tuple(tree.gen_leaves())\n output = tuple(tree.output)\n\n while len(leaves) > groupsize:\n # choose number of partitions so that each\n # has approximately ``groupsize`` nodes\n parts = max(2, len(leaves) // groupsize)\n\n # partition! get community membership list\n inputs = [tuple(tree.get_legs(node)) for node in leaves]\n membership = self.partition_fn(\n inputs,\n output,\n rand_size_dict,\n parts=parts,\n **partition_opts,\n )\n\n # group leaves according to partition label\n partitions = separate(leaves, membership)\n\n if len(partitions) == 1:\n # only found one group, move to final contraction\n break\n\n # contract each group into a new leaf\n leaves = [\n tree.contract_nodes(\n partition,\n check=check,\n optimize=sub_optimize,\n )\n for partition in partitions\n ]\n\n if len(leaves) > 1:\n # contract any remaining leaves together\n tree.contract_nodes(\n leaves,\n check=check,\n optimize=sub_optimize,\n grandparent=tree.root,\n )\n\n if check:\n assert tree.is_complete()\n\n return tree\n\n def trial_fn(self, inputs, output, size_dict, **partition_opts):\n return self.build_divide(inputs, output, size_dict, **partition_opts)\n\n def trial_fn_agglom(self, inputs, output, size_dict, **partition_opts):\n return self.build_agglom(inputs, output, size_dict, **partition_opts)\n\n\ndef jitter(x, strength, rng):\n return x * (1 + strength * rng.expovariate(1.0))\n\n\ndef jitter_dict(d, strength, seed=None):\n rng = get_rng(seed)\n return {k: jitter(v, strength, rng) for k, v in d.items()}\n\n\ndef separate(xs, blocks):\n \"\"\"Partition ``xs`` into ``n`` different list based on the corresponding\n labels in ``blocks``.\n \"\"\"\n sorter = collections.defaultdict(list)\n for x, b in zip(xs, blocks):\n sorter[b].append(x)\n x_b = list(sorter.items())\n x_b.sort()\n return [x[1] for x in x_b]\n"
},
{
"path": "cotengra/core_multi.py",
"content": "import math\n\nfrom .core import ContractionTree, cached_node_property\n\n\nclass ContractionTreeMulti(ContractionTree):\n def __init__(\n self,\n inputs,\n output,\n size_dict,\n sliced_inds,\n objective,\n track_cache=False,\n ):\n super().__init__(inputs, output, size_dict, objective=objective)\n self.sliced_inds = {ix: None for ix in sliced_inds}\n\n self._track_cache = track_cache\n if track_cache:\n self._cache_est = 0\n\n def set_state_from(self, other):\n super().set_state_from(other)\n self._track_cache = other._track_cache\n if other._track_cache:\n self._cache_est = other._cache_est\n\n def _remove_node(self, node):\n if self._track_cache:\n self._cache_est -= self.get_cache_contrib(node)\n super()._remove_node(node)\n\n def _update_tracked(self, node):\n if self._track_cache:\n self._cache_est += self.get_cache_contrib(node)\n super()._update_tracked(node)\n\n @cached_node_property(\"node_var_inds\")\n def get_node_var_inds(self, node):\n \"\"\"Get the set of variable indices that a node depends on.\"\"\"\n if self.is_leaf(node):\n i = self.node_to_input(node)\n term = self.inputs[i]\n return {ix: None for ix in term if ix in self.sliced_inds}\n\n try:\n l, r = self.children[node]\n return self.get_node_var_inds(l) | self.get_node_var_inds(r)\n except KeyError:\n return {\n ix: None\n for term in self.node_to_terms(node)\n for ix in term\n if ix in self.sliced_inds\n }\n\n @cached_node_property(\"node_is_bright\")\n def get_node_is_bright(self, node):\n \"\"\"Get whether a node is 'bright', i.e. contains a different set of\n variable indices to either of its children, if a node is not bright\n then its children never have to be stored in the cache.\n \"\"\"\n if self.is_leaf(node):\n i = self.node_to_input(node)\n term = self.inputs[i]\n return any(ix in self.sliced_inds for ix in term)\n\n l, r = self.children[node]\n return (self.get_node_var_inds(node) != self.get_node_var_inds(l)) or (\n self.get_node_var_inds(node) != self.get_node_var_inds(r)\n )\n\n @cached_node_property(\"node_mult\")\n def get_node_mult(self, node):\n \"\"\"Get the estimated 'multiplicity' of a node, i.e. the number of times\n it will have to be recomputed for different index configurations.\n \"\"\"\n return self.get_default_objective().estimate_node_mult(self, node)\n\n def get_node_cache_mult(self, node, sliced_ind_ordering):\n \"\"\"Get the estimated 'cache multiplicity' of a node, i.e. the total\n number of versions with different index configurations that must be\n stored simultaneously in the cache.\n \"\"\"\n return self.get_default_objective().estimate_node_cache_mult(\n self, node, sliced_ind_ordering\n )\n\n # @cached_node_property(\"multi_flops\")\n def get_flops(self, node):\n \"\"\"The the estimated total cost of computing a node for all index\n configurations.\n \"\"\"\n return super().get_flops(node) * self.get_node_mult(node)\n\n @cached_node_property(\"cache_contrib\")\n def get_cache_contrib(self, node):\n l, r = self.children[node]\n lr_peak = 0\n if self.get_node_is_bright(l):\n lr_peak += self.get_size(l)\n if self.get_node_is_bright(r):\n lr_peak += self.get_size(r) * self.get_node_mult(r)\n\n rl_peak = 0\n if self.get_node_is_bright(r):\n rl_peak += self.get_size(r)\n if self.get_node_is_bright(l):\n rl_peak += self.get_size(l) * self.get_node_mult(l)\n\n if lr_peak < rl_peak:\n return lr_peak\n else:\n self.children[node] = (r, l)\n return rl_peak\n\n def peak_size(self, log=None):\n if not self._track_cache:\n self._cache_est = 0\n for (\n node,\n _,\n _,\n ) in self.traverse():\n self._cache_est += self.get_cache_contrib(node)\n self._track_cache = True\n peak = self._cache_est\n\n if log is not None:\n peak = math.log(peak, log)\n\n return peak\n\n def reorder_contractions_for_peak_est(self):\n \"\"\"Reorder the contractions to try and reduce the peak memory usage.\"\"\"\n swapped = False\n\n for p, l, r in self.descend():\n lr_peak = 0\n if self.get_node_is_bright(l):\n lr_peak += self.get_size(l)\n if self.get_node_is_bright(r):\n lr_peak += self.get_size(r) * self.get_node_mult(r)\n\n rl_peak = 0\n if self.get_node_is_bright(r):\n rl_peak += self.get_size(r)\n if self.get_node_is_bright(l):\n rl_peak += self.get_size(l) * self.get_node_mult(l)\n\n if rl_peak < lr_peak:\n self.children[p] = (r, l)\n swapped = True\n\n return swapped\n\n def reorder_sliced_inds(self):\n \"\"\" \"\"\"\n sliced_ind_ordering = dict()\n\n for node, _, _ in self.traverse():\n sliced_ind_ordering.update(self.get_node_var_inds(node))\n\n self.sliced_inds = {ix: None for ix in sliced_ind_ordering}\n\n def exact_multi_stats(self, configs):\n # ragged list of lists (configs and contractions)\n cons = []\n\n # build this for efficiency\n plr = tuple(self.traverse())\n\n def to_key(node, config):\n subconfig = tuple(\n map(config.__getitem__, self.get_node_var_inds(node))\n )\n return hash((node, subconfig))\n\n # iterate forward, recording only when we first need to produce a 'parent'\n seen = set()\n for config in configs:\n cons_i = []\n for p, l, r in plr:\n pkey = to_key(p, config)\n first = pkey not in seen\n if first:\n seen.add(pkey)\n cons_i.append(\n {\n \"p\": p,\n \"l\": l,\n \"r\": r,\n \"pkey\": pkey,\n \"lkey\": to_key(l, config),\n \"rkey\": to_key(r, config),\n }\n )\n cons.append(cons_i)\n del seen\n\n # iterate backward, checking the last\n # time a 'child' is seen -> can delete\n deleted = set()\n for cons_i in reversed(cons):\n for con in cons_i:\n rkey = con[\"rkey\"]\n rdel = rkey not in deleted\n if rdel:\n deleted.add(rkey)\n con[\"rdel\"] = rdel\n\n lkey = con[\"lkey\"]\n ldel = lkey not in deleted\n if ldel:\n deleted.add(lkey)\n con[\"ldel\"] = ldel\n del deleted\n\n # iterate forward again if we want to compute flops and memory usage:\n # not needed if we already know these & just want to contract\n flops = 0\n mems = []\n mem_current = 0\n mem_peak = 0\n mem_write = 0\n\n for cons_i in cons:\n for con in cons_i:\n p = con[\"p\"]\n flops += super().get_flops(p)\n psize = self.get_size(p)\n mem_current += psize\n mem_write += psize\n\n mems.append(mem_current)\n mem_peak = max(mem_peak, mem_current)\n\n l, r = con[\"l\"], con[\"r\"]\n if con[\"ldel\"] and not self.is_leaf(l):\n mem_current -= self.get_size(l)\n if con[\"rdel\"] and not self.is_leaf(r):\n mem_current -= self.get_size(r)\n\n # final output of each config is always deletable\n mem_current -= self.get_size(p)\n\n return {\n \"flops\": flops,\n \"write\": mem_write,\n \"size\": self.max_size(),\n \"peak\": mem_peak,\n }\n"
},
{
"path": "cotengra/experimental/__init__.py",
"content": "\"\"\"Potentially useful but experimental (untested) features,\"\"\"\n"
},
{
"path": "cotengra/experimental/hyper_de.py",
"content": "\"\"\"Hyper optimization using a pure Python differential evolution strategy.\"\"\"\n\nfrom ..utils import get_rng\nfrom ._param_mapping import (\n LCBOptimizer,\n build_params,\n convert_raw,\n num_params,\n)\nfrom .hyper import HyperOptLib, register_hyper_optlib\n\n\nclass HyperDESampler:\n \"\"\"A lightweight differential evolution optimizer operating in raw\n ``[-1, 1]`` parameter space.\n\n Each generation maintains a population of candidate vectors. New trial\n vectors are created using ``DE/rand/1/bin`` mutation and binomial\n crossover, then kept only if they improve on their parent.\n\n Parameters\n ----------\n space : dict[str, dict]\n The search space for a single contraction method.\n seed : None or int, optional\n Random seed.\n population_size : int or \"auto\", optional\n The population size. When ``\"auto\"`` it is chosen based on the mapped\n parameter dimension.\n mutation : float, optional\n The differential weight (F) applied to the difference vector.\n crossover : float, optional\n The crossover probability (CR) for binomial crossover.\n mutation_decay : float, optional\n Multiplicative decay applied to ``mutation`` after each completed\n generation.\n mutation_min : float, optional\n Lower bound for ``mutation``.\n mutation_max : float, optional\n Upper bound for ``mutation``.\n exponential_param_power : float, optional\n Passed through to the shared parameter mapping for ``FLOAT_EXP``\n parameters.\n \"\"\"\n\n def __init__(\n self,\n space,\n seed=None,\n population_size=\"auto\",\n mutation=0.8,\n crossover=0.7,\n mutation_decay=1.0,\n mutation_min=0.1,\n mutation_max=1.5,\n exponential_param_power=None,\n ):\n self.rng = get_rng(seed)\n self.params = build_params(\n space, exponential_param_power=exponential_param_power\n )\n self.ndim = num_params(self.params)\n\n if population_size == \"auto\":\n population_size = max(8, 5 * self.ndim)\n self.population_size = population_size\n self.mutation = mutation\n self.crossover = crossover\n self.mutation_decay = mutation_decay\n self.mutation_min = mutation_min\n self.mutation_max = mutation_max\n\n # initialize population uniformly in [-1, 1]\n self._population = [\n tuple(self.rng.uniform(-1.0, 1.0) for _ in range(self.ndim))\n for _ in range(self.population_size)\n ]\n self._scores = [float(\"inf\")] * self.population_size\n\n self._trial_counter = 0\n self._target_index = 0\n self._generation = None\n self._trial_map = {}\n\n def _mutate(self, target_idx):\n \"\"\"Create a trial vector via DE/rand/1/bin.\"\"\"\n # pick three distinct indices, all different from target\n indices = list(range(self.population_size))\n indices.remove(target_idx)\n r0, r1, r2 = self.rng.sample(indices, 3)\n\n x_r0 = self._population[r0]\n x_r1 = self._population[r1]\n x_r2 = self._population[r2]\n\n # mutation: v = x_r0 + F * (x_r1 - x_r2)\n v = []\n for d in range(self.ndim):\n vi = x_r0[d] + self.mutation * (x_r1[d] - x_r2[d])\n v.append(min(max(vi, -1.0), 1.0))\n\n # binomial crossover\n x_target = self._population[target_idx]\n j_rand = self.rng.randrange(self.ndim)\n trial = []\n for d in range(self.ndim):\n if self.rng.random() < self.crossover or d == j_rand:\n trial.append(v[d])\n else:\n trial.append(x_target[d])\n\n return tuple(trial)\n\n def _sample_generation(self):\n \"\"\"Prepare trial vectors for all population members.\"\"\"\n self._generation = {\n \"trials\": [],\n \"trial_numbers\": [],\n \"target_indices\": [],\n \"scores\": {},\n \"next_index\": 0,\n }\n for i in range(self.population_size):\n self._extend_generation(i)\n\n def _extend_generation(self, target_idx=None):\n \"\"\"Append one more trial to the current generation.\"\"\"\n if target_idx is None:\n # wrap around if we need more trials than population\n target_idx = len(self._generation[\"trials\"]) % self.population_size\n\n trial_number = self._trial_counter\n self._trial_counter += 1\n\n trial_vec = self._mutate(target_idx)\n slot = len(self._generation[\"trials\"])\n\n self._generation[\"trials\"].append(trial_vec)\n self._generation[\"trial_numbers\"].append(trial_number)\n self._generation[\"target_indices\"].append(target_idx)\n self._trial_map[trial_number] = slot\n\n def ask(self):\n \"\"\"Return the next candidate from the current generation.\n\n If all prepared candidates have been issued, grow the generation\n by one more sample.\n \"\"\"\n if self._generation is None:\n self._sample_generation()\n\n if self._generation[\"next_index\"] >= len(self._generation[\"trials\"]):\n self._extend_generation()\n\n i = self._generation[\"next_index\"]\n self._generation[\"next_index\"] += 1\n\n trial_number = self._generation[\"trial_numbers\"][i]\n x = self._generation[\"trials\"][i]\n\n return trial_number, convert_raw(self.params, x)\n\n def tell(self, trial_number, score):\n \"\"\"Record a completed trial and perform selection if the\n generation is complete.\n\n For each trial vector, if it scores better than (or equal to) its\n target parent, it replaces the parent in the population.\n \"\"\"\n slot = self._trial_map.pop(trial_number)\n self._generation[\"scores\"][slot] = score\n\n if len(self._generation[\"scores\"]) != self._generation[\"next_index\"]:\n return\n\n # selection: compare each trial against its target\n # (only iterate over issued slots, not all pre-sampled ones)\n for slot_i in range(self._generation[\"next_index\"]):\n target_idx = self._generation[\"target_indices\"][slot_i]\n trial_score = self._generation[\"scores\"][slot_i]\n if trial_score <= self._scores[target_idx]:\n self._population[target_idx] = self._generation[\"trials\"][\n slot_i\n ]\n self._scores[target_idx] = trial_score\n\n # decay mutation factor\n self.mutation *= self.mutation_decay\n self.mutation = min(\n max(self.mutation, self.mutation_min), self.mutation_max\n )\n\n self._generation = None\n\n\nclass DEOptLib(HyperOptLib):\n \"\"\"Hyper-optimization using differential evolution.\"\"\"\n\n def setup(\n self,\n methods,\n space,\n optimizer=None,\n population_size=\"auto\",\n mutation=0.8,\n crossover=0.7,\n mutation_decay=1.0,\n mutation_min=0.1,\n mutation_max=1.5,\n method_exploration=1.0,\n method_temperature=1.0,\n exponential_param_power=None,\n seed=None,\n **kwargs,\n ):\n \"\"\"Initialize DE optimizers for each contraction method.\n\n Parameters\n ----------\n methods : list[str]\n The contraction methods to optimize over.\n space : dict[str, dict[str, dict]]\n The per-method hyperparameter search space.\n optimizer : HyperOptimizer, optional\n The parent optimizer. Used to size the initial population\n large enough for parallel pre-dispatch.\n population_size : int or \"auto\", optional\n The population size for each method-specific DE sampler.\n mutation : float, optional\n Differential weight (F).\n crossover : float, optional\n Crossover probability (CR).\n mutation_decay, mutation_min, mutation_max : float, optional\n Parameters controlling mutation scale over generations.\n method_exploration : float, optional\n Exploration strength for the LCB-based method chooser.\n method_temperature : float, optional\n Noise temperature for the LCB-based method chooser.\n exponential_param_power : float, optional\n Passed to the shared parameter mapping for ``FLOAT_EXP``.\n seed : None or int, optional\n Random seed.\n \"\"\"\n if population_size == \"auto\":\n max_ndim = max(\n num_params(\n build_params(\n space[m],\n exponential_param_power=exponential_param_power,\n )\n )\n for m in methods\n )\n population_size = max(\n 8,\n max(1, getattr(optimizer, \"pre_dispatch\", 1)),\n 5 * max_ndim,\n )\n\n self._method_chooser = LCBOptimizer(\n options=methods,\n exploration=method_exploration,\n temperature=method_temperature,\n seed=seed,\n )\n self._optimizers = {\n method: HyperDESampler(\n space[method],\n seed=seed,\n population_size=population_size,\n mutation=mutation,\n crossover=crossover,\n mutation_decay=mutation_decay,\n mutation_min=mutation_min,\n mutation_max=mutation_max,\n exponential_param_power=exponential_param_power,\n )\n for method in methods\n }\n\n def get_setting(self):\n \"\"\"Choose a contraction method, then request its next setting.\"\"\"\n method = self._method_chooser.ask()\n params_token, params = self._optimizers[method].ask()\n\n return {\n \"method\": method,\n \"params_token\": params_token,\n \"params\": params,\n }\n\n def report_result(self, setting, trial, score):\n \"\"\"Report a completed trial back to the method chooser and DE.\"\"\"\n self._method_chooser.tell(setting[\"method\"], score)\n self._optimizers[setting[\"method\"]].tell(\n setting[\"params_token\"], score\n )\n\n\nregister_hyper_optlib(\"de\", DEOptLib)\nregister_hyper_optlib(\"diffev\", DEOptLib)\n"
},
{
"path": "cotengra/experimental/hyper_pe.py",
"content": "\"\"\"Hyper optimization using parallel evolution with ranked sigma assignment.\"\"\"\n\nimport math\n\nfrom ..utils import get_rng\nfrom ._param_mapping import (\n LCBOptimizer,\n build_params,\n convert_raw,\n num_params,\n)\nfrom .hyper import HyperOptLib, register_hyper_optlib\n\n\nclass HyperPESampler:\n \"\"\"A parallel evolution optimizer operating in raw ``[-1, 1]`` space.\n\n Multiple workers each maintain their own solution. Perturbation scales\n (sigmas) are distributed across an evenly spaced range and reassigned\n by rank after each generation: the best-scoring worker gets the lowest\n sigma (exploit) and the worst gets the highest (explore).\n\n Parameters\n ----------\n space : dict[str, dict]\n The search space for a single contraction method.\n seed : None or int, optional\n Random seed.\n population_size : int or \"auto\", optional\n The number of parallel workers. When ``\"auto\"`` it is chosen\n based on the mapped parameter dimension.\n sigma_min : float, optional\n The smallest perturbation scale (assigned to the best worker).\n sigma_max : float, optional\n The largest perturbation scale (assigned to the worst worker).\n elite_migrate_prob : float, optional\n Probability each generation of copying the best worker's\n solution to the worst worker's slot.\n differential_prob : float, optional\n Per-sample probability of using a differential perturbation\n (``x_best - x_rand``) instead of Gaussian noise.\n patience : int or None, optional\n If a worker has not improved for this many generations,\n re-randomize its solution. ``None`` or ``0`` disables.\n exponential_param_power : float, optional\n Passed through to the shared parameter mapping for ``FLOAT_EXP``\n parameters.\n \"\"\"\n\n def __init__(\n self,\n space,\n seed=None,\n population_size=8,\n sigma_min=0.01,\n sigma_max=0.5,\n elite_migrate_prob=0.0,\n differential_prob=0.0,\n patience=None,\n exponential_param_power=None,\n ):\n self.rng = get_rng(seed)\n self.params = build_params(\n space, exponential_param_power=exponential_param_power\n )\n self.ndim = num_params(self.params)\n\n if population_size == \"auto\":\n population_size = max(8, 4 * self.ndim)\n self.population_size = population_size\n self.sigma_min = sigma_min\n self.sigma_max = sigma_max\n self.elite_migrate_prob = elite_migrate_prob\n self.differential_prob = differential_prob\n self.patience = patience\n\n # initialize each worker's solution uniformly in [-1, 1]\n self._solutions = [\n tuple(self.rng.uniform(-1.0, 1.0) for _ in range(self.ndim))\n for _ in range(self.population_size)\n ]\n self._scores = [float(\"inf\")] * self.population_size\n self._stagnation = [0] * self.population_size\n\n # evenly spaced sigmas, assigned by rank (best -> lowest)\n self._sigmas = self._make_sigmas()\n\n self._trial_counter = 0\n self._generation = None\n self._trial_map = {}\n\n def _make_sigmas(self):\n \"\"\"Create geometrically spaced sigmas from sigma_min to sigma_max.\"\"\"\n n = self.population_size\n if n == 1:\n return [math.sqrt(self.sigma_min * self.sigma_max)]\n log_min = math.log(self.sigma_min)\n log_max = math.log(self.sigma_max)\n return [\n math.exp(log_min + i * (log_max - log_min) / (n - 1))\n for i in range(n)\n ]\n\n def _sample_candidate(self, worker_idx, noise=None):\n \"\"\"Perturb a worker's current solution with its assigned sigma.\"\"\"\n sigma = self._sigmas[worker_idx]\n sol = self._solutions[worker_idx]\n\n if noise is None:\n # possibly use differential perturbation\n if (\n self.differential_prob > 0.0\n and self.population_size >= 3\n and self.rng.random() < self.differential_prob\n ):\n best_idx = min(\n range(self.population_size),\n key=lambda i: self._scores[i],\n )\n others = [\n j\n for j in range(self.population_size)\n if j != worker_idx and j != best_idx\n ]\n rand_idx = self.rng.choice(others)\n noise = [\n self._solutions[best_idx][d] - self._solutions[rand_idx][d]\n for d in range(self.ndim)\n ]\n else:\n noise = [self.rng.gauss(0.0, 1.0) for _ in range(self.ndim)]\n\n x = []\n for si, ni in zip(sol, noise):\n xi = si + sigma * ni\n x.append(min(max(xi, -1.0), 1.0))\n return tuple(x)\n\n def _sample_generation(self):\n \"\"\"Start a new generation with one trial per worker.\"\"\"\n self._generation = {\n \"worker_indices\": [],\n \"xs\": [],\n \"trial_numbers\": [],\n \"scores\": {},\n \"next_index\": 0,\n }\n for i in range(self.population_size):\n self._extend_generation(i)\n\n def _extend_generation(self, worker_idx=None, noise=None):\n \"\"\"Append one more trial to the current generation.\"\"\"\n if worker_idx is None:\n worker_idx = len(self._generation[\"xs\"]) % self.population_size\n\n trial_number = self._trial_counter\n self._trial_counter += 1\n\n slot = len(self._generation[\"xs\"])\n self._generation[\"worker_indices\"].append(worker_idx)\n self._generation[\"xs\"].append(\n self._sample_candidate(worker_idx, noise=noise)\n )\n self._generation[\"trial_numbers\"].append(trial_number)\n self._trial_map[trial_number] = slot\n\n def ask(self):\n \"\"\"Return the next candidate from the current generation.\n\n If all prepared candidates have been issued, grow the generation\n by one more sample.\n \"\"\"\n if self._generation is None:\n self._sample_generation()\n\n if self._generation[\"next_index\"] >= len(self._generation[\"xs\"]):\n self._extend_generation()\n\n i = self._generation[\"next_index\"]\n self._generation[\"next_index\"] += 1\n\n trial_number = self._generation[\"trial_numbers\"][i]\n x = self._generation[\"xs\"][i]\n\n return trial_number, convert_raw(self.params, x)\n\n def tell(self, trial_number, score):\n \"\"\"Record a completed trial and update workers if the generation\n is complete.\n\n For each trial, if it scores better than (or equal to) its\n worker's current best, the worker adopts the new solution.\n Then sigmas are reassigned by rank: best worker gets lowest\n sigma, worst gets highest.\n \"\"\"\n slot = self._trial_map.pop(trial_number)\n self._generation[\"scores\"][slot] = score\n\n if len(self._generation[\"scores\"]) != self._generation[\"next_index\"]:\n return\n\n # greedy update: adopt better solutions and track stagnation\n improved = set()\n for slot_i in range(self._generation[\"next_index\"]):\n worker_idx = self._generation[\"worker_indices\"][slot_i]\n trial_score = self._generation[\"scores\"][slot_i]\n if trial_score <= self._scores[worker_idx]:\n self._solutions[worker_idx] = self._generation[\"xs\"][slot_i]\n self._scores[worker_idx] = trial_score\n improved.add(worker_idx)\n\n for i in range(self.population_size):\n if i in improved:\n self._stagnation[i] = 0\n else:\n self._stagnation[i] += 1\n\n # stagnation restart\n if self.patience:\n for i in range(self.population_size):\n if self._stagnation[i] >= self.patience:\n self._solutions[i] = tuple(\n self.rng.uniform(-1.0, 1.0) for _ in range(self.ndim)\n )\n self._scores[i] = float(\"inf\")\n self._stagnation[i] = 0\n\n # rank workers by score (best to worst) and reassign sigmas\n ranking = sorted(\n range(self.population_size),\n key=lambda i: self._scores[i],\n )\n\n # elite migration\n if self.elite_migrate_prob > 0.0:\n if self.rng.random() < self.elite_migrate_prob:\n worst = ranking[-1]\n best = ranking[0]\n self._solutions[worst] = self._solutions[best]\n\n new_sigmas = [0.0] * self.population_size\n base_sigmas = self._make_sigmas()\n for rank, worker_idx in enumerate(ranking):\n new_sigmas[worker_idx] = base_sigmas[rank]\n self._sigmas = new_sigmas\n\n self._generation = None\n\n\nclass PEOptLib(HyperOptLib):\n \"\"\"Hyper-optimization using parallel evolution with ranked sigmas.\"\"\"\n\n def setup(\n self,\n methods,\n space,\n optimizer=None,\n population_size=\"auto\",\n sigma_min=0.01,\n sigma_max=0.5,\n elite_migrate_prob=0.5,\n differential_prob=0.5,\n patience=8,\n method_exploration=1.0,\n method_temperature=1.0,\n exponential_param_power=None,\n seed=None,\n **kwargs,\n ):\n \"\"\"Initialize PE optimizers for each contraction method.\n\n Parameters\n ----------\n methods : list[str]\n The contraction methods to optimize over.\n space : dict[str, dict[str, dict]]\n The per-method hyperparameter search space.\n optimizer : HyperOptimizer, optional\n The parent optimizer. Used to size the initial population\n large enough for parallel pre-dispatch.\n population_size : int or \"auto\", optional\n The number of parallel workers for each method.\n sigma_min : float, optional\n Smallest perturbation scale (for the best-ranked worker).\n sigma_max : float, optional\n Largest perturbation scale (for the worst-ranked worker).\n elite_migrate_prob : float, optional\n Probability of copying best solution to worst worker.\n differential_prob : float, optional\n Per-sample probability of differential perturbation.\n patience : int or None, optional\n Generations without improvement before restart.\n method_exploration : float, optional\n Exploration strength for the LCB-based method chooser.\n method_temperature : float, optional\n Noise temperature for the LCB-based method chooser.\n exponential_param_power : float, optional\n Passed to the shared parameter mapping for ``FLOAT_EXP``.\n seed : None or int, optional\n Random seed.\n \"\"\"\n if population_size == \"auto\":\n max_ndim = max(\n num_params(\n build_params(\n space[m],\n exponential_param_power=exponential_param_power,\n )\n )\n for m in methods\n )\n population_size = max(\n 8,\n max(1, getattr(optimizer, \"pre_dispatch\", 1)),\n 4 * max_ndim,\n )\n\n self._method_chooser = LCBOptimizer(\n options=methods,\n exploration=method_exploration,\n temperature=method_temperature,\n seed=seed,\n )\n self._optimizers = {\n method: HyperPESampler(\n space[method],\n seed=seed,\n population_size=population_size,\n sigma_min=sigma_min,\n sigma_max=sigma_max,\n elite_migrate_prob=elite_migrate_prob,\n differential_prob=differential_prob,\n patience=patience,\n exponential_param_power=exponential_param_power,\n )\n for method in methods\n }\n\n def get_setting(self):\n \"\"\"Choose a contraction method, then request its next setting.\"\"\"\n method = self._method_chooser.ask()\n params_token, params = self._optimizers[method].ask()\n\n return {\n \"method\": method,\n \"params_token\": params_token,\n \"params\": params,\n }\n\n def report_result(self, setting, trial, score):\n \"\"\"Report a completed trial back to the method chooser and PE.\"\"\"\n self._method_chooser.tell(setting[\"method\"], score)\n self._optimizers[setting[\"method\"]].tell(\n setting[\"params_token\"], score\n )\n\n\nregister_hyper_optlib(\"pe\", PEOptLib)\nregister_hyper_optlib(\"parallelev\", PEOptLib)\n"
},
{
"path": "cotengra/experimental/hyper_pymoo.py",
"content": "\"\"\"Hyper optimization using pymoo single-objective algorithms.\n\nThis backend currently supports serial optimization only. Pymoo ask/tell\nalgorithms operate on generations/batches rather than individual trials, so\nthe integration buffers one full batch at a time and feeds it back when all\nbatch members have been evaluated.\n\"\"\"\n\nfrom ._param_mapping import (\n LCBOptimizer,\n build_params,\n convert_raw,\n num_params,\n)\nfrom .hyper import HyperOptLib, register_hyper_optlib\n\n\ndef _get_pymoo_algorithm(name):\n if name == \"de\":\n from pymoo.algorithms.soo.nonconvex.de import DE\n\n return DE\n if name == \"ga\":\n from pymoo.algorithms.soo.nonconvex.ga import GA\n\n return GA\n if name == \"pso\":\n from pymoo.algorithms.soo.nonconvex.pso import PSO\n\n return PSO\n\n if name == \"brkga\":\n from pymoo.algorithms.soo.nonconvex.brkga import BRKGA\n\n return BRKGA\n\n if name == \"es\":\n from pymoo.algorithms.soo.nonconvex.es import ES\n\n return ES\n\n if name == \"sres\":\n from pymoo.algorithms.soo.nonconvex.sres import SRES\n\n return SRES\n\n if name == \"isres\":\n from pymoo.algorithms.soo.nonconvex.isres import ISRES\n\n return ISRES\n\n if name == \"nrbo\":\n from pymoo.algorithms.soo.nonconvex.nrbo import NRBO\n\n return NRBO\n\n raise ValueError(f\"Unknown pymoo sampler {name}.\")\n\n\nclass HyperPymooSampler:\n \"\"\"Per-method ask/tell wrapper around a pymoo algorithm.\"\"\"\n\n def __init__(\n self,\n space,\n sampler=\"de\",\n sampler_opts=None,\n exponential_param_power=None,\n seed=None,\n ):\n import numpy as np\n from pymoo.core.evaluator import Evaluator\n from pymoo.core.problem import Problem\n from pymoo.core.termination import NoTermination\n from pymoo.problems.static import StaticProblem\n\n self._np = np\n self._Evaluator = Evaluator\n self._StaticProblem = StaticProblem\n\n self.params = build_params(\n space, exponential_param_power=exponential_param_power\n )\n self._ndim = num_params(self.params)\n self._problem = Problem(\n n_var=self._ndim,\n n_obj=1,\n n_constr=0,\n xl=np.full(self._ndim, -1.0),\n xu=np.full(self._ndim, 1.0),\n )\n\n sampler_opts = {} if sampler_opts is None else dict(sampler_opts)\n Algorithm = _get_pymoo_algorithm(sampler)\n self.algorithm = Algorithm(**sampler_opts)\n self.algorithm.setup(\n self._problem,\n termination=NoTermination(),\n seed=seed,\n verbose=False,\n )\n\n self._trial_counter = 0\n self._active_batch = None\n\n def ask(self):\n if self._active_batch is None:\n pop = self.algorithm.ask()\n xs = pop.get(\"X\")\n trial_numbers = tuple(\n range(\n self._trial_counter,\n self._trial_counter + len(xs),\n )\n )\n self._trial_counter += len(xs)\n settings = tuple(convert_raw(self.params, x.copy()) for x in xs)\n self._active_batch = {\n \"pop\": pop,\n \"trial_numbers\": trial_numbers,\n \"settings\": settings,\n \"scores\": {},\n \"next_index\": 0,\n }\n\n i = self._active_batch[\"next_index\"]\n self._active_batch[\"next_index\"] += 1\n return (\n self._active_batch[\"trial_numbers\"][i],\n self._active_batch[\"settings\"][i],\n )\n\n def tell(self, trial_number, score):\n batch = self._active_batch\n batch[\"scores\"][trial_number] = score\n\n if len(batch[\"scores\"]) != len(batch[\"trial_numbers\"]):\n return\n\n f = self._np.asarray(\n [batch[\"scores\"][t] for t in batch[\"trial_numbers\"]],\n dtype=float,\n ).reshape(-1, 1)\n static = self._StaticProblem(self._problem, F=f)\n self._Evaluator().eval(static, batch[\"pop\"])\n self.algorithm.tell(infills=batch[\"pop\"])\n self._active_batch = None\n\n\nclass PymooOptLib(HyperOptLib):\n \"\"\"Hyper-optimization using pymoo algorithms with LCB method choice.\"\"\"\n\n def setup(\n self,\n methods,\n space,\n optimizer=None,\n sampler=\"de\",\n sampler_opts=None,\n method_exploration=1.0,\n method_temperature=1.0,\n exponential_param_power=None,\n seed=None,\n **kwargs,\n ):\n if getattr(optimizer, \"_pool\", None) is not None:\n raise ValueError(\n \"The 'pymoo' optlib currently only supports serial \"\n \"hyper-optimization (`parallel=False`).\"\n )\n\n self._method_chooser = LCBOptimizer(\n options=methods,\n exploration=method_exploration,\n temperature=method_temperature,\n seed=seed,\n )\n self._optimizers = {\n method: HyperPymooSampler(\n space[method],\n sampler=sampler,\n sampler_opts=sampler_opts,\n exponential_param_power=exponential_param_power,\n seed=seed,\n )\n for method in methods\n }\n\n def get_setting(self):\n method = self._method_chooser.ask()\n params_token, params = self._optimizers[method].ask()\n\n return {\n \"method\": method,\n \"params_token\": params_token,\n \"params\": params,\n }\n\n def report_result(self, setting, trial, score):\n method = setting[\"method\"]\n self._method_chooser.tell(method, score)\n self._optimizers[method].tell(setting[\"params_token\"], score)\n\n\nregister_hyper_optlib(\"pymoo\", PymooOptLib)\n"
},
{
"path": "cotengra/experimental/hyper_scipy.py",
"content": "\"\"\"Hyper optimization using scipy gradient-free optimizers.\n\nSupported methods: ``differential_evolution``, ``dual_annealing``,\n``direct``, ``shgo``. Since these optimizers use a callback-style\nobjective, a background thread is used to invert the control flow into\nan ask/tell interface. Multiple workers are maintained per method to\nsupport parallel pre-dispatch (ask-ask-...-tell-tell-...).\n\"\"\"\n\nimport queue\nimport threading\n\nfrom ._param_mapping import (\n LCBOptimizer,\n build_params,\n convert_raw,\n num_params,\n)\nfrom .hyper import HyperOptLib, register_hyper_optlib\n\n_OPTIMIZER_NAMES = {\n \"differential_evolution\",\n \"dual_annealing\",\n \"direct\",\n \"shgo\",\n}\n\n\nclass _StopOptimization(Exception):\n \"\"\"Raised inside the objective to abort the scipy optimizer.\"\"\"\n\n\nclass ScipyAskTell:\n \"\"\"Ask/tell wrapper around a scipy global optimizer.\n\n The optimizer runs in a background thread. Each time it needs an objective\n evaluation it posts the candidate vector to ``_ask_q`` and blocks on\n ``_tell_q``. The caller drives progress by alternating\n ``ask()`` / ``tell()`` calls from the main thread.\n\n Parameters\n ----------\n method : str\n One of the supported scipy optimizer names.\n bounds : list[tuple[float, float]]\n Bounds for every raw dimension.\n kwargs\n Forwarded to the underlying scipy optimizer.\n \"\"\"\n\n def __init__(self, method, bounds, **kwargs):\n self.method = method\n self.bounds = bounds\n self.kwargs = kwargs\n\n self._ask_q = queue.Queue()\n self._tell_q = queue.Queue()\n self._thread = None\n self._stop = threading.Event()\n self.done = False\n\n # ---- internal --------------------------------------------------------\n\n def _get_optimizer_fn(self):\n from scipy import optimize\n\n return getattr(optimize, self.method)\n\n def _objective(self, x):\n if self._stop.is_set():\n raise _StopOptimization\n self._ask_q.put(x)\n val = self._tell_q.get()\n if self._stop.is_set():\n raise _StopOptimization\n return float(val)\n\n def _run(self):\n try:\n fn = self._get_optimizer_fn()\n fn(self._objective, self.bounds, **self.kwargs)\n except _StopOptimization:\n pass\n except Exception:\n pass\n finally:\n self.done = True\n # unblock any pending ask()\n self._ask_q.put(None)\n\n # ---- public ----------------------------------------------------------\n\n def start(self):\n \"\"\"Launch the optimizer in a background thread.\"\"\"\n self._thread = threading.Thread(target=self._run, daemon=True)\n self._thread.start()\n\n def ask(self):\n \"\"\"Block until the optimizer requests an evaluation.\n\n Returns\n -------\n x : ndarray or None\n The candidate vector, or ``None`` if the optimizer finished.\n \"\"\"\n x = self._ask_q.get()\n if self.done and x is None:\n return None\n return x\n\n def tell(self, score):\n \"\"\"Provide the objective value back to the optimizer.\"\"\"\n self._tell_q.put(score)\n\n def stop(self):\n \"\"\"Signal the background thread to stop and wait for it.\"\"\"\n self._stop.set()\n # unblock if the thread is waiting in _objective for a tell\n try:\n self._tell_q.put_nowait(float(\"inf\"))\n except queue.Full:\n pass\n if self._thread is not None:\n self._thread.join(timeout=5.0)\n\n\nclass HyperScipySampler:\n \"\"\"Per-method optimizer that wraps a pool of ``ScipyAskTell``\n workers with the ``Param``-based space mapping.\n\n Each ``ScipyAskTell`` worker is strictly serial\n (ask-tell-ask-tell), so to support parallel pre-dispatch\n (multiple ``ask`` calls before any ``tell``), we maintain a pool\n of workers. Each ``ask`` grabs a candidate from the next idle\n worker; each ``tell`` feeds the score back to the specific worker\n that produced that candidate.\n\n Parameters\n ----------\n space : dict[str, dict]\n The search space for one method.\n method : str\n Which scipy optimizer to use.\n n_workers : int, optional\n Number of concurrent ``ScipyAskTell`` threads to run.\n exponential_param_power : float, optional\n Power for ``ParamFloatExp``.\n kwargs\n Extra keyword arguments forwarded to the scipy optimizer.\n \"\"\"\n\n def __init__(\n self,\n space,\n method=\"differential_evolution\",\n n_workers=1,\n exponential_param_power=None,\n **kwargs,\n ):\n self.params = build_params(\n space, exponential_param_power=exponential_param_power\n )\n self._ndim = num_params(self.params)\n self._method = method\n self._scipy_opts = kwargs\n self._n_workers = n_workers\n\n # trial bookkeeping\n self._trial_counter = 0\n self._trial_to_worker = {}\n\n # worker pool: round-robin index for asks\n self._workers = []\n self._worker_idx = 0\n for _ in range(self._n_workers):\n self._workers.append(self._make_worker())\n\n def _make_worker(self):\n \"\"\"Create and start a fresh ``ScipyAskTell`` instance.\"\"\"\n w = ScipyAskTell(\n method=self._method,\n bounds=[(-1.0, 1.0)] * self._ndim,\n **self._scipy_opts,\n )\n w.start()\n return w\n\n def _next_worker(self):\n \"\"\"Pick the next worker in round-robin order, restarting any\n that have finished.\n \"\"\"\n for _ in range(self._n_workers):\n w = self._workers[self._worker_idx]\n self._worker_idx = (self._worker_idx + 1) % self._n_workers\n if w.done:\n # optimizer converged / exhausted – replace it\n w.stop()\n w = self._make_worker()\n # the index was already advanced, so store at previous\n prev = (self._worker_idx - 1) % self._n_workers\n self._workers[prev] = w\n return w\n\n # should never get here\n return self._workers[0]\n\n def ask(self):\n \"\"\"Return ``(trial_number, params_dict)``.\n\n Picks the next idle worker (round-robin), blocking until it\n has a candidate ready. If a worker has converged, it is\n automatically restarted.\n \"\"\"\n w = self._next_worker()\n x = w.ask()\n\n if x is None:\n # worker finished between _next_worker check and ask\n prev = (self._worker_idx - 1) % self._n_workers\n w = self._make_worker()\n self._workers[prev] = w\n x = w.ask()\n\n trial_number = self._trial_counter\n self._trial_to_worker[trial_number] = w\n self._trial_counter += 1\n\n return trial_number, convert_raw(self.params, x)\n\n def tell(self, trial_number, score):\n \"\"\"Report a score back to the specific worker that produced\n this trial.\n \"\"\"\n w = self._trial_to_worker.pop(trial_number, None)\n if w is not None and not w.done:\n w.tell(score)\n\n def stop(self):\n \"\"\"Stop all background threads.\"\"\"\n for w in self._workers:\n w.stop()\n\n\nclass ScipyOptLib(HyperOptLib):\n \"\"\"Hyper-optimization using scipy gradient-free optimizers with\n an LCB method selector.\n \"\"\"\n\n def setup(\n self,\n methods,\n space,\n optimizer=None,\n method=\"differential_evolution\",\n method_exploration=1.0,\n method_temperature=1.0,\n exponential_param_power=None,\n **scipy_opts,\n ):\n \"\"\"Initialize per-method scipy optimizers.\n\n Parameters\n ----------\n methods : list[str]\n The contraction methods to optimize over.\n space : dict[str, dict[str, dict]]\n The search space.\n optimizer : HyperOptimizer, optional\n The parent optimizer instance.\n method : str, optional\n Which scipy global optimizer to use. One of\n ``'differential_evolution'``, ``'dual_annealing'``,\n ``'direct'``, ``'shgo'``.\n method_exploration : float, optional\n Exploration parameter for the LCB method selector.\n method_temperature : float, optional\n Temperature parameter for the LCB method selector.\n exponential_param_power : float, optional\n Power for ``ParamFloatExp``.\n scipy_opts\n Extra keyword arguments forwarded to the scipy optimizer.\n \"\"\"\n if method not in _OPTIMIZER_NAMES:\n raise ValueError(\n f\"method must be one of {sorted(_OPTIMIZER_NAMES)}, \"\n f\"got {method!r}\"\n )\n\n n_workers = getattr(optimizer, \"_num_workers\", 1)\n # need at least as many workers per method as there are\n # parallel pre-dispatch slots, so that asks never block on a\n # worker waiting for its tell\n pre_dispatch = getattr(optimizer, \"pre_dispatch\", 1)\n n_workers = max(n_workers, pre_dispatch)\n\n self._method_chooser = LCBOptimizer(\n options=methods,\n exploration=method_exploration,\n temperature=method_temperature,\n )\n self._optimizers = {\n m: HyperScipySampler(\n space[m],\n method=method,\n n_workers=n_workers,\n exponential_param_power=exponential_param_power,\n **scipy_opts,\n )\n for m in methods\n }\n\n def get_setting(self):\n method = self._method_chooser.ask()\n trial_number, params = self._optimizers[method].ask()\n\n return {\n \"method\": method,\n \"params_token\": trial_number,\n \"params\": params,\n }\n\n def report_result(self, setting, trial, score):\n self._method_chooser.tell(setting[\"method\"], score)\n self._optimizers[setting[\"method\"]].tell(\n setting[\"params_token\"], score\n )\n\n def cleanup(self):\n \"\"\"Stop all background optimizer threads.\"\"\"\n for sampler in self._optimizers.values():\n sampler.stop()\n\n\nregister_hyper_optlib(\"scipy\", ScipyOptLib)\n"
},
{
"path": "cotengra/experimental/hyper_smac.py",
"content": "\"\"\"Hyper parameter optimization using SMAC3.\n\nhttps://automl.github.io/SMAC3/latest/\n\"\"\"\n\nfrom .hyper import HyperOptLib, register_hyper_optlib\n\n\ndef build_config_space(method, space):\n \"\"\"Build a SMAC ``ConfigurationSpace`` from a cotengra space dict.\n\n Parameters\n ----------\n method : str\n The method name (used as a prefix to avoid name collisions when\n multiple methods share the same parameter name).\n space : dict[str, dict]\n The search space for a single method.\n\n Returns\n -------\n cs : ConfigurationSpace\n \"\"\"\n from ConfigSpace import (\n CategoricalHyperparameter,\n UniformFloatHyperparameter,\n UniformIntegerHyperparameter,\n )\n from smac.configspace import ConfigurationSpace\n\n cs = ConfigurationSpace()\n for name, param in space.items():\n ptype = param[\"type\"]\n if ptype == \"FLOAT\":\n hp = UniformFloatHyperparameter(\n name, lower=param[\"min\"], upper=param[\"max\"]\n )\n elif ptype == \"FLOAT_EXP\":\n hp = UniformFloatHyperparameter(\n name, lower=param[\"min\"], upper=param[\"max\"], log=True\n )\n elif ptype == \"INT\":\n hp = UniformIntegerHyperparameter(\n name, lower=param[\"min\"], upper=param[\"max\"]\n )\n elif ptype == \"STRING\":\n hp = CategoricalHyperparameter(name, choices=param[\"options\"])\n elif ptype == \"BOOL\":\n hp = CategoricalHyperparameter(name, choices=[False, True])\n else:\n raise ValueError(f\"Unknown parameter type: {ptype!r}\")\n cs.add_hyperparameter(hp)\n return cs\n\n\ndef config_to_params(config):\n \"\"\"Convert a SMAC ``Configuration`` to a plain dict of parameters.\"\"\"\n return dict(config)\n\n\nclass SMACOptLib(HyperOptLib):\n \"\"\"Hyper-optimization using SMAC3 with per-method facades and\n a Lower Confidence Bound method selector.\n \"\"\"\n\n def setup(\n self,\n methods,\n space,\n optimizer=None,\n facade=\"BlackBoxFacade\",\n n_trials=10000,\n seed=0,\n method_exploration=1.0,\n method_temperature=1.0,\n **facade_opts,\n ):\n from smac import BlackBoxFacade, HyperparameterOptimizationFacade\n from smac.scenario import Scenario\n\n from ._param_mapping import LCBOptimizer\n\n self._method_chooser = LCBOptimizer(\n options=methods,\n exploration=method_exploration,\n temperature=method_temperature,\n )\n\n if isinstance(facade, str):\n facade_cls = {\n \"BlackBoxFacade\": BlackBoxFacade,\n \"HyperparameterOptimizationFacade\": HyperparameterOptimizationFacade,\n }[facade]\n else:\n facade_cls = facade\n\n self._facades = {}\n self._trial_infos = {}\n for method in methods:\n cs = build_config_space(method, space[method])\n scenario = Scenario(\n cs,\n n_trials=n_trials,\n seed=seed,\n deterministic=True,\n )\n self._facades[method] = facade_cls(\n scenario,\n target_function=lambda cfg, seed: 0.0,\n overwrite=True,\n logging_level=False,\n **facade_opts,\n )\n\n def get_setting(self):\n method = self._method_chooser.ask()\n smac = self._facades[method]\n info = smac.ask()\n trial_key = (method, id(info))\n self._trial_infos[trial_key] = info\n return {\n \"method\": method,\n \"trial_key\": trial_key,\n \"params\": config_to_params(info.config),\n }\n\n def report_result(self, setting, trial, score):\n from smac.runhistory import TrialValue\n\n method = setting[\"method\"]\n trial_key = setting[\"trial_key\"]\n info = self._trial_infos.pop(trial_key)\n self._method_chooser.tell(method, score)\n value = TrialValue(cost=score)\n self._facades[method].tell(info, value)\n\n\nregister_hyper_optlib(\"smac\", SMACOptLib)\n"
},
{
"path": "cotengra/experimental/multi.ipynb",
"content": "{\n \"cells\": [\n {\n \"cell_type\": \"code\",\n \"execution_count\": 1,\n \"id\": \"9576af0f-adc3-44e4-88d0-7a57dc97a1d5\",\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"%config InlineBackend.figure_formats = ['svg']\\n\",\n \"import math # noqa\\n\",\n \"import random # noqa\\n\",\n \"import itertools # noqa\\n\",\n \"import functools # noqa\\n\",\n \"import collections # noqa\\n\",\n \"import tqdm # noqa\\n\",\n \"import numpy as np # noqa\\n\",\n \"import matplotlib as mpl # noqa\\n\",\n \"import matplotlib.pyplot as plt # noqa\\n\",\n \"import quimb as qu # noqa\\n\",\n \"import xyzpy as xyz # noqa\\n\",\n \"import autoray as ar # noqa\\n\",\n \"import cotengra as ctg # noqa\\n\",\n \"import quimb.tensor as qtn # noqa\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 2,\n \"id\": \"881a2afe-11e9-43f4-a62f-0bda257b7750\",\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"tn = qtn.PEPS.rand(4, 5, 4)\\n\",\n \"# tn = qtn.MPS_rand_state(20, 7)\\n\",\n \"\\n\",\n \"# tn = qtn.TN2D_rand(32, 32, D=2)\\n\",\n \"\\n\",\n \"\\n\",\n \"inputs, output, size_dict = tn.get_inputs_output_size_dict()\\n\",\n \"symbol_map = tn.get_symbol_map()\\n\",\n \"\\n\",\n \"rng = ctg.utils.get_rng(42)\\n\",\n \"\\n\",\n \"M = 200\\n\",\n \"\\n\",\n \"sliced_inds = output\\n\",\n \"\\n\",\n \"configs = [{ix: rng.choice([0, 1]) for ix in output} for _ in range(M)]\\n\",\n \"\\n\",\n \"tree = ctg.array_contract_tree(\\n\",\n \" inputs,\\n\",\n \" output,\\n\",\n \" size_dict,\\n\",\n \")\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 11,\n \"id\": \"46dd3357-18b0-4a80-afd5-935f6a4029f9\",\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"# inputs, output, shapes, size_dict = ctg.utils.randreg_equation(100, 4, seed=42)\\n\",\n \"inputs, output, shapes, size_dict = ctg.utils.lattice_equation(\\n\",\n \" [28, 28], d_max=2\\n\",\n \")\\n\",\n \"\\n\",\n \"opt = ctg.HyperOptimizer(\\n\",\n \" simulated_annealing_opts=dict(target_size=2**24),\\n\",\n \" max_repeats=8,\\n\",\n \" progbar=True,\\n\",\n \")\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 12,\n \"id\": \"92b7e150-cffd-40d7-ab26-7cf34af79476\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stderr\",\n \"output_type\": \"stream\",\n \"text\": [\n \"F=14.05 C=14.72 S=24 P=25.02 $=8192: 100%|██████████| 16/16 [01:56<00:00, 7.31s/it] \\n\"\n ]\n }\n ],\n \"source\": [\n \"tree = opt.search(inputs, output, size_dict)\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 13,\n \"id\": \"d47c757e-ecb8-43c8-91a7-d464cb146353\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"data\": {\n \"text/plain\": [\n \"8192\"\n ]\n },\n \"execution_count\": 13,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"sliced_inds = tree.sliced_inds\\n\",\n \"configs = [\\n\",\n \" dict(zip(sliced_inds, x))\\n\",\n \" for x in itertools.product(*[range(size_dict[i]) for i in sliced_inds])\\n\",\n \"]\\n\",\n \"len(configs)\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 14,\n \"id\": \"771de335-5945-4909-902f-0f77b13d6c7e\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"data\": {\n \"text/plain\": [\n \"1.573196922320367\"\n ]\n },\n \"execution_count\": 14,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"tree.contraction_cost() / tree.unslice_all().contraction_cost()\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 15,\n \"id\": \"b020426e-19a8-4700-ab4a-80dc8d714680\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"data\": {\n \"text/plain\": [\n \"24.0\"\n ]\n },\n \"execution_count\": 15,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"tree.max_size(log=2)\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 3,\n \"id\": \"c3802312-f19a-481b-9b17-371c28b49d37\",\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"from cotengra.core import (\\n\",\n \" ContractionTree,\\n\",\n \" get_score_fn,\\n\",\n \")\\n\",\n \"from cotengra.scoring import Objective\\n\",\n \"\\n\",\n \"\\n\",\n \"class MultiObjective(Objective):\\n\",\n \" __slots__ = (\\\"num_configs\\\",)\\n\",\n \"\\n\",\n \" def __init__(self, num_configs):\\n\",\n \" self.num_configs = num_configs\\n\",\n \"\\n\",\n \" def compute_mult(self, dims):\\n\",\n \" raise NotImplementedError\\n\",\n \"\\n\",\n \" def estimate_node_mult(self, tree, node):\\n\",\n \" return self.compute_mult(\\n\",\n \" [tree.size_dict[ix] for ix in tree.get_node_var_inds(node)]\\n\",\n \" )\\n\",\n \"\\n\",\n \" def estimate_node_cache_mult(self, tree, node, sliced_ind_ordering):\\n\",\n \" node_var_inds = tree.get_node_var_inds(node)\\n\",\n \"\\n\",\n \" # indices which are the first 'k' in the sliced ordering\\n\",\n \" non_heavy_inds = [\\n\",\n \" ix\\n\",\n \" for ix in tree.get_node_var_inds(node)\\n\",\n \" if ix not in sliced_ind_ordering[: len(node_var_inds)]\\n\",\n \" ]\\n\",\n \"\\n\",\n \" # each of these cycles 'out of sync' and thus must be kept\\n\",\n \" return self.compute_mult([tree.size_dict[ix] for ix in non_heavy_inds])\\n\",\n \"\\n\",\n \"\\n\",\n \"class MultiObjectiveDense(MultiObjective):\\n\",\n \" \\\"\\\"\\\"Number of intermediate configurations is expected to scale as if all\\n\",\n \" configurations are present.\\n\",\n \" \\\"\\\"\\\"\\n\",\n \"\\n\",\n \" __slots__ = (\\\"num_configs\\\",)\\n\",\n \"\\n\",\n \" def compute_mult(self, dims):\\n\",\n \" return math.prod(dims)\\n\",\n \"\\n\",\n \"\\n\",\n \"def expected_coupons(num_sub, num_total):\\n\",\n \" \\\"\\\"\\\"If we draw a random 'coupon` which can take `num_sub` different values\\n\",\n \" `num_total` times, how many unique coupons will we expect?\\n\",\n \" \\\"\\\"\\\"\\n\",\n \" # return min(num_sub, num_total)\\n\",\n \" return num_sub * (1 - (1 - 1 / num_sub) ** num_total)\\n\",\n \"\\n\",\n \"\\n\",\n \"class MultiObjectiveUniform(MultiObjective):\\n\",\n \" \\\"\\\"\\\"Number of intermediate configurations is expected to scale as if all\\n\",\n \" configurations are randomly draw from a uniform distribution.\\n\",\n \" \\\"\\\"\\\"\\n\",\n \"\\n\",\n \" __slots__ = (\\\"num_configs\\\",)\\n\",\n \"\\n\",\n \" def compute_mult(self, dims):\\n\",\n \" return expected_coupons(math.prod(dims), self.num_configs)\\n\",\n \"\\n\",\n \"\\n\",\n \"class MultiObjectiveLinear(MultiObjective):\\n\",\n \" \\\"\\\"\\\"Number of intermediate configurations is expected to scale linearly with\\n\",\n \" respect to number of variable indices (e.g. VMC like 'locally connected'\\n\",\n \" configurations).\\n\",\n \" \\\"\\\"\\\"\\n\",\n \"\\n\",\n \" __slots__ = (\\\"num_configs\\\", \\\"coeff\\\")\\n\",\n \"\\n\",\n \" def __init__(self, num_configs, coeff=1):\\n\",\n \" self.coeff = coeff\\n\",\n \" super().__init__(num_configs=num_configs)\\n\",\n \"\\n\",\n \" def compute_mult(self, dims):\\n\",\n \" return min(self.coeff * len(dims), self.num_configs)\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 18,\n \"id\": \"2a6d0a84-e0a0-4e4a-a92b-fd70e89ab75b\",\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"M = len(configs)\\n\",\n \"# objective = MultiObjectiveUniform(num_configs=M)\\n\",\n \"# objective = MultiObjectiveDense(num_configs=M)\\n\",\n \"objective = MultiObjectiveLinear(num_configs=M, coeff=M / len(sliced_inds))\\n\",\n \"\\n\",\n \"mtree = ctg.ContractionTreeMulti.from_path(\\n\",\n \" inputs,\\n\",\n \" output,\\n\",\n \" size_dict,\\n\",\n \" sliced_inds=sliced_inds,\\n\",\n \" objective=objective,\\n\",\n \" path=tree.get_path(),\\n\",\n \")\\n\",\n \"mtree.reorder_contractions_for_peak_est()\\n\",\n \"mtree.reorder_sliced_inds()\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 19,\n \"id\": \"0a89762d-f34b-4c42-b187-c291be640341\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"data\": {\n \"text/plain\": [\n \"1\"\n ]\n },\n \"execution_count\": 19,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"mtree.multiplicity\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 20,\n \"id\": \"8482cdf3-7f5f-423d-8bff-51188fd4272e\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"data\": {\n \"text/plain\": [\n \"{'flops': 5273600.0, 'write': 6913, 'size': 1024}\"\n ]\n },\n \"execution_count\": 20,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"mtree.contract_stats()\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 21,\n \"id\": \"2b380e14-f8ab-4e18-b83e-3c6177790856\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"data\": {\n \"text/plain\": [\n \"8.5302276561314\"\n ]\n },\n \"execution_count\": 21,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"math.log10(339021824)\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 22,\n \"id\": \"d7d99d38-bf07-40f4-9230-9893ad0fd939\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"data\": {\n \"text/plain\": [\n \"4.920123326290724\"\n ]\n },\n \"execution_count\": 22,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"math.log10(\\n\",\n \" sum(\\n\",\n \" ContractionTree.get_flops(mtree, node)\\n\",\n \" for node, _, _ in mtree.descend()\\n\",\n \" )\\n\",\n \")\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 23,\n \"id\": \"f342e882-ab08-4707-86a7-ad93bfa8ab2a\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"data\": {\n \"text/plain\": [\n \"6.722107185681003\"\n ]\n },\n \"execution_count\": 23,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"math.log10(sum(mtree.get_flops(node) for node, _, _ in mtree.descend()))\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 24,\n \"id\": \"38609b23-ef7c-44eb-b382-977f79d98d8c\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"data\": {\n \"text/plain\": [\n \"6.722107185681002\"\n ]\n },\n \"execution_count\": 24,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"mtree.contraction_cost(log=10)\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 25,\n \"id\": \"bb422dce-0da9-4cc0-80c0-02657691d71c\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"data\": {\n \"text/plain\": [\n \"{'flops': 6203392, 'write': 248008, 'size': 1024, 'peak': 111360}\"\n ]\n },\n \"execution_count\": 25,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"stats = mtree.exact_multi_stats(configs)\\n\",\n \"stats\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 12,\n \"id\": \"84a0c184-6e73-4a39-88c3-5b0c8c7edafa\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"data\": {\n \"text/plain\": [\n \"6.792629225636649\"\n ]\n },\n \"execution_count\": 12,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"math.log10(stats[\\\"flops\\\"])\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 16,\n \"id\": \"8449c4a8-4e1c-43dd-b614-7dba4b28e2af\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"data\": {\n \"text/plain\": [\n \"16.98299357469431\"\n ]\n },\n \"execution_count\": 16,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"mtree.peak_size(log=2)\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 17,\n \"id\": \"d32c40a1-19c3-4f62-9b38-aa3d067eb838\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"data\": {\n \"text/plain\": [\n \"16.764871590736092\"\n ]\n },\n \"execution_count\": 17,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"math.log2(stats[\\\"peak\\\"])\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 19,\n \"id\": \"7fd0ac7d-ad07-4654-8fbc-278e9f3d07d3\",\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"from cotengra.pathfinders.path_simulated_annealing import (\\n\",\n \" _describe_tree,\\n\",\n \" _slice_tree_basic,\\n\",\n \" _slice_tree_drift,\\n\",\n \" _slice_tree_reslice,\\n\",\n \" get_rng,\\n\",\n \" linspace_generator,\\n\",\n \")\\n\",\n \"\\n\",\n \"\\n\",\n \"def compute_contracted_info(\\n\",\n \" legsa,\\n\",\n \" legsb,\\n\",\n \" appearances,\\n\",\n \" size_dict,\\n\",\n \" var_inds_a,\\n\",\n \" var_inds_b,\\n\",\n \"):\\n\",\n \" \\\"\\\"\\\"Compute the contracted legs, cost and size of a pair of legs.\\n\",\n \"\\n\",\n \" Parameters\\n\",\n \" ----------\\n\",\n \" legsa : dict[str, int]\\n\",\n \" The legs of the first tensor.\\n\",\n \" legsb : dict[str, int]\\n\",\n \" The legs of the second tensor.\\n\",\n \" appearances : dict[str, int]\\n\",\n \" The total number of appearances of each index in the contraction.\\n\",\n \" size_dict : dict[str, int]\\n\",\n \" The size of each index.\\n\",\n \"\\n\",\n \" Returns\\n\",\n \" -------\\n\",\n \" legsab : dict[str, int]\\n\",\n \" The contracted legs.\\n\",\n \" cost : int\\n\",\n \" The cost of the contraction.\\n\",\n \" size : int\\n\",\n \" The size of the resulting tensor.\\n\",\n \" \\\"\\\"\\\"\\n\",\n \" legsab = {}\\n\",\n \" cost = 1\\n\",\n \" size = 1\\n\",\n \"\\n\",\n \" # handle all left indices\\n\",\n \" for ix, ix_count in legsa.items():\\n\",\n \" d = size_dict[ix]\\n\",\n \" # all involved indices contribute to cost\\n\",\n \" cost *= d\\n\",\n \" if ix in legsb:\\n\",\n \" ix_count += legsb[ix]\\n\",\n \" if ix_count < appearances[ix]:\\n\",\n \" # index appears on output\\n\",\n \" legsab[ix] = ix_count\\n\",\n \" # and so contributes to size\\n\",\n \" size *= d\\n\",\n \"\\n\",\n \" # now handle right indices that we haven't seen yet\\n\",\n \" for ix, ix_count in legsb.items():\\n\",\n \" if ix not in legsa:\\n\",\n \" d = size_dict[ix]\\n\",\n \" cost *= d\\n\",\n \" if ix_count < appearances[ix]:\\n\",\n \" legsab[ix] = ix_count\\n\",\n \" size *= d\\n\",\n \"\\n\",\n \" return legsab, cost, size, var_inds_a | var_inds_b\\n\",\n \"\\n\",\n \"\\n\",\n \"def simulated_anneal_tree(\\n\",\n \" tree,\\n\",\n \" tfinal=0.05,\\n\",\n \" tstart=2,\\n\",\n \" tsteps=50,\\n\",\n \" numiter=50,\\n\",\n \" minimize=None,\\n\",\n \" target_size=None,\\n\",\n \" target_size_initial=None,\\n\",\n \" slice_mode=\\\"basic\\\",\\n\",\n \" seed=None,\\n\",\n \" progbar=False,\\n\",\n \" inplace=False,\\n\",\n \"):\\n\",\n \" \\\"\\\"\\\"Perform a simulated annealing optimization of this contraction\\n\",\n \" tree, based on \\\"Multi-Tensor Contraction for XEB Verification of\\n\",\n \" Quantum Circuits\\\" by Gleb Kalachev, Pavel Panteleev, Man-Hong Yung\\n\",\n \" (arXiv:2108.05665), and the \\\"treesa\\\" implementation in\\n\",\n \" OMEinsumContractionOrders.jl by Jin-Guo Liu and Pan Zhang.\\n\",\n \"\\n\",\n \" Parameters\\n\",\n \" ----------\\n\",\n \" tfinal : float, optional\\n\",\n \" The final temperature.\\n\",\n \" tstart : float, optional\\n\",\n \" The starting temperature.\\n\",\n \" tsteps : int, optional\\n\",\n \" The number of temperature steps.\\n\",\n \" numiter : int, optional\\n\",\n \" The number of sweeps at each temperature step.\\n\",\n \" minimize : {'flops', 'combo', 'write', 'size', ...}, optional\\n\",\n \" The objective function to minimize.\\n\",\n \" target_size : int, optional\\n\",\n \" The target size to slice the contraction to. A schedule is used to\\n\",\n \" reach this only at the final temperature step.\\n\",\n \" target_size_initial : int, optional\\n\",\n \" The initial target size to use in the slicing schedule. If None, then\\n\",\n \" the current size is used.\\n\",\n \" slice_mode : {'basic', 'reslice', 'drift'}, optional\\n\",\n \" The mode for slicing the contraction tree within each annealing\\n\",\n \" iteration. 'basic' always unslices a random index and then slices to\\n\",\n \" the target size. 'reslice' unslices all indices and then slices to the\\n\",\n \" target size. 'drift' unslices a random index with probability 1/4 and\\n\",\n \" slices to the target size with probability 3/4. It is therefore not\\n\",\n \" guaranteed to reach the target size, but may be more explorative for\\n\",\n \" long annealing schedules.\\n\",\n \" seed : int, optional\\n\",\n \" A random seed.\\n\",\n \" progbar : bool, optional\\n\",\n \" Whether to show live progress.\\n\",\n \" inplace : bool, optional\\n\",\n \" Whether to perform the optimization inplace.\\n\",\n \"\\n\",\n \" Returns\\n\",\n \" -------\\n\",\n \" ContractionTree\\n\",\n \" \\\"\\\"\\\"\\n\",\n \"\\n\",\n \" tree = tree if inplace else tree.copy()\\n\",\n \" # ensure stats tracking is on\\n\",\n \" tree.contract_stats()\\n\",\n \"\\n\",\n \" if minimize is None:\\n\",\n \" minimize = tree.get_default_objective()\\n\",\n \" scorer = get_score_fn(minimize)\\n\",\n \" rng = get_rng(seed)\\n\",\n \"\\n\",\n \" # create a schedule for annealing temperatures\\n\",\n \" temps = linspace_generator(tstart, tfinal, tsteps, log=True)\\n\",\n \"\\n\",\n \" if target_size is not None:\\n\",\n \" # create a schedule for slicing target sizes\\n\",\n \" if target_size_initial is None:\\n\",\n \" # start with the current size\\n\",\n \" current_size = max(tree.contraction_width(log=None), target_size)\\n\",\n \" else:\\n\",\n \" current_size = max(target_size_initial, target_size)\\n\",\n \"\\n\",\n \" target_sizes = linspace_generator(\\n\",\n \" current_size,\\n\",\n \" target_size,\\n\",\n \" tsteps,\\n\",\n \" log=True,\\n\",\n \" )\\n\",\n \" _slice_tree = {\\n\",\n \" \\\"basic\\\": _slice_tree_basic,\\n\",\n \" \\\"reslice\\\": _slice_tree_reslice,\\n\",\n \" \\\"drift\\\": _slice_tree_drift,\\n\",\n \" }[slice_mode]\\n\",\n \" else:\\n\",\n \" target_sizes = itertools.repeat(None)\\n\",\n \"\\n\",\n \" def _slice_tree(tree, current_target_size, rng):\\n\",\n \" pass\\n\",\n \"\\n\",\n \" if progbar:\\n\",\n \" import tqdm\\n\",\n \"\\n\",\n \" pbar = tqdm.tqdm(total=tsteps)\\n\",\n \" pbar.set_description(_describe_tree(tree))\\n\",\n \"\\n\",\n \" for temp in temps:\\n\",\n \" # handle slicing\\n\",\n \" _slice_tree(tree, next(target_sizes), rng)\\n\",\n \"\\n\",\n \" for _ in range(numiter):\\n\",\n \" candidates = [tree.root]\\n\",\n \"\\n\",\n \" while candidates:\\n\",\n \" p = candidates.pop(0)\\n\",\n \" l, r = tree.children[p]\\n\",\n \"\\n\",\n \" # check which local moves are possible\\n\",\n \" if len(l) == 1:\\n\",\n \" if len(r) == 1:\\n\",\n \" # both are leaves\\n\",\n \" continue\\n\",\n \" else:\\n\",\n \" # left is leaf\\n\",\n \" rule = rng.randint(2, 3)\\n\",\n \" elif len(r) == 1:\\n\",\n \" # right is leaf\\n\",\n \" rule = rng.randint(0, 1)\\n\",\n \" else:\\n\",\n \" # neither are leaves\\n\",\n \" rule = rng.randint(0, 3)\\n\",\n \"\\n\",\n \" if rule < 2:\\n\",\n \" # ((AB)C)\\n\",\n \" x, c = l, r\\n\",\n \" a, b = tree.children[x]\\n\",\n \" if rule == 0:\\n\",\n \" # -> ((AC)B)\\n\",\n \" new_order = [a, c, b]\\n\",\n \" else:\\n\",\n \" # -> (A(BC))\\n\",\n \" new_order = [b, c, a]\\n\",\n \" else:\\n\",\n \" # (A(BC))\\n\",\n \" a, x = l, r\\n\",\n \" b, c = tree.children[x]\\n\",\n \" if rule == 2:\\n\",\n \" # -> (B(AC))\\n\",\n \" new_order = [a, c, b]\\n\",\n \" else:\\n\",\n \" # -> (C(AB))\\n\",\n \" new_order = [a, b, c]\\n\",\n \"\\n\",\n \" current_score = math.log2(\\n\",\n \" tree.get_flops(p) + tree.get_flops(x)\\n\",\n \" )\\n\",\n \"\\n\",\n \" # current_score = scorer.score_local(\\n\",\n \" # flops=[tree.get_flops(p), tree.get_flops(x)],\\n\",\n \" # size=[tree.get_size(p), tree.get_size(x)],\\n\",\n \" # )\\n\",\n \"\\n\",\n \" # legs0 = tree.get_legs(new_order[0])\\n\",\n \" # legs1 = tree.get_legs(new_order[1])\\n\",\n \" # if any(ix0 in legs1 for ix0 in legs0):\\n\",\n \"\\n\",\n \" # compute new intermediate\\n\",\n \" new_legs0, new_cost0, new_size0, new_var_inds0 = (\\n\",\n \" compute_contracted_info(\\n\",\n \" tree.get_legs(new_order[0]),\\n\",\n \" tree.get_legs(new_order[1]),\\n\",\n \" tree.appearances,\\n\",\n \" tree.size_dict,\\n\",\n \" tree.get_node_var_inds(new_order[0]),\\n\",\n \" tree.get_node_var_inds(new_order[1]),\\n\",\n \" )\\n\",\n \" )\\n\",\n \"\\n\",\n \" new_mult_cost0 = new_cost0 * scorer.estimate_mult(\\n\",\n \" [tree.size_dict[ix] for ix in new_var_inds0]\\n\",\n \" )\\n\",\n \"\\n\",\n \" # compute new parent costs\\n\",\n \" new_legs1, new_cost1, new_size1, new_var_inds1 = (\\n\",\n \" compute_contracted_info(\\n\",\n \" new_legs0,\\n\",\n \" tree.get_legs(new_order[2]),\\n\",\n \" tree.appearances,\\n\",\n \" tree.size_dict,\\n\",\n \" new_var_inds0,\\n\",\n \" tree.get_node_var_inds(new_order[2]),\\n\",\n \" )\\n\",\n \" )\\n\",\n \"\\n\",\n \" new_mult_cost1 = new_cost1 * scorer.estimate_mult(\\n\",\n \" [tree.size_dict[ix] for ix in new_var_inds1]\\n\",\n \" )\\n\",\n \"\\n\",\n \" proposed_score = math.log2(new_mult_cost0 + new_mult_cost1)\\n\",\n \"\\n\",\n \" # proposed_score = scorer.score_local(\\n\",\n \" # flops=[new_cost0, new_cost1],\\n\",\n \" # size=[new_size0, new_size1],\\n\",\n \" # )\\n\",\n \"\\n\",\n \" dE = proposed_score - current_score\\n\",\n \" accept = (dE <= 0) or (math.log(rng.random()) < -dE / temp)\\n\",\n \"\\n\",\n \" if accept:\\n\",\n \" tree._remove_node(p)\\n\",\n \" tree._remove_node(x)\\n\",\n \"\\n\",\n \" tree.contract_nodes_pair(\\n\",\n \" tree.contract_nodes_pair(\\n\",\n \" new_order[0],\\n\",\n \" new_order[1],\\n\",\n \" legs=new_legs0,\\n\",\n \" cost=new_cost0,\\n\",\n \" size=new_size0,\\n\",\n \" ),\\n\",\n \" new_order[2],\\n\",\n \" legs=new_legs1,\\n\",\n \" cost=new_cost1,\\n\",\n \" size=new_size1,\\n\",\n \" )\\n\",\n \"\\n\",\n \" if progbar:\\n\",\n \" pbar.set_description(\\n\",\n \" f\\\"T: {temp:.2e} \\\" + _describe_tree(tree),\\n\",\n \" refresh=False,\\n\",\n \" )\\n\",\n \"\\n\",\n \" # check which children to recurse into\\n\",\n \" l, r = tree.children[p]\\n\",\n \" if len(l) > 2:\\n\",\n \" candidates.append(l)\\n\",\n \" if len(r) > 2:\\n\",\n \" candidates.append(r)\\n\",\n \"\\n\",\n \" if progbar:\\n\",\n \" pbar.update()\\n\",\n \"\\n\",\n \" return tree\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 32,\n \"id\": \"d61c20c0-e3a4-489a-ad78-4b9e2097bc69\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stderr\",\n \"output_type\": \"stream\",\n \"text\": [\n \"T: 5.00e-02 F=6.708 C=6.744 S=10 P=11.74 $=1: 100%|██████████| 50/50 [00:00<00:00, 55.26it/s]\\n\"\n ]\n },\n {\n \"data\": {\n \"text/plain\": [\n \"\"\n ]\n },\n \"execution_count\": 32,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"simulated_anneal_tree(mtree, inplace=True, progbar=True)\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 33,\n \"id\": \"0edc1201-8e64-4023-982f-6b5df8a28a4e\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"data\": {\n \"text/plain\": [\n \"{'flops': 7598080, 'write': 475336, 'size': 1024, 'peak': 61953}\"\n ]\n },\n \"execution_count\": 33,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"stats = mtree.exact_multi_stats(configs)\\n\",\n \"stats\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 34,\n \"id\": \"8b3ce36c-f79d-4304-b585-40e049ab3915\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"data\": {\n \"text/plain\": [\n \"6.880703861918839\"\n ]\n },\n \"execution_count\": 34,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"math.log10(stats[\\\"flops\\\"])\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 7,\n \"id\": \"dc812dca-b634-490b-9db9-9fb1e63f46b1\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"data\": {\n \"text/plain\": [\n \"MultiObjectiveLinear(num_configs=200, coeff=2.0)\"\n ]\n },\n \"execution_count\": 7,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"mtree.get_default_objective()\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 10,\n \"id\": \"61d8fa34-d7da-44ab-bb61-e613c4b62b98\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"data\": {\n \"text/plain\": [\n \"0.3460989920591954\"\n ]\n },\n \"execution_count\": 10,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"mtree.estimate_multi_total_flops() / (ContractionTree.total_flops(mtree) * M)\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 12,\n \"id\": \"9f2d113c-44fd-43a1-bc3c-b26739687042\",\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"configs.sort(key=lambda c: tuple(c[ix] for ix in sliced_ind_ordering))\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 126,\n \"id\": \"cd241688-eda3-4d53-9b4b-cd2712ecf6e3\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"data\": {\n \"text/plain\": [\n \"3.920360304532837e-21\"\n ]\n },\n \"execution_count\": 126,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"stats[\\\"flops\\\"] / tree.total_flops()\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 127,\n \"id\": \"56abeb3e-a518-4faa-a7f2-087e5cb111c7\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"data\": {\n \"text/plain\": [\n \"3.920360304532837e-21\"\n ]\n },\n \"execution_count\": 127,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"stats[\\\"flops\\\"] / tree.unslice_all().total_flops()\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 128,\n \"id\": \"d5f5ce47-0408-463b-b4a6-d7980efde2d7\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"data\": {\n \"text/plain\": [\n \"104.0\"\n ]\n },\n \"execution_count\": 128,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"tree.unslice_all().peak_size(log=2)\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 129,\n \"id\": \"611f63a0-3f6b-448e-98ba-b0177ace9a27\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"data\": {\n \"text/plain\": [\n \"28.0\"\n ]\n },\n \"execution_count\": 129,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"ContractionTree.max_size(mtree, log=2)\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 130,\n \"id\": \"67ddde83-1531-4a60-af64-2ff03add40a0\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"data\": {\n \"text/plain\": [\n \"29.067312900309663\"\n ]\n },\n \"execution_count\": 130,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"math.log2(stats[\\\"peak\\\"])\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 120,\n \"id\": \"1d32e5ae-9583-4ffe-a2ed-a9116e188920\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"data\": {\n \"image/svg+xml\": [\n \"
![\"cotengra\"](\"https://imgur.com/OM5XyaD.png\")
![\"cotengra\"](\"https://imgur.com/jMO138y.png\")