The following shows an example code for adding new layers to SOL’s PyTorch frontend.
import sol.pytorch
def parse_inputs(node, scope):
return tuple(scope[i.unique()] for i in node.inputs())
def my_handler(node, scope):
a, b, c = parse_inputs(node, scope)
x = sol.hlir.Tensor(my_backend_lib.add_my_layer(a, b, c))
scope[node.output().unique()] = x
sol.pytorch.add_handler("aten::not_implemented_by_sol", my_handler)
my_handler
gets two arguments: 1. the TorchScript node of the layer, 2. a
dictionary that contains all variables and tensors visible on the current scope.
Within your handler you can either use sol.hlir
API or your own backend to
add layers to the neural network representation within SOL. Last, don’t forget
to add the output tensor of your layer to the scope, so that following layers
can access it!