Sharing IntermediatesΒΆ
If you want to compute multiple similar contractions with common terms, you can embed them in a shared_intermediates()
context. Computations of subexpressions in this context will be memoized, and will be garbage collected when the contexts exits.
For example, suppose we want to compute marginals at each point in a factor chain:
inputs = 'ab,bc,cd,de,ef'
factors = [np.random.rand(1000, 1000) for _ in range(5)]
%%timeit
marginals = {output: contract('{}->{}'.format(inputs, output), *factors)
for output in 'abcdef'}
1 loop, best of 3: 5.82 s per loop
To share this computation, we can perform all contractions in a shared context:
%%timeit
with shared_intermediates():
marginals = {output: contract('{}->{}'.format(inputs, output), *factors)
for output in 'abcdef'}
1 loop, best of 3: 1.55 s per loop
If it is difficult to fit your code into a context, you can instead save the sharing cache for later reuse.
with shared_intermediates() as cache: # create a cache
pass
marginals = {}
for output in 'abcdef':
with shared_intermediates(cache): # reuse a common cache
marginals[output] = contract('{}->{}'.format(inputs, output), *factors)
del cache # garbage collect intermediates
Note that sharing contexts can be nested, so it is safe to to use shared_intermediates()
in library code without leaking intermediates into user caches.
Note
By default a cache is thread safe, to share intermediates between threads explicitly pass the same cache to each thread.