SplitInfinity:


How can we give it a custom callable type for a function and force it to recompile a copy of the function body with the new types (but a new name)?

Could you provide a code sample to clarify this question? What is it that you want to do?


I’d like to be able to say something like:
def generic_loop(self, inputs, state):
    " ^^^ I'm trying to reuse this. ^^^ "
    outs = []
    for x in inputs:
        out, state = self.fwd(x, state)
        outs.append(out)
    return outs, state
class Cell(nn.Module):
    def __init__(self, forward_sig):
        self.forward = apply_cell_sig(generic_loop, forward_sig)
        # self.forward : (TCustomCell, List[TCustomInput], TCustomState) -> TCustomOut

class CustomCell(Cell):
    def __init__(self):
        super().__init__(
            get_cell_sig(
                CustomCell, TCustomInput, TCustomState, TCustomOut))
    def fwd(self, input: TCustomInput, state: TCustomState):
        return blah(input, state)

More generally, I’m looking for a way to treat a function body as a “template” that I can provide specific types for, and ask the JIT explicitly to give me a new version for it. So it doesn’t just see its full name and go “I’ve already seen this; here’s the cached version”.
By your confirmation though, this’s probably very difficult. I couldn’t get the JIT to accept a function that takes in a nn.Module instance as an argument either (besides self. accesses) so there’re multiple issues here.