I wrote a tool to help converting Einstein notation assignment directly to numpy.einsum. It looks like:
import numpy as np
from easy_ast.tensor_contract import TensorContract
b = np.array([1, 2])
c = np.array([1, 2])
@TensorContract().exec
def _(i, j):
a[i, j] = -b[i] * c[j]
print(a)
Where TensorContract().exec get the source of function and do ast transformer and exec code immediately.
However, I pass dummy index manaully to my transformer currently. If I could get the information about dummy index directly from symbol without assginment before(symbol i and j in this example), my interface could be better, since declaring dummy index is annoying.
Hmm. As far as Python is concerned, names like b and c are just as external as i and j - and even print. How do you know the difference, if they’re not arguments?
Define “assigned”. You mean when it’s used as a subscript on the left of the equals sign? There’s a distinction between the assignment in a = 1 (which assigns to the name a) and what you have with a[i, j] = ... (which does a subscript assignment to the object a). You could search the bytecode for the STORE_SUBSCR operation, but it might not be easy.
In my example, i and j are unassigned names, but b and c were assigned in b = np.array([1, 2]) and c = np.array([1, 2]), print can be found in global namespace. a is the Name with ctx=ast.Store() but not ast.Load(), so skip to check it.
Ahh, I see. So the module globals and the builtins are relevant, but anything that isn’t in one of those needs to become a local. That makes sense.
(By the way, since _ is a really awkward name, so I’m going to refer to it as func hereunder. Feel free to suggest a better name.)
So here’s what I’m thinking.
import builtins
class Magic(dict):
def __missing__(self, key):
if key in func.__globals__: return func.__globals__[key]
if hasattr(builtins, key): return getattr(builtins, key)
print("It's a local!")
exec("func()", func.__globals__, Magic())