Struggling to understand this error: ValueError: [eval] Attempting to eval an array without a primitive.

[ali-shihab]

I've gone through the docs but I just don't understand what causes the "ValueError: [eval] Attempting to eval an array without a primitive." error? For example, I'm training a mamba implementation that can be found at: https://github.com/alxndrTL/mamba.py/tree/main/mlx, in which I replaced the stock linear output with a custom classification output. In the training code, the following function raised this error:

@partial(mx.compile, inputs=state, outputs=state)
def step(inputs, targets):
    loss_and_grad_fn = mnn.value_and_grad(model, loss_fn)
    loss, grads = loss_and_grad_fn(model, inputs, targets)
    optimizer.update(model, grads)
    return loss

Commenting out the compilation line fixed the issue after a whole day of trying to figure it out; I still don't know why. Similarly, the following vectorised map function (I'm new to vmap), for two identically shaped arrays, throws the same value error:

def reshape_labels(self, y, out):
  if y[0] >= 0: out[y[0]] = 1
  else: out = y

map_labels = mx.vmap(self.reshape_labels, in_axes=(-1, -1))

map_labels(labels_a, labels_b)

And I'm yet again confused. What am I doing wrong and what exactly does the error mean? The docs say something about primitives defining the way arrays are evaluated, but I'm not sure I really understand the explanation.

[awni]

Sorry that message is not so helpful, we should add a documentation on that. There's basically two cases in which you could see that message:

1. During compilation if the state or function parameters/returns don't capture all the input/output arrays. For example:

For example, updating z inside the compiled function here is not allowed unless you specify it as an output.

z = []

@mx.compile
def fun(x, y):
    z.append(x + y)
    return x * y

out = fun(mx.array(1.0), mx.array(1.0))
print(z) # ValueError: [eval] Attempting to eval an array without a primitive.

2. If you eval during a transformation in which graph evaluation is not allowed. For example the condition y[0] >= 0 is an implicit graph evaluation because it has to read y[0] in order to do the control flow in Python. So vmap or compile around that is not allowed and will cause that error.

[ali-shihab]

Thanks for the quick response. In the case of the step function, isn't state implicitly captured as input & output in compile ? Since state = [model.state, optimizer.state].

Additionally, the Error is raised at this point in the call stack:

File [/opt/homebrew/lib/python3.11/site-packages/mlx/nn/utils.py:34](https://file+.vscode-resource.vscode-cdn.net/opt/homebrew/lib/python3.11/site-packages/mlx/nn/utils.py:34), in value_and_grad.<locals>.wrapped_value_grad_fn(*args, **kwargs)
     [32](https://file+.vscode-resource.vscode-cdn.net/opt/homebrew/lib/python3.11/site-packages/mlx/nn/utils.py:32) @wraps(fn)
     [33](https://file+.vscode-resource.vscode-cdn.net/opt/homebrew/lib/python3.11/site-packages/mlx/nn/utils.py:33) def wrapped_value_grad_fn(*args, **kwargs):
---> [34](https://file+.vscode-resource.vscode-cdn.net/opt/homebrew/lib/python3.11/site-packages/mlx/nn/utils.py:34)     value, grad = value_grad_fn(model.trainable_parameters(), *args, **kwargs)
     [35](https://file+.vscode-resource.vscode-cdn.net/opt/homebrew/lib/python3.11/site-packages/mlx/nn/utils.py:35)     return value, grad

Which seems to me that it's the module.trainable_parameters() call that is raising this error - is this expected behaviour, given the context?

[awni]

In the case of the step function, isn't state implicitly captured as input & output in compile ? Since state = [model.state, optimizer.state]

There should be some other state you are using. The common one that people miss is the RNG state if they are using the global RNG. See here for example. If that's not it, I don't know why it's throwing that error for you. I'd be happy to take a look if you are able to share some code to reproduce.

[ali-shihab]

Ah yes, you're right, I wasn't using the RNG but it was due to the array being indexed inside the function call. The solution was as follows:

• Converting to arrays to numpy, and performing the map with numpy's apply_along_axis was a workaround for vmap;
• The map to reshape_labels function was being called inside the model's __call__ function, which of course needed to eval before it was being computed. Moving it outside the function fixed it.

Thanks a lot! :)