mx.eval(l, a) array without primitive error

[aimaregana]

Hello everyone,

I'm trying to go from PyTorch to MLX due I'm using Mac mini M4 Pro computer and I know the speed improvement using MLX.

I tested speed comparison "same code" for PyTorch and mlx and I got:

about 300 img/s with PyTorch
about 850 img/s with MLX.

My main problem is the mx.eval(l, a) that returns an error of "attempting to eval without primitive" after first epoch.

As you can see there are 877 img/s at first epoch and then the error comes.

Anyone could help me hot to solve this problem?

Best regards in advance!!!!

If I remove @mx.compile in train_step and eval_step everything works but the speed goes down to 130 img/s.

My code:

kaggle MLX 2.py

[awni]

When you compile it's important to capture any implicit inputs and outputs to the function:

from functools import partial

   state = [model, optimizer.state]

    @partial(mx.compile, inputs=state, outputs=state)
    def train_step(x, y):
        (loss, acc), grads = loss_and_grad_fn(model, x, y)
        optimizer.update(model, grads)
        return loss, acc

And below:

            loss, acc = train_step(x, y)
            mx.eval(loss, acc, state)

[awni]

There is more on this in the documentation: https://ml-explore.github.io/mlx/build/html/usage/compile.html#pure-functions

[aimaregana]

Thanks awn!!!!

This way works but goes from 877img/s to 150 img/s. With PyTorch I have twice img/s. Is it usual?

[awni]

Yes it might be expected, I'm not certain. I do know we have seen slow resnet training compared to PyTorch with MPS in the past. We haven't had the bandwidth to optimize it yet. See e.g. #2770