SplitInfinity:

I don’t think it is a publicly advertised feature, but you could try using @torch.jit._overload . You can find examples of how to use this decorator in test/test_jit.py , like test_function_overloads . The caveat here is that the one body you supply will be used with all overloaded signatures because there is no concept of TypeVar in the JIT typing system. So you might have do type refinement to write one body that works for all types (see my_conv in test_function_overloading_isinstance , also in test_jit.py ).


This looks promising!
Can you elaborate on what’s happening here?
        # TODO: pyflakes currently does not compose @overload annotation with other
        # decorators. This is fixed on master but not on version 2.1.1.
        # Next version update remove noqa and add @typing.overload annotation

        @torch.jit._overload  # noqa: F811
        def test_simple(x1):  # noqa: F811
            # type: (int) -> int
            pass

        @torch.jit._overload  # noqa: F811
        def test_simple(x1):  # noqa: F811
            # type: (float) -> float
            pass

        def test_simple(x1):  # noqa: F811
            return x1

        def invoke_function():
            return test_simple(1.0), test_simple(.5)

        self.checkScript(invoke_function, ())

        # testing that the functions are cached
        compiled_fns_1 = torch.jit._script._get_overloads(test_simple)
        compiled_fns_2 = torch.jit._script._get_overloads(test_simple)
        # ^^^ HERE HERE HERE ^^^ #
        for a, b in zip(compiled_fns_1, compiled_fns_2):
            self.assertIs(a.graph, b.graph)

Why could the two successive invocations of torch.jit._script._get_overloads(test_simple) get different results? At that point, there’s only one thing attached to the test_simple name.