Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Use typing and pytypes for type checking and multiple dispatch #434

Closed
wants to merge 5 commits into from

Conversation

eb8680
Copy link
Member

@eb8680 eb8680 commented Jan 20, 2021

Addresses #351, #355.

This is a more ambitious version of the changes in #433 that uses the typing-compatible runtime type checking library pytypes for type checking and inference (via pytypes.is_subtype, a typing-compatible version of issubclass) and for pattern matching (via a pytypes-based fork of multipledispatch that uses pytypes.is_subtype in place of issubclass).

After these changes, it should be possible to use all of typing directly when describing funsor input and output types or when writing patterns, e.g. we could simply replace Product from #430 with typing.Tuple:

class Tuple(Funsor):
    def __init__(self, args):
        ...
        output = typing.Tuple[tuple(arg.output for arg in args)]
        ...

...or use typing.Tuple, including variadic cases, in patterns like the ones needed for Finitary in #423:

@eager.register(Finitary)
def eager_finitary_generic(
        op: EinsumOp, 
        args: typing.Tuple[typing.Union[Tensor, Number], ...]  # variadic
    ) -> typing.Union[Tensor, Number]:
    ...

As suggested in #355, we could even replace many of the assertions in our Funsor.__init__ methods with runtime type checking in interpreters, eliminating lots of boilerplate code and potentially boosting performance (by making many assertions optional):

class FunsorMeta(type):
    def __call__(cls, *args, **kwargs):
        ...
        if interpreter.TYPECHECK:
            assert all(pytypes.is_of_type(arg, static_arg_type)  # is_of_type == isinstance
                        for arg, static_arg_type in zip(args, get_type_hints(cls.__init__)))
        return interpret(cls, *args)
# example
class Reduce(Funsor):
-     def __init__(self, op, arg, reduced_vars):
+     def __init__(self, op: ops.AssociativeOp, arg: Funsor, reduced_vars: typing.FrozenSet[Variable]):
-         assert callable(op)
-         assert isinstance(arg, Funsor)
-         assert isinstance(reduced_vars, frozenset)
-         assert all(isinstance(v, Variable) for v in reduced_vars)
          reduced_names = frozenset(v.name for v in reduced_vars)
          ...

I have also taken the opportunity to remove some of the uglier code in FunsorMeta and domains made obsolete by the use of pytypes and typing, including the custom introspection and subtype checking logic in FunsorMeta.

In #355 I had originally suggested going even further toward integrating typing with Funsor interpretations than the changes here by making parametric Funsor subtypes use typing.Generic. I am somewhat skeptical about taking that additional step because of how difficult it is to interact with typing.Generic at runtime, and have chosen in this PR to retain a simplified version of the GenericTypeMeta class as a substrate for Funsor and FunsorMeta.

One critical question for this approach is how well the multipledispatch fork holds up in terms of correctness and performance and how much additional maintenance effort it would entail. In an ideal world, someone else would implement a performant, fully typing-compatible multipledispatch completely independent of Funsor and we could just use their solution, since that's clearly a useful and interesting direction in its own right and there are a number of new runtime type checkers (1, 2) that could be repurposed to support it.

However, since that is unlikely (AFAIK) in the near term and the developers of multipledispatch are not interested in merging the pytypes-based version (which I took and lightly modified from an old pull request there), we'd have to maintain the fork ourselves. The core functionality seems to be largely complete - almost all of the multipledispatch and Funsor tests pass locally, and pytypes itself is extensively tested. We might need to add a dedicated typing-compatible Dispatcher upon which to base funsor.ops.Op and funsor.registry.KeyedRegistry that makes sure signatures use typing syntax, e.g. converting tuples to typing.Unions or object to typing.Any. As for performance, running tests locally shows there is definitely a slowdown relative to the current multipledispatch, although it's much less severe than the discussion in that multipledispatch PR would suggest.

Another practical issue is the compatibility of pytypes with various Python versions, especially 3.7-3.9. It seems like that's been addressed to a large extent in their master branch, at least for the subset of typing we would use in Funsor, but I'm not entirely sure what its status is or when any outstanding issues might be resolved.

I am cautiously optimistic about these questions, but it's too early to say right now whether they'd end up being show-stoppers relative to the status quo or the simpler approach in #433.

@eb8680
Copy link
Member Author

eb8680 commented Feb 2, 2021

Closing this, I believe we can work around depending on pytypes or forking multipledispatch.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
None yet
Development

Successfully merging this pull request may close these issues.

1 participant