Expanded support for imports / namespaces

[iancero]

Hi Kartik,

As my models develop, it often helps to try to organize my code in different files. When agents are heterogeneous (e.g., different available actions, reward functions), one of the intuitive options is to divide them into separate python files (e.g., alice.py, bob.py), then import the necessary classes and functions from those files into one another.

When I use memo, this works with from A import a, but not import A. Is there a way to allow memo to use both syntaxes? I've got a reproducible example similar to the Memonomicon below:

We start with a simple model of Alice in alice.py with no frills (no errors).

from memo import memo
from enum import IntEnum

class Cards(IntEnum):
    JACK = 0
    QUEEN = 1
    KING = 2

@memo
def chooses_card[c: Cards](beta):
    alice: chooses(c in Cards, wpp=exp(beta * c))
    return E[alice.c == c]

Then, we attempt to import alice into bob.py in two ways. The first works, the second does not.

from memo import memo

# import for Version 1
from alice import Cards as AlicesCards
from alice import chooses_card as alice_chooses_card

# import for Version 2
import alice

# Version 1 - works
@memo
def chooses_card_v1[c: AlicesCards]():
    bob: thinks[partner : chooses(c in AlicesCards, wpp=alice_chooses_card[c](3.0))]
    bob: chooses(c in AlicesCards, wpp=E[~(c == partner.c)])
    return E[bob.c == c]
    # No errors, works as expected


# Version 2 - fails
@memo
def chooses_card_v2[c: alice.Cards]():
    bob: thinks[partner : chooses(c in alice.Cards, wpp=alice.chooses_card[c](3.0))]
    bob: chooses(c in alice.Cards, wpp=E[~(c == partner.c)])
    return E[bob.c == c]
    # AssertionError
    # line 633, in parse_memo
    # assert isinstance(tp.bound, ast.Name)

Is there a way to allow memo to work with version 2 also? Although neither import style is perfect, and support for both would be great, I think version two will often lead to clearer and more readable code. As models get larger and more complex, I think the benefit of version two also grows. Thoughts?

Thanks for your help,
Ian

[kach]

Hi Ian! Thanks for asking about this. You are absolutely right that importing namespaced packages helps keep code clean and maintainable. I haven't actually met anyone who has a model complicated enough to need this, but I can certainly imagine it.

The reason I am reluctant to support dot-chained references in memo is that dot is already used as shorthand for an agent's choice (e.g. partner.c). I think overloading dots to also mean "class/module name lookup" will lead to confusion among users. In fact, in general, I have been trying to find ways to emphasize that memo ≠ Python: even though memo lives inside Python and uses superficially-similar syntax, it is simply a different language, and things that work in Python are not guaranteed to work in memo.

Another issue is that types are currently compared by name, not by value. For example, if you say "alice: chooses(x in X, wpp=1)", and "bob: chooses(y in Y, wpp=1)," then alice.x and bob.y have different types, even if Y is defined via Y = X in Python. This is desirable to give different semantically-meaningful names to structually-equivalent domains (e.g. think about how many different domains the set {0, 1} could represent). I think this aspect of memo would interact poorly with the possibility of aliasing caused by imports. In general, how would we decide at compile-time whether the types alice.X and bob.Y should be treated the same?

With all of that in mind, I'm wondering if your current problem could be addressed by prefixing your variable names directly in the source and then doing a * import?

from memo import memo
from enum import IntEnum

class AliceCards(IntEnum):
    JACK = 0
    QUEEN = 1
    KING = 2

@memo
def alice_chooses_card[c: AliceCards](beta):
    alice: chooses(c in AliceCards, wpp=exp(beta * c))
    return E[alice.c == c]

If that is too much of a hassle, let me know, and I'll think more about what can be done here…