ValueError in mx.grad When Using mx.grad

[alien2327]

Hello,

I'm experiencing an issue with the mx.grad function in the mlx library. My goal is to use automatic differentiation feature for a project, but I'm encountering a ValueError when I attempt to execute the following code:

import scipy.stats as sp
import numpy as np
import mx.core as mx
import mx.nn as nn

# Some model I defined
class Model(nn.Module):
    def __init__(self):
        super().__init__()
        # input 2 nodes -> output 1 node, all nn.Linear
        ...
    def __call__(self, x, t): ...
model = Model()

# Code for sampling x and t
samples = sp.stats.qmc.LatinHypercube(d=1).random(100)
x_bnd_cond = sp.stats.qmc.scale(samples, [0.0], [1.0])
t_bnd_cond = np.zeros_like(x_bnd_cond)

# Change numpy array to mlx array
x_bnd_cond = mx.array(x_bnd_cond)
t_bnd_cond = mx.array(t_bnd_cond)

# Error occurs here
u_x = mx.grad(model, argnums=0)(x_bnd_cond, t_bnd_cond)
u_t = mx.grad(model, argnums=1)(x_bnd_cond, t_bnd_cond)

The error message I receive is:

ValueError: [vjp] Output shape (100,1) does not match cotangent shape ().

I'm not sure why the output shape and the cotangent shape are not matching. Could this be a bug, or am I possibly misusing the mx.grad function? Any guidance on how to resolve this issue would be greatly appreciated.

Thank you for your assistance.

Best regards,

[awni]

The output of the model.__call__ needs to be a scalar (a.ndim==0) for mx.grad to work.