tensorscript compiler

Imagine a framework-agnostic DSL with strong typing, dimension and ownership checking and lots of syntax sugar. What would it be like? As interesting as it is, here is why there needs to a langauge for neural network:

1. Multi-target. Write once, run everywhere(interpreted in PyTorch or compiled to Tensorflow MetaGraphDef).

2. Typechecking with good type annotation.

3. Parallellize with language-level directives

4. Composition(constructing larger systems with building blocks) is easier

5. Ownership system because GC in CUDA is a nightmare

6. No more frustration from deciphering undocumented code written by researchers. The issue is, an overwhelming majority of researchers are not programmers, who care about the aesthetics of clean code and helpful documentation. I get it - people are lazy and unsafe unless the compiler forces them to annotate their code.

About syntax:

1. Stateful symbols are capitalized

2. Impure functions must be annotated

3. Local variables cannot be shadowed

4. Function args have keyword

use conv::{Conv2d, Dropout2d, maxpool2d};
use nonlin::relu;
use lin::Linear;

// three levels of abstraction: node, weights, graph
// ? stands for batch size, channel, height, width
// outputs tensor of shape [batch_size, 10]
node Mnist<?,c,h,w -> ?,OUT> {
    // define symbol level macro for both type/dimension and scope variable
    FC1 = 320;
    FC2 = 50;
    OUT = 10;
}
weights Mnist<?,c,h,w -> ?,OUT> {
    conv1 = Conv2d::<?,c,h,w -> ?,c,h,w>::new(in_ch=1, out_ch=10, kernel_size=5);
    conv2 = Conv2d::<?,c,h,w -> ?,c,h,w>::new(in_ch=10, out_ch=20, kernel_size=5);
    dropout = Dropout2d::<?,c,h,w -> ?,c,h,w>::new(p=0.5);
    fc1 = Linear::<?,FC1 -> ?,FC2>::new();
    fc2 = Linear::<?,FC2 -> ?,OUT>::new();
}
graph Mnist<?,c,h,w -> ?,OUT> {
    fn new() {
        fc1.init_normal(std=1.);
        fc2.init_normal(std=1.);
    }

    fn forward(self, x) {
        x
        |> conv1            |> maxpool2d(kernel_size=2)
        |> conv2 |> dropout |> maxpool2d(kernel_size=2)
        |> view(?, FC1)
        |> fc1 |> relu
        |> self.fc2()
        |> log_softmax(dim=1)
    }

    fn fc2(self, x: <?,FC2>) -> <?,OUT> {
        x |> fc2 |> relu
    }
}

compiles to

from torch import nn, F

class Mnist(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
        self.conv2_drop = nn.Dropout2d()
        self.fc1 = nn.Linear(320, 50)
        self.fc2 = nn.Linear(50, 10)

    def forward(self, x):
        x = F.relu(F.max_pool2d(self.conv1(x), 2))
        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
        x = x.view(-1, 320)
        x = F.relu(self.fc1(x))
        x = F.dropout(x, training=self.training)
        x = self.fc2(x)
        return F.log_softmax(x, dim=1)

Built With

• pytorch
• rust