I think it's good since it makes @torch.jit.script more consistent, namely that you can add it as a decorator on types that inherit from object but not nn.Module, and for a function you can add @torch.jit.script or call torch.jit.script() yourself, but this is different for modules. I agree that it's becoming a pretty overloaded name in that it does too many things through 1 entry point and it's not exactly clear / easily describable how it can be used. Maybe we should re-name the decorator used for user defined classes to something like @torch.jit.class.

If the problem you're trying to solve is to avoid having to call the recursive scripting func all the time, why can't we just recommend something like

class _MyMod(torch.nn.Module):
    ...

MyMod = torch.compile(_MyMod)

I don't think it's a big deal either way, but we should minimize the number of ways to do the same thing I think.

I think it's good since it makes @torch.jit.script more consistent, namely that you can add it as a decorator on types that inherit from object but not nn.Module, and for a function you can add @torch.jit.script or call torch.jit.script() yourself, but this is different for modules.

That's a pretty good argument, I'm fine with it either way I think.

If the problem you're trying to solve is to avoid having to call the recursive scripting func all the time, why can't we just recommend something like

class _MyMod(torch.nn.Module):
    ...

MyMod = torch.compile(_MyMod)

I don't think it's a big deal either way, but we should minimize the number of ways to do the same thing I think.

Having a function like that is the same as having a decorator (it just calls the function with the decorated type under the hood), similar to how you can use @script on functions or pass a Python function to script() already.

I would really really hope that this can be updated and get through some day!

For me, the following simple patch works:

import functools

import torch
import torch.nn as nn


def compile(kind='script', *args, **kwargs):
    if isinstance(kind, type) and issubclass(kind, nn.Module):
        raise RuntimeError('Use @torch.jit.compile(kind, ...) as decorator')

    if callable(kind):
        compiler = kind 
    elif kind == 'script':
        compiler = torch.jit.script
    else:
        assert kind == 'trace'
        compiler = torch.jit.trace

    class CompiledModuleMeta(type(nn.Module)):
        def __call__(cls, *ctor_args, **ctor_kwargs):
            return torch.jit.script(super(CompiledModuleMeta, cls).__call__(*ctor_args, **ctor_kwargs), *args, **kwargs)

    def decorator(cls):
        return CompiledModuleMeta(cls.__name__, (cls,), {})

    return decorator


assert not hasattr(torch.jit, 'compile')
torch.jit.compile = compile

@torch.jit.compile()
class MLP(nn.Module):
    r"""my mlp"""
    def __init__(self):
        super().__init__()
        self.model = nn.Sequential(nn.Linear(3, 4), nn.ReLU(True), nn.Linear(4, 3))

    def forward(self, x):
        return self.model(x)

@ssnl We’re going with the contract everything under script is compiled or at least looked at by the compiler, but that doesn’t hold for modules so this API isn’t going to get landed

@driazati That is reasonable, but results in a really ugly API. Could another API with a different name be added?