Investigate the possibility of compiling a subset of Kotlin + KMath to TorchScript

[CommanderTvis]

trafficstars

In Torch's Python API, it is possible to do JIT compilation of Python source code to high-performant Torchscript functions:

import torch
import torchvision

class MyScriptModule(torch.nn.Module):
    def __init__(self):
        super(MyScriptModule, self).__init__()
        self.means = torch.nn.Parameter(torch.tensor([103.939, 116.779, 123.68])
                                        .resize_(1, 3, 1, 1))
        self.resnet = torch.jit.trace(torchvision.models.resnet18(),
                                      torch.rand(1, 3, 224, 224))

    def forward(self, input):
        return self.resnet(input - self.means)

my_script_module = torch.jit.script(MyScriptModule())

The task is to investigate if it is possible to analyze a Kotlin class with a compiler plugin or KSP to build a Torchscript module. The result would be a syntax prototype that would possibly involve KMath's Tensor classes.

[CommanderTvis]

Blocked by https://github.com/pytorch/pytorch/issues/50734

[CommanderTvis]

Well, now there is some documentation on the TorchScript format: https://github.com/pytorch/pytorch/blob/master/torch/csrc/jit/docs/serialization.md