Support opting in to spikes on the forward pass

[arvoelke]

Currently there isn't an easily accessible or documented way to train with spikes on the forward pass. It seems like it can be done by patching in your own builder method that overrides the training_step for each neuron model that is being used. This isn't thoroughly tested but it seems to work (e.g., for the nengo.StochasticSpiking wrapper in Nengo 3.1.0):

from nengo_dl import compat
from nengo_dl.neuron_builders import SimNeuronsBuilder, StochasticSpikingBuilder

class CustomStochasticSpikingBuilder(StochasticSpikingBuilder):

    def training_step(self, J, dt):
        return self.step(J, dt)


SimNeuronsBuilder.TF_NEURON_IMPL[compat.StochasticSpiking] = CustomStochasticSpikingBuilder