[alanma]

A couple core commands in our ISA detailed on our GitHub, map your problem to matrix ops, here's a brief excerpt, but our tpu_compiler and tpu_driver are the core to programming your own:

from tpu_compiler import TPUCompiler, TPURuntime

class Custom(nn.Module):

    def __init__(self):
        super().__init__()

        self.layer1 = nn.Linear(2, 2, bias=False)
        self.layer2 = nn.Linear(2, 2, bias=False)

    def forward(self, x):
        x = self.layer1(x)
        x = torch.relu(x)
        x = self.layer2(x)
        return x

model = train_model(your_data)

# compile to the tiny tiny TPU format

compiler = TPUCompiler()

compiled = compiler.compile(model)

# run and enjoy :)

runtime = TPURuntime(tpu)

result = runtime.inference(compiled, input_data)

Will update soon with some better documentation, but hopefully this will get you started!

- Alan and Abiral