Factor out auto path optimizer + allow shapes in contract_path #84
Conversation
jcmgray
changed the title
opt_einsum/paths.py
Outdated
| def add_path_fn(name, fn): | ||
| """Add path finding function ``fn`` as an option with ``name``. | ||
| """ | ||
| _PATH_OPTIONS[name] = fn |
There was a problem hiding this comment.
Can we 1) lowercase name and 2) make sure the key does not already exist?
opt_einsum/path_random.py
Outdated
| @@ -372,3 +372,6 @@ def random_greedy(inputs, output, idx_dict, memory_limit=None, **optimizer_kwarg | |||
| """ | |||
| optimizer = RandomGreedy(**optimizer_kwargs) | |||
| return optimizer(inputs, output, idx_dict, memory_limit) | |||
|
|
|||
|
|
|||
| paths.add_path_fn('random-greedy', random_greedy) | |||
There was a problem hiding this comment.
I do like the registration, but the potential circular imports worry me a bit. What do you think of imports the path finders into paths.py and register them there?
There was a problem hiding this comment.
So the registration was my way of getting round the circular import problem since path_random depends on paths quite a bit. However switching to from . import paths in this file might have sorted that already. Alternatively, could perform the registration in the top __init__?
There was a problem hiding this comment.
The from . imports paths did prevent the circular depends in this particular case, but is quite vulnerable to problems in the future. Perhaps we could split off registration into another file or as you said, add it to the init.
opt_einsum/paths.py
Outdated
| } | ||
|
|
||
|
|
||
| def add_path_fn(name, fn): |
There was a problem hiding this comment.
Could we call this register_path_fn, would make the language slightly clearer I believe.
|
I've made those changes @dgasmith and moved the random-greedy registration into the top |
opt_einsum/paths.py
Outdated
| } | ||
|
|
||
|
|
||
| def register_path_fn(name, fn, overwrite=False): |
There was a problem hiding this comment.
I would vote against an overwrite option here. This could give the ability for folks to override the definitions of paths in our code. Choosing another name is equally easy.
There was a problem hiding this comment.
Sure thing, I was thinking here one might want to update a function, but this can still be explicitly done by deleting the key first.
Description
A couple of minor-ish path tweaks:
'auto'path logic to an actual path function (closes FR: separate 'auto' strategy out of contract_path #82) cc @fritzocontract_pathto accept shapes with the optionshapes=True(rather than requiring actual array-like objects as now)The path options and auto path choosing based on size I've implemented as dictionaries, which as well as for efficiencies sake potentially allows them to be customized in the future. E.g. users adding their own named optimizers options or setting custom 'auto' strategies.
Status