Thanks again for the repost and all the support!! Been a blast and super cool to see the interest, if you want to follow along for more of our writeups, our blog can be found here: https://chewingonchips.substack.com/
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!
