How to efficiently use nnx.jit for methods?

[cisprague]

Hi! I am a big fan of NNX so far, great job!

I am trying to make a "training" module that includes everything (model, dataset, metrics, etc.), but I find that nnx.jiting the train_step method when using self is much slower. As a result, I resorted to using staticmethod. Are there some best practices on how to do this properly? Please see the code below (not the full code, but I hope it gets the point across).

Thank you!

@dataclass(kw_only=True)
class Trainer0:

    optimizer: nnx.Optimizer
    dataset: Dataset # my own dataset class

    def __post_init__(self):
        self.train_step = nnx.jit(self._train_step)

    @staticmethod
    def _train_step(optimizer, batch):
        f = lambda model: model.loss(batch)
        loss, grads = nnx.value_and_grad(f)(optimizer.model)
        optimizer.update(grads)
        return loss
    
    def __call__(self, n_epochs):
        for _ in range(n_epochs):
            for batch in self.dataset.train_batches():
                loss = self.train_step(self.optimizer, batch)
                print(f'Loss: {loss:.2f}')
    
@dataclass(kw_only=True)
class Trainer1(nnx.Module):

    optimizer: nnx.Optimizer
    dataset: Dataset # my own dataset class

    @nnx.jit
    def train_step(self, batch):
        f = lambda model: model.loss(batch)
        loss, grads = nnx.value_and_grad(f)(self.optimizer.model)
        self.optimizer.update(grads)
        return loss
    
    def __call__(self, n_epochs):
        for _ in range(n_epochs):
            for batch in self.dataset.train_batches():
                loss = self.train_step(batch)
                print(f'Loss: {loss:.2f}')

>>> trainer = Trainer0(optimizer=optimizer, dataset=dataset)
>>> trainer(10) # much faster

>>> trainer = Trainer1(optimizer=optimizer, dataset=dataset)
>>> trainer(10) # much slower

[cgarciae]

Hey @cisprague, can you verify that Trainer1.train_step is not constantly recompiling by adding a print statement and making sure you only see the message once?