The retracing is avoid by making the number of calls a tensor instead of a python/numpy primitive.
By doing that, however, one needs to ensure that the entire thing is flexible with the number of calls. Basically anything that is explicitly compiled would need a signature.
(I've added a crappy warning which should obviously be improved)
Sorry about the whitespaces and formatting, makes the changes more difficult to see.
The retracing is avoid by making the number of calls a tensor instead of a python/numpy primitive. By doing that, however, one needs to ensure that the entire thing is flexible with the number of calls. Basically anything that is explicitly compiled would need a signature.
(I've added a crappy warning which should obviously be improved)
Sorry about the whitespaces and formatting, makes the changes more difficult to see.