opt_einsum.contractEvaluates the Einstein summation convention on the operands. A drop in +replacement for NumPy's einsum function that optimizes the order of contraction +to reduce overall scaling at the cost of several intermediate arrays.
+ + +Parameters:
+| Name | +Type | +Description | +Default | +
|---|---|---|---|
subscripts |
+
+ Union[str, ArrayType]
+ |
+
+
+
+ Specifies the subscripts for summation. + |
+ + required + | +
*operands |
+
+ Union[ArrayType, Collection[int]]
+ |
+
+
+
+ These are the arrays for the operation. + |
+
+ ()
+ |
+
out |
+
+ Optional[ArrayType]
+ |
+
+
+
+ A output array in which set the resulting output. + |
+
+ None
+ |
+
dtype |
+
+ Optional[str]
+ |
+
+
+
+ The dtype of the given contraction, see np.einsum. + |
+
+ None
+ |
+
order |
+
+ _OrderKACF
+ |
+
+
+
+ The order of the resulting contraction, see np.einsum. + |
+
+ 'K'
+ |
+
casting |
+
+ _Casting
+ |
+
+
+
+ The casting procedure for operations of different dtype, see np.einsum. + |
+
+ 'safe'
+ |
+
use_blas |
+
+ bool
+ |
+
+
+
+ Do you use BLAS for valid operations, may use extra memory for more intermediates. + |
+
+ True
+ |
+
optimize |
+
+ OptimizeKind
+ |
+
+
+
+
|
+
+ True
+ |
+
memory_limit |
+
+ _MemoryLimit
+ |
+
+
+
+
The default is None. Note that imposing a limit can make contractions +exponentially slower to perform. + |
+
+ None
+ |
+
backend |
+
+ BackendType
+ |
+
+
+
+ Which library to use to perform the required |
+
+ 'auto'
+ |
+
Returns:
+| Type | +Description | +
|---|---|
+ ArrayType
+ |
+
+
+
+ The result of the einsum expression. + |
+
This function should produce a result identical to that of NumPy's einsum
+function. The primary difference is contract will attempt to form
+intermediates which reduce the overall scaling of the given einsum contraction.
+By default the worst intermediate formed will be equal to that of the largest
+input array. For large einsum expressions with many input arrays this can
+provide arbitrarily large (1000 fold+) speed improvements.
For contractions with just two tensors this function will attempt to use +NumPy's built-in BLAS functionality to ensure that the given operation is +performed optimally. When NumPy is linked to a threaded BLAS, potential +speedups are on the order of 20-100 for a six core machine.
+opt_einsum/contract.py509 +510 +511 +512 +513 +514 +515 +516 +517 +518 +519 +520 +521 +522 +523 +524 +525 +526 +527 +528 +529 +530 +531 +532 +533 +534 +535 +536 +537 +538 +539 +540 +541 +542 +543 +544 +545 +546 +547 +548 +549 +550 +551 +552 +553 +554 +555 +556 +557 +558 +559 +560 +561 +562 +563 +564 +565 +566 +567 +568 +569 +570 +571 +572 +573 +574 +575 +576 +577 +578 +579 +580 +581 +582 +583 +584 +585 +586 +587 +588 +589 +590 +591 +592 +593 +594 +595 +596 +597 +598 +599 +600 +601 +602 +603 +604 +605 +606 +607 +608 +609 +610 +611 +612 +613 +614 +615 +616 +617 +618 +619 +620 | |
opt_einsum.contract_pathFind a contraction order path, without performing the contraction.
Parameters:
+| Name | +Type | +Description | +Default | +
|---|---|---|---|
subscripts |
+
+ Any
+ |
+
+
+
+ Specifies the subscripts for summation. + |
+ + required + | +
*operands |
+
+ Any
+ |
+
+
+
+ These are the arrays for the operation. + |
+
+ ()
+ |
+
use_blas |
+
+ bool
+ |
+
+
+
+ Do you use BLAS for valid operations, may use extra memory for more intermediates. + |
+
+ True
+ |
+
optimize |
+
+ OptimizeKind
+ |
+
+
+
+ Choose the type of path the contraction will be optimized with.
+- if a list is given uses this as the path.
+- |
+
+ True
+ |
+
memory_limit |
+
+ _MemoryLimit
+ |
+
+
+
+ Give the upper bound of the largest intermediate tensor contract will build.
+- None or -1 means there is no limit
+- The default is None. Note that imposing a limit can make contractions +exponentially slower to perform. + |
+
+ None
+ |
+
shapes |
+
+ bool
+ |
+
+
+
+ Whether |
+
+ False
+ |
+
Returns:
+| Name | Type | +Description | +
|---|---|---|
path |
+ PathType
+ |
+
+
+
+ The optimized einsum contraciton path + |
+
PathInfo |
+ PathInfo
+ |
+
+
+
+ A printable object containing various information about the path found. + |
+
The resulting path indicates which terms of the input contraction should be +contracted first, the result of this contraction is then appended to the end of +the contraction list.
+Examples:
+We can begin with a chain dot example. In this case, it is optimal to
+contract the b and c tensors represented by the first element of the path (1,
+2). The resulting tensor is added to the end of the contraction and the
+remaining contraction, (0, 1), is then executed.
a = np.random.rand(2, 2)
+b = np.random.rand(2, 5)
+c = np.random.rand(5, 2)
+path_info = opt_einsum.contract_path('ij,jk,kl->il', a, b, c)
+print(path_info[0])
+#> [(1, 2), (0, 1)]
+print(path_info[1])
+#> Complete contraction: ij,jk,kl->il
+#> Naive scaling: 4
+#> Optimized scaling: 3
+#> Naive FLOP count: 1.600e+02
+#> Optimized FLOP count: 5.600e+01
+#> Theoretical speedup: 2.857
+#> Largest intermediate: 4.000e+00 elements
+#> -------------------------------------------------------------------------
+#> scaling current remaining
+#> -------------------------------------------------------------------------
+#> 3 kl,jk->jl ij,jl->il
+#> 3 jl,ij->il il->il
+A more complex index transformation example.
+I = np.random.rand(10, 10, 10, 10)
+C = np.random.rand(10, 10)
+path_info = oe.contract_path('ea,fb,abcd,gc,hd->efgh', C, C, I, C, C)
+
+print(path_info[0])
+#> [(0, 2), (0, 3), (0, 2), (0, 1)]
+print(path_info[1])
+#> Complete contraction: ea,fb,abcd,gc,hd->efgh
+#> Naive scaling: 8
+#> Optimized scaling: 5
+#> Naive FLOP count: 8.000e+08
+#> Optimized FLOP count: 8.000e+05
+#> Theoretical speedup: 1000.000
+#> Largest intermediate: 1.000e+04 elements
+#> --------------------------------------------------------------------------
+#> scaling current remaining
+#> --------------------------------------------------------------------------
+#> 5 abcd,ea->bcde fb,gc,hd,bcde->efgh
+#> 5 bcde,fb->cdef gc,hd,cdef->efgh
+#> 5 cdef,gc->defg hd,defg->efgh
+#> 5 defg,hd->efgh efgh->efgh
+opt_einsum/contract.py173 +174 +175 +176 +177 +178 +179 +180 +181 +182 +183 +184 +185 +186 +187 +188 +189 +190 +191 +192 +193 +194 +195 +196 +197 +198 +199 +200 +201 +202 +203 +204 +205 +206 +207 +208 +209 +210 +211 +212 +213 +214 +215 +216 +217 +218 +219 +220 +221 +222 +223 +224 +225 +226 +227 +228 +229 +230 +231 +232 +233 +234 +235 +236 +237 +238 +239 +240 +241 +242 +243 +244 +245 +246 +247 +248 +249 +250 +251 +252 +253 +254 +255 +256 +257 +258 +259 +260 +261 +262 +263 +264 +265 +266 +267 +268 +269 +270 +271 +272 +273 +274 +275 +276 +277 +278 +279 +280 +281 +282 +283 +284 +285 +286 +287 +288 +289 +290 +291 +292 +293 +294 +295 +296 +297 +298 +299 +300 +301 +302 +303 +304 +305 +306 +307 +308 +309 +310 +311 +312 +313 +314 +315 +316 +317 +318 +319 +320 +321 +322 +323 +324 +325 +326 +327 +328 +329 +330 +331 +332 +333 +334 +335 +336 +337 +338 +339 +340 +341 +342 +343 +344 +345 +346 +347 +348 +349 +350 +351 +352 +353 +354 +355 +356 +357 +358 +359 +360 +361 +362 +363 +364 +365 +366 +367 +368 +369 +370 +371 +372 +373 +374 +375 +376 +377 +378 +379 +380 +381 +382 +383 +384 +385 +386 +387 +388 +389 +390 +391 +392 +393 +394 +395 +396 +397 +398 +399 +400 +401 +402 +403 +404 +405 +406 +407 +408 +409 +410 +411 +412 +413 +414 +415 +416 +417 +418 +419 +420 +421 +422 +423 +424 +425 +426 +427 +428 +429 +430 | |
opt_einsum.contract_expressionGenerate a reusable expression for a given contraction with +specific shapes, which can, for example, be cached.
+Parameters:
+subscripts: Specifies the subscripts for summation.
+shapes: Shapes of the arrays to optimize the contraction for.
+constants: The indices of any constant arguments in `shapes`, in which case the
+ actual array should be supplied at that position rather than just a
+ shape. If these are specified, then constant parts of the contraction
+ between calls will be reused. Additionally, if a GPU-enabled backend is
+ used for example, then the constant tensors will be kept on the GPU,
+ minimizing transfers.
+kwargs: Passed on to `contract_path` or `einsum`. See `contract`.
+Returns:
+| Type | +Description | +
|---|---|
+ ContractExpression
+ |
+
+
+
+ Callable with signature |
+
The out keyword argument should be supplied to the generated expression
+rather than this function.
+The backend keyword argument should also be supplied to the generated
+expression. If numpy arrays are supplied, if possible they will be
+converted to and back from the correct backend array type.
+The generated expression will work with any arrays which have
+the same rank (number of dimensions) as the original shapes, however, if
+the actual sizes are different, the expression may no longer be optimal.
+Constant operations will be computed upon the first call with a particular
+backend, then subsequently reused.
Examples: +Basic usage:
+expr = contract_expression("ab,bc->ac", (3, 4), (4, 5))
+a, b = np.random.rand(3, 4), np.random.rand(4, 5)
+c = expr(a, b)
+np.allclose(c, a @ b)
+#> True
+Supply a as a constant:
expr = contract_expression("ab,bc->ac", a, (4, 5), constants=[0])
+expr
+#> <ContractExpression('[ab],bc->ac', constants=[0])>
+
+c = expr(b)
+np.allclose(c, a @ b)
+#> True
+opt_einsum/contract.py984 + 985 + 986 + 987 + 988 + 989 + 990 + 991 + 992 + 993 + 994 + 995 + 996 + 997 + 998 + 999 +1000 +1001 +1002 +1003 +1004 +1005 +1006 +1007 +1008 +1009 +1010 +1011 +1012 +1013 +1014 +1015 +1016 +1017 +1018 +1019 +1020 +1021 +1022 +1023 +1024 +1025 +1026 +1027 +1028 +1029 +1030 +1031 +1032 +1033 +1034 +1035 +1036 +1037 +1038 +1039 +1040 +1041 +1042 +1043 +1044 +1045 +1046 +1047 +1048 +1049 +1050 +1051 +1052 +1053 +1054 +1055 +1056 +1057 +1058 +1059 +1060 +1061 +1062 +1063 +1064 +1065 +1066 +1067 +1068 +1069 +1070 +1071 | |
opt_einsum.contract.ContractExpressionHelper class for storing an explicit contraction_list which can
+then be repeatedly called solely with the array arguments.
opt_einsum/contract.py761 +762 +763 +764 +765 +766 +767 +768 +769 +770 +771 +772 +773 +774 +775 +776 +777 +778 +779 +780 +781 +782 +783 +784 +785 +786 +787 +788 +789 +790 +791 +792 +793 +794 +795 +796 +797 +798 +799 +800 +801 +802 +803 +804 +805 +806 +807 +808 +809 +810 +811 +812 +813 +814 +815 +816 +817 +818 +819 +820 +821 +822 +823 +824 +825 +826 +827 +828 +829 +830 +831 +832 +833 +834 +835 +836 +837 +838 +839 +840 +841 +842 +843 +844 +845 +846 +847 +848 +849 +850 +851 +852 +853 +854 +855 +856 +857 +858 +859 +860 +861 +862 +863 +864 +865 +866 +867 +868 +869 +870 +871 +872 +873 +874 +875 +876 +877 +878 +879 +880 +881 +882 +883 +884 +885 +886 +887 +888 +889 +890 +891 +892 +893 +894 +895 +896 +897 +898 +899 +900 +901 +902 +903 +904 +905 +906 +907 +908 +909 +910 +911 +912 +913 +914 +915 +916 +917 +918 +919 +920 +921 +922 +923 +924 +925 +926 +927 +928 +929 +930 +931 +932 +933 +934 +935 +936 +937 +938 +939 +940 +941 +942 +943 +944 +945 +946 +947 +948 | |
__call__(*arrays, out=None, backend='auto', evaluate_constants=False)
+
+Evaluate this expression with a set of arrays.
+ + +Parameters:
+| Name | +Type | +Description | +Default | +
|---|---|---|---|
arrays |
+
+ ArrayType
+ |
+
+
+
+ The arrays to supply as input to the expression. + |
+
+ ()
+ |
+
out |
+
+ Union[None, ArrayType]
+ |
+
+
+
+ If specified, output the result into this array. + |
+
+ None
+ |
+
backend |
+
+ str
+ |
+
+
+
+ Perform the contraction with this backend library. If numpy arrays +are supplied then try to convert them to and from the correct +backend array type. + |
+
+ 'auto'
+ |
+
evaluate_constants |
+
+ bool
+ |
+
+
+
+ Pre-evaluates constants with the appropriate backend. + |
+
+ False
+ |
+
Returns:
+| Type | +Description | +
|---|---|
+ ArrayType
+ |
+
+
+
+ The contracted result. + |
+
opt_einsum/contract.py879 +880 +881 +882 +883 +884 +885 +886 +887 +888 +889 +890 +891 +892 +893 +894 +895 +896 +897 +898 +899 +900 +901 +902 +903 +904 +905 +906 +907 +908 +909 +910 +911 +912 +913 +914 +915 +916 +917 +918 +919 +920 +921 +922 +923 +924 +925 +926 +927 +928 +929 +930 +931 +932 | |
evaluate_constants(backend='auto')
+
+Convert any constant operands to the correct backend form, and
+perform as many contractions as possible to create a new list of
+operands, stored in self._evaluated_constants[backend]. This also
+makes sure self.contraction_list only contains the remaining,
+non-const operations.
opt_einsum/contract.py792 +793 +794 +795 +796 +797 +798 +799 +800 +801 +802 +803 +804 +805 +806 +807 +808 +809 +810 | |
opt_einsum.contract.PathInfoA printable object to contain information about a contraction path.
+ +opt_einsum/contract.py35 + 36 + 37 + 38 + 39 + 40 + 41 + 42 + 43 + 44 + 45 + 46 + 47 + 48 + 49 + 50 + 51 + 52 + 53 + 54 + 55 + 56 + 57 + 58 + 59 + 60 + 61 + 62 + 63 + 64 + 65 + 66 + 67 + 68 + 69 + 70 + 71 + 72 + 73 + 74 + 75 + 76 + 77 + 78 + 79 + 80 + 81 + 82 + 83 + 84 + 85 + 86 + 87 + 88 + 89 + 90 + 91 + 92 + 93 + 94 + 95 + 96 + 97 + 98 + 99 +100 +101 +102 | |
opt_einsum.get_symbolGet the symbol corresponding to int i - runs through the usual 52
+letters before resorting to unicode characters, starting at chr(192) and skipping surrogates.
Examples:
+get_symbol(2)
+#> 'c'
+
+get_symbol(200)
+#> 'Å”'
+
+get_symbol(20000)
+#> '京'
+opt_einsum/parser.py57 +58 +59 +60 +61 +62 +63 +64 +65 +66 +67 +68 +69 +70 +71 +72 +73 +74 +75 +76 +77 +78 +79 +80 | |
opt_einsum.shared_intermediatesContext in which contract intermediate results are shared.
+Note that intermediate computations will not be garbage collected until +1. this context exits, and +2. the yielded cache is garbage collected (if it was captured).
+Parameters:
+Returns:
+opt_einsum/sharing.py54 +55 +56 +57 +58 +59 +60 +61 +62 +63 +64 +65 +66 +67 +68 +69 +70 +71 +72 +73 +74 +75 +76 +77 +78 +79 +80 +81 | |
opt_einsum.paths.optimalComputes all possible pair contractions in a depth-first recursive manner,
+sieving results based on memory_limit and the best path found so far.
Parameters:
+| Name | +Type | +Description | +Default | +
|---|---|---|---|
inputs |
+
+ List[ArrayIndexType]
+ |
+
+
+
+ List of sets that represent the lhs side of the einsum subscript. + |
+ + required + | +
output |
+
+ ArrayIndexType
+ |
+
+
+
+ Set that represents the rhs side of the overall einsum subscript. + |
+ + required + | +
size_dict |
+
+ Dict[str, int]
+ |
+
+
+
+ Dictionary of index sizes. + |
+ + required + | +
memory_limit |
+
+ Optional[int]
+ |
+
+
+
+ The maximum number of elements in a temporary array. + |
+
+ None
+ |
+
Returns:
+| Name | Type | +Description | +
|---|---|---|
path |
+ PathType
+ |
+
+
+
+ The optimal contraction order within the memory limit constraint. + |
+
Examples: +
isets = [set('abd'), set('ac'), set('bdc')]
+oset = set('')
+idx_sizes = {'a': 1, 'b':2, 'c':3, 'd':4}
+optimal(isets, oset, idx_sizes, 5000)
+#> [(0, 2), (0, 1)]
+opt_einsum/paths.py201 +202 +203 +204 +205 +206 +207 +208 +209 +210 +211 +212 +213 +214 +215 +216 +217 +218 +219 +220 +221 +222 +223 +224 +225 +226 +227 +228 +229 +230 +231 +232 +233 +234 +235 +236 +237 +238 +239 +240 +241 +242 +243 +244 +245 +246 +247 +248 +249 +250 +251 +252 +253 +254 +255 +256 +257 +258 +259 +260 +261 +262 +263 +264 +265 +266 +267 +268 +269 +270 +271 +272 +273 +274 +275 +276 +277 +278 +279 +280 +281 +282 +283 | |
opt_einsum.paths.greedyFinds the path by a three stage algorithm:
+removed_sizeThis algorithm scales quadratically with respect to the +maximum number of elements sharing a common dim.
+ + +Parameters:
+| Name | +Type | +Description | +Default | +
|---|---|---|---|
inputs |
+
+ List[ArrayIndexType]
+ |
+
+
+
+ List of sets that represent the lhs side of the einsum subscript + |
+ + required + | +
output |
+
+ ArrayIndexType
+ |
+
+
+
+ Set that represents the rhs side of the overall einsum subscript + |
+ + required + | +
size_dict |
+
+ Dict[str, int]
+ |
+
+
+
+ Dictionary of index sizes + |
+ + required + | +
memory_limit |
+
+ Optional[int]
+ |
+
+
+
+ The maximum number of elements in a temporary array + |
+
+ None
+ |
+
choose_fn |
+
+ Any
+ |
+
+
+
+ A function that chooses which contraction to perform from the queue + |
+
+ None
+ |
+
cost_fn |
+
+ str
+ |
+
+
+
+ A function that assigns a potential contraction a cost. + |
+
+ 'memory-removed'
+ |
+
Returns:
+| Name | Type | +Description | +
|---|---|---|
path |
+ PathType
+ |
+
+
+
+ The contraction order (a list of tuples of ints). + |
+
Examples:
+isets = [set('abd'), set('ac'), set('bdc')]
+oset = set('')
+idx_sizes = {'a': 1, 'b':2, 'c':3, 'd':4}
+greedy(isets, oset, idx_sizes)
+#> [(0, 2), (0, 1)]
+opt_einsum/paths.py728 +729 +730 +731 +732 +733 +734 +735 +736 +737 +738 +739 +740 +741 +742 +743 +744 +745 +746 +747 +748 +749 +750 +751 +752 +753 +754 +755 +756 +757 +758 +759 +760 +761 +762 +763 +764 +765 +766 +767 +768 +769 | |
opt_einsum.paths.branchopt_einsum/paths.py497 +498 +499 +500 +501 +502 +503 +504 +505 | |
opt_einsum.paths.PathOptimizerBase class for different path optimizers to inherit from.
+Subclassed optimizers should define a call method with signature:
+def __call__(self, inputs: List[ArrayIndexType], output: ArrayIndexType, size_dict: dict[str, int], memory_limit: int | None = None) -> list[tuple[int, ...]]:
+ \"\"\"
+ Parameters:
+ inputs: The indices of each input array.
+ outputs: The output indices
+ size_dict: The size of each index
+ memory_limit: If given, the maximum allowed memory.
+ \"\"\"
+ # ... compute path here ...
+ return path
+where path is a list of int-tuples specifying a contraction order.
opt_einsum/paths.py32 +33 +34 +35 +36 +37 +38 +39 +40 +41 +42 +43 +44 +45 +46 +47 +48 +49 +50 +51 +52 +53 +54 +55 +56 +57 +58 +59 +60 +61 +62 +63 +64 +65 +66 +67 +68 +69 +70 +71 +72 +73 +74 +75 +76 +77 +78 +79 | |
opt_einsum.paths.BranchBound
+ Bases: PathOptimizer
opt_einsum/paths.py330 +331 +332 +333 +334 +335 +336 +337 +338 +339 +340 +341 +342 +343 +344 +345 +346 +347 +348 +349 +350 +351 +352 +353 +354 +355 +356 +357 +358 +359 +360 +361 +362 +363 +364 +365 +366 +367 +368 +369 +370 +371 +372 +373 +374 +375 +376 +377 +378 +379 +380 +381 +382 +383 +384 +385 +386 +387 +388 +389 +390 +391 +392 +393 +394 +395 +396 +397 +398 +399 +400 +401 +402 +403 +404 +405 +406 +407 +408 +409 +410 +411 +412 +413 +414 +415 +416 +417 +418 +419 +420 +421 +422 +423 +424 +425 +426 +427 +428 +429 +430 +431 +432 +433 +434 +435 +436 +437 +438 +439 +440 +441 +442 +443 +444 +445 +446 +447 +448 +449 +450 +451 +452 +453 +454 +455 +456 +457 +458 +459 +460 +461 +462 +463 +464 +465 +466 +467 +468 +469 +470 +471 +472 +473 +474 +475 +476 +477 +478 +479 +480 +481 +482 +483 +484 +485 +486 +487 +488 +489 +490 +491 +492 +493 +494 | |
__call__(inputs_, output_, size_dict, memory_limit=None)
+
+Parameters:
+| Name | +Type | +Description | +Default | +
|---|---|---|---|
inputs_ |
+
+ List[ArrayIndexType]
+ |
+
+
+
+ List of sets that represent the lhs side of the einsum subscript + |
+ + required + | +
output_ |
+
+ ArrayIndexType
+ |
+
+
+
+ Set that represents the rhs side of the overall einsum subscript + |
+ + required + | +
size_dict |
+
+ Dict[str, int]
+ |
+
+
+
+ Dictionary of index sizes + |
+ + required + | +
memory_limit |
+
+ Optional[int]
+ |
+
+
+
+ The maximum number of elements in a temporary array. + |
+
+ None
+ |
+
Returns:
+| Name | Type | +Description | +
|---|---|---|
path |
+ PathType
+ |
+
+
+
+ The contraction order within the memory limit constraint. + |
+
Examples: +```python +isets = [set('abd'), set('ac'), set('bdc')] +oset = set('') +idx_sizes = {'a': 1, 'b':2, 'c':3, 'd':4} +optimal(isets, oset, idx_sizes, 5000)
+opt_einsum/paths.py373 +374 +375 +376 +377 +378 +379 +380 +381 +382 +383 +384 +385 +386 +387 +388 +389 +390 +391 +392 +393 +394 +395 +396 +397 +398 +399 +400 +401 +402 +403 +404 +405 +406 +407 +408 +409 +410 +411 +412 +413 +414 +415 +416 +417 +418 +419 +420 +421 +422 +423 +424 +425 +426 +427 +428 +429 +430 +431 +432 +433 +434 +435 +436 +437 +438 +439 +440 +441 +442 +443 +444 +445 +446 +447 +448 +449 +450 +451 +452 +453 +454 +455 +456 +457 +458 +459 +460 +461 +462 +463 +464 +465 +466 +467 +468 +469 +470 +471 +472 +473 +474 +475 +476 +477 +478 +479 +480 +481 +482 +483 +484 +485 +486 +487 +488 +489 +490 +491 +492 +493 +494 | |
__init__(nbranch=None, cutoff_flops_factor=4, minimize='flops', cost_fn='memory-removed')
+
+Explores possible pair contractions in a depth-first recursive manner like
+the optimal approach, but with extra heuristic early pruning of branches
+as well sieving by memory_limit and the best path found so far.
Parameters:
+| Name | +Type | +Description | +Default | +
|---|---|---|---|
nbranch |
+
+ Optional[int]
+ |
+
+
+
+ How many branches to explore at each contraction step. If None, explore +all possible branches. If an integer, branch into this many paths at +each step. Defaults to None. + |
+
+ None
+ |
+
cutoff_flops_factor |
+
+ int
+ |
+
+
+
+ If at any point, a path is doing this much worse than the best path +found so far was, terminate it. The larger this is made, the more paths +will be fully explored and the slower the algorithm. Defaults to 4. + |
+
+ 4
+ |
+
minimize |
+
+ str
+ |
+
+
+
+ Whether to optimize the path with regard primarily to the total +estimated flop-count, or the size of the largest intermediate. The +option not chosen will still be used as a secondary criterion. + |
+
+ 'flops'
+ |
+
cost_fn |
+
+ str
+ |
+
+
+
+ A function that returns a heuristic 'cost' of a potential contraction
+with which to sort candidates. Should have signature
+ |
+
+ 'memory-removed'
+ |
+
opt_einsum/paths.py331 +332 +333 +334 +335 +336 +337 +338 +339 +340 +341 +342 +343 +344 +345 +346 +347 +348 +349 +350 +351 +352 +353 +354 +355 +356 +357 +358 +359 +360 +361 +362 +363 +364 +365 +366 +367 | |
opt_einsum.path_random.RandomOptimizer
+ Bases: PathOptimizer
Base class for running any random path finder that benefits
+from repeated calling, possibly in a parallel fashion. Custom random
+optimizers should subclass this, and the setup method should be
+implemented with the following signature:
def setup(self, inputs, output, size_dict):
+ # custom preparation here ...
+ return trial_fn, trial_args
+Where trial_fn itself should have the signature::
def trial_fn(r, *trial_args):
+ # custom computation of path here
+ return ssa_path, cost, size
+Where r is the run number and could for example be used to seed a
+random number generator. See RandomGreedy for an example.
Parameters:
+| Name | +Type | +Description | +Default | +
|---|---|---|---|
max_repeats |
+
+ int
+ |
+
+
+
+ The maximum number of repeat trials to have. + |
+
+ 32
+ |
+
max_time |
+
+ Optional[float]
+ |
+
+
+
+ The maximum amount of time to run the algorithm for. + |
+
+ None
+ |
+
minimize |
+
+ str
+ |
+
+
+
+ Whether to favour paths that minimize the total estimated flop-count or +the size of the largest intermediate created. + |
+
+ 'flops'
+ |
+
parallel |
+
+ Union[bool, Decimal, int]
+ |
+
+
+
+ Whether to parallelize the random trials, by default |
+
+ False
+ |
+
pre_dispatch |
+
+ int
+ |
+
+
+
+ If running in parallel, how many jobs to pre-dispatch so as to avoid +submitting all jobs at once. Should also be more than twice the number +of workers to avoid under-subscription. Default: 128. + |
+
+ 128
+ |
+
Attributes:
+| Name | +Type | +Description | +
|---|---|---|
path |
+
+ PathType
+ |
+
+
+
+ The best path found so far. + |
+
costs |
+
+ List[int]
+ |
+
+
+
+ The list of each trial's costs found so far. + |
+
sizes |
+
+ List[int]
+ |
+
+
+
+ The list of each trial's largest intermediate size so far. + |
+
opt_einsum/path_random.py19 + 20 + 21 + 22 + 23 + 24 + 25 + 26 + 27 + 28 + 29 + 30 + 31 + 32 + 33 + 34 + 35 + 36 + 37 + 38 + 39 + 40 + 41 + 42 + 43 + 44 + 45 + 46 + 47 + 48 + 49 + 50 + 51 + 52 + 53 + 54 + 55 + 56 + 57 + 58 + 59 + 60 + 61 + 62 + 63 + 64 + 65 + 66 + 67 + 68 + 69 + 70 + 71 + 72 + 73 + 74 + 75 + 76 + 77 + 78 + 79 + 80 + 81 + 82 + 83 + 84 + 85 + 86 + 87 + 88 + 89 + 90 + 91 + 92 + 93 + 94 + 95 + 96 + 97 + 98 + 99 +100 +101 +102 +103 +104 +105 +106 +107 +108 +109 +110 +111 +112 +113 +114 +115 +116 +117 +118 +119 +120 +121 +122 +123 +124 +125 +126 +127 +128 +129 +130 +131 +132 +133 +134 +135 +136 +137 +138 +139 +140 +141 +142 +143 +144 +145 +146 +147 +148 +149 +150 +151 +152 +153 +154 +155 +156 +157 +158 +159 +160 +161 +162 +163 +164 +165 +166 +167 +168 +169 +170 +171 +172 +173 +174 +175 +176 +177 +178 +179 +180 +181 +182 +183 +184 +185 +186 +187 +188 +189 +190 +191 +192 +193 +194 +195 +196 +197 +198 +199 +200 +201 +202 +203 +204 +205 +206 +207 +208 | |
path: PathType
+
+
+ property
+
+
+The best path found so far.
+opt_einsum.path_random.RandomGreedy
+ Bases: RandomOptimizer
opt_einsum/path_random.py326 +327 +328 +329 +330 +331 +332 +333 +334 +335 +336 +337 +338 +339 +340 +341 +342 +343 +344 +345 +346 +347 +348 +349 +350 +351 +352 +353 +354 +355 +356 +357 +358 +359 +360 +361 +362 +363 +364 +365 +366 +367 +368 +369 +370 +371 +372 +373 +374 +375 +376 +377 +378 +379 +380 +381 +382 +383 | |
choose_fn: Any
+
+
+ property
+
+
+The function that chooses which contraction to take - make this a
+property so that temperature and nbranch etc. can be updated
+between runs.
__init__(cost_fn='memory-removed-jitter', temperature=1.0, rel_temperature=True, nbranch=8, **kwargs)
+
+cost_fn: A function that returns a heuristic 'cost' of a potential contraction
+ with which to sort candidates. Should have signature
+ cost_fn(size12, size1, size2, k12, k1, k2).
+temperature: When choosing a possible contraction, its relative probability will be
+ proportional to exp(-cost / temperature). Thus the larger
+ temperature is, the further random paths will stray from the normal
+ 'greedy' path. Conversely, if set to zero, only paths with exactly the
+ same cost as the best at each step will be explored.
+rel_temperature: Whether to normalize the temperature at each step to the scale of
+ the best cost. This is generally beneficial as the magnitude of costs
+ can vary significantly throughout a contraction. If False, the
+ algorithm will end up branching when the absolute cost is low, but
+ stick to the 'greedy' path when the cost is high - this can also be
+ beneficial.
+nbranch: How many potential paths to calculate probability for and choose from at each step.
+kwargs: Supplied to RandomOptimizer.
opt_einsum/path_random.py327 +328 +329 +330 +331 +332 +333 +334 +335 +336 +337 +338 +339 +340 +341 +342 +343 +344 +345 +346 +347 +348 +349 +350 +351 +352 +353 +354 +355 +356 +357 | |
opt_einsum.paths.DynamicProgramming
+ Bases: PathOptimizer
Finds the optimal path of pairwise contractions without intermediate outer +products based a dynamic programming approach presented in +Phys. Rev. E 90, 033315 (2014) (the corresponding preprint is publicly +available at https://arxiv.org/abs/1304.6112). This method is especially +well-suited in the area of tensor network states, where it usually +outperforms all the other optimization strategies.
+This algorithm shows exponential scaling with the number of inputs +in the worst case scenario (see example below). If the graph to be +contracted consists of disconnected subgraphs, the algorithm scales +linearly in the number of disconnected subgraphs and only exponentially +with the number of inputs per subgraph.
+ + +Parameters:
+| Name | +Type | +Description | +Default | +
|---|---|---|---|
minimize |
+
+ str
+ |
+
+
+
+ What to minimize:
+- 'flops' - minimize the number of flops
+- 'size' - minimize the size of the largest intermediate
+- 'write' - minimize the size of all intermediate tensors
+- 'combo' - minimize |
+
+ 'flops'
+ |
+
cost_cap |
+
+ Union[bool, int]
+ |
+
+
+
+ How to implement cost-capping: +- True - iteratively increase the cost-cap +- False - implement no cost-cap at all +- int - use explicit cost cap + |
+
+ True
+ |
+
search_outer |
+
+ bool
+ |
+
+
+
+ In rare circumstances the optimal contraction may involve an outer +product, this option allows searching such contractions but may well +slow down the path finding considerably on all but very small graphs. + |
+
+ False
+ |
+
opt_einsum/paths.py1103 +1104 +1105 +1106 +1107 +1108 +1109 +1110 +1111 +1112 +1113 +1114 +1115 +1116 +1117 +1118 +1119 +1120 +1121 +1122 +1123 +1124 +1125 +1126 +1127 +1128 +1129 +1130 +1131 +1132 +1133 +1134 +1135 +1136 +1137 +1138 +1139 +1140 +1141 +1142 +1143 +1144 +1145 +1146 +1147 +1148 +1149 +1150 +1151 +1152 +1153 +1154 +1155 +1156 +1157 +1158 +1159 +1160 +1161 +1162 +1163 +1164 +1165 +1166 +1167 +1168 +1169 +1170 +1171 +1172 +1173 +1174 +1175 +1176 +1177 +1178 +1179 +1180 +1181 +1182 +1183 +1184 +1185 +1186 +1187 +1188 +1189 +1190 +1191 +1192 +1193 +1194 +1195 +1196 +1197 +1198 +1199 +1200 +1201 +1202 +1203 +1204 +1205 +1206 +1207 +1208 +1209 +1210 +1211 +1212 +1213 +1214 +1215 +1216 +1217 +1218 +1219 +1220 +1221 +1222 +1223 +1224 +1225 +1226 +1227 +1228 +1229 +1230 +1231 +1232 +1233 +1234 +1235 +1236 +1237 +1238 +1239 +1240 +1241 +1242 +1243 +1244 +1245 +1246 +1247 +1248 +1249 +1250 +1251 +1252 +1253 +1254 +1255 +1256 +1257 +1258 +1259 +1260 +1261 +1262 +1263 +1264 +1265 +1266 +1267 +1268 +1269 +1270 +1271 +1272 +1273 +1274 +1275 +1276 +1277 +1278 +1279 +1280 +1281 +1282 +1283 +1284 +1285 +1286 +1287 +1288 +1289 +1290 +1291 +1292 +1293 +1294 +1295 +1296 +1297 +1298 +1299 +1300 +1301 | |
__call__(inputs_, output_, size_dict_, memory_limit_=None)
+
+Parameters:
+| Name | +Type | +Description | +Default | +
|---|---|---|---|
inputs_ |
+
+ List[ArrayIndexType]
+ |
+
+
+
+ List of sets that represent the lhs side of the einsum subscript + |
+ + required + | +
output_ |
+
+ ArrayIndexType
+ |
+
+
+
+ Set that represents the rhs side of the overall einsum subscript + |
+ + required + | +
size_dict_ |
+
+ Dict[str, int]
+ |
+
+
+
+ Dictionary of index sizes + |
+ + required + | +
memory_limit_ |
+
+ Optional[int]
+ |
+
+
+
+ The maximum number of elements in a temporary array. + |
+
+ None
+ |
+
Returns:
+| Name | Type | +Description | +
|---|---|---|
path |
+ PathType
+ |
+
+
+
+ The contraction order (a list of tuples of ints). + |
+
Examples:
+n_in = 3 # exponential scaling
+n_out = 2 # linear scaling
+s = dict()
+i_all = []
+for _ in range(n_out):
+ i = [set() for _ in range(n_in)]
+ for j in range(n_in):
+ for k in range(j+1, n_in):
+ c = oe.get_symbol(len(s))
+ i[j].add(c)
+ i[k].add(c)
+ s[c] = 2
+ i_all.extend(i)
+o = DynamicProgramming()
+o(i_all, set(), s)
+#> [(1, 2), (0, 4), (1, 2), (0, 2), (0, 1)]
+opt_einsum/paths.py1143 +1144 +1145 +1146 +1147 +1148 +1149 +1150 +1151 +1152 +1153 +1154 +1155 +1156 +1157 +1158 +1159 +1160 +1161 +1162 +1163 +1164 +1165 +1166 +1167 +1168 +1169 +1170 +1171 +1172 +1173 +1174 +1175 +1176 +1177 +1178 +1179 +1180 +1181 +1182 +1183 +1184 +1185 +1186 +1187 +1188 +1189 +1190 +1191 +1192 +1193 +1194 +1195 +1196 +1197 +1198 +1199 +1200 +1201 +1202 +1203 +1204 +1205 +1206 +1207 +1208 +1209 +1210 +1211 +1212 +1213 +1214 +1215 +1216 +1217 +1218 +1219 +1220 +1221 +1222 +1223 +1224 +1225 +1226 +1227 +1228 +1229 +1230 +1231 +1232 +1233 +1234 +1235 +1236 +1237 +1238 +1239 +1240 +1241 +1242 +1243 +1244 +1245 +1246 +1247 +1248 +1249 +1250 +1251 +1252 +1253 +1254 +1255 +1256 +1257 +1258 +1259 +1260 +1261 +1262 +1263 +1264 +1265 +1266 +1267 +1268 +1269 +1270 +1271 +1272 +1273 +1274 +1275 +1276 +1277 +1278 +1279 +1280 +1281 +1282 +1283 +1284 +1285 +1286 +1287 +1288 +1289 +1290 +1291 +1292 +1293 +1294 +1295 +1296 +1297 +1298 +1299 +1300 +1301 | |