3 ms·
Great! How do you program it?
by fooblaster 9mo ago
Great! How do you program it?
- alanma 9mo agoA 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