Replace domains -> types, .inputs+.output -> type hints

[fritzo]

This issue proposes to replace funsor Domain objects with new subscripted types a la Python 3's typing module.

The goal is to make Funsor more Pythonic and to make funsors act more like Python functions (but with support for pointwise operations). In particular many of our design questions about multiple inputs or outputs could be resolved by "following Python". A simple example of more-Pythonic syntax is:

# decorator syntax to read .inputs and .outputs from a type annotation
@to_funsor
def f(x: Real[3], y: Real[3,3], t: Real) -> Real:
    ...

# explicitly specify type as Callable[..., ...]
g = to_funsor(lambda x: x.sum(), Callable[[Real[3]], Real)
    ...

h = f + g  # materialized view of lambda x, y, t: f(x, y, t) + g(x)

Interface

Because the Python 3 typing library is moving and not natively inspectable, we should implement a minimal interface for inspection:

• [x] #352 creating domain types, Bint[n], Reals[m,n]
• [x] #352 isinstance(-, Domain) (with Domain = type)
• [x] #352 issubclass(-, -) (by overriding __subclasscheck__)
• [x] #354 typing.get_type_hints() for funsors
• [ ] typing.get_origin()
• [ ] typing.get_args()

Tasks

• [x] #352 Refactor domains to be typing-compatible type objects: Bint[n], Real, Real[m,n], ...
• [x] #356 Generalize Bint to nontrivial shape #322 and add an Array supertype
• [x] #354 Update interfaces to make use of typing-compatible types (e.g. allow @function, @of_type to read type hints)
• [x] #357 Cleanup: replace all uses of reals() and bint() with Reals[] and Bint[]
• [x] #421 Cleanup: unify op-caching metaclasses
• [x] #422 Support dynamic op creation and registration in find_domain
• [x] #423 Add a Finitary term for lazy op application
• [x] #430 Promote LazyTuple to a subclass of Funsor
• [x] #431 Cleanup: automate invocation of Unary and Binary for ops applied to funsors
• [ ] Replace Function with Finitary (or merge Unary and Binary into Function as in the paper); and replace @function with a dynamic Op factory.
• [ ] Support limited (variables but no arithmetic) shape polymorphism in Array domains (#214)

  @function
  def matmul(x: Reals["i", "j"], y: Reals["j", "k"]) -> Reals["i", "k"]:
      return x @ y
  

  This will require a unification algorithm in find_domain() and substitution.
• [ ] Add type .__annotations__ to all Ops and refactor find_domain(). We may be able to read many of these from gufunc signatures.
• [ ] Add a mypy test stage and ensure the Funsor codebase typechecks.
• [ ] Register to_funsor to use either @function or @symbolic (requires first class Tuples, memoization)?
• [ ] Possibly promote Funsor types themselves to subtypes of Callable

[eb8680]

This definitely seems like the right way to go. Some design questions:

1. Could we replace find_domain with some standard Python type inference library? Are there existing special array types in NumPy or elsewhere we should use rather than making a custom Domain?
2. Would some version of this also get us shape polymorphism (as in #214) for free? If not, does it at least not put up an intrinsic barrier to shape variables (typing.Literal is currently very limited)?
3. How would this interact with surrounding typed Python code? That is, what is the type signature of to_funsor applied to a Callable as in the above examples?
4. Should we distinguish between a black-box funsor.function that creates a custom new funsor.Function term and a transparent funsor.to_funsor that attempts to extract a funsor expression by evaluating the function on appropriately typed funsor.Variables?

Possibly promote Funsor types themselves to subtypes of Callable

I was more thinking about making all funsor.Funsors subtypes of typing.Generic, so that the parametrized types generated by FunsorMeta and used in our interpreters would be proper Python types and we could typecheck our interpreters with standard Python typecheckers. That seems orthogonal to this issue - it is about the types of Funsor interpreters, rather than the types of Funsors themselves as in this issue. Making funsor.Funsors subtypes of Callable, and presumably also including Domains in interpreter type signatures and thus removing the boundary between the two issues, would be much more ambitious and would likely break the first-order restriction.

[fritzo]

@eb8680 great questions!

1. I like the idea of being not only more Pythonic but closer to the NumPy ecosystem. Let's see what we can find. EDIT it looks like numpy-stubs was merged into numpy, providing nascent typing support
2. Yes I believe we can get polymorphism using type variables 🎉
3. Maybe we can make Funsor a subclass of Callable, i.e. a callable that supports pointwise evaluation.
4. I suppose to_funsor() behavior could depend on interpretation...

[fritzo]

@eb8680 suggested this will be easier after #491 , e.g. we could add an @ops.make decorator and use it in the VAE example e.g.

...
encoder = Encoder()
decoder = Decoder()

# Version 0. current state
encode = funsor.function(Reals[28, 28], (Reals[20], Reals[20]))(encoder)
decode = funsor.function(Reals[20], Reals[28, 28])(decoder)

# Version 1. replace Function with and Op
@UnaryOp.make
def encode(image):
    return encoder(image)

@find_domain.register(encode)
def _(op, image):
    return typing.Tuple[Reals[20], Reals[20]]

# Version 2. read signature from the nn.Module
encode = UnaryOp.make(encoder)
# ...find_domain.register

# Version 3. registers find_domain based on type annotations
@UnaryOp.make
def encode(image: Reals[28, 28]) -> typing.Tuple[Reals[20], Reals[20]]:
    return encoder(image)

# Version 4. automatically determines arity
@ops.make(arity=1)  # manually specify
@ops.make  # assume arity = nargs
def encode(image: Reals[28, 28]) -> typing.Tuple[Reals[20], Reals[20]]:
    return encoder(image)