mx.vmap throws error "[tree_flatten] The argument should contain only arrays"

[dasayan05]

I am implementing a simple loss function called "Implicit Score Matching" that involves a parametric score function (simple MLP in our case) and data. For reference, this equation sums it up

I am trying to implement it in mlx. I have a non-batched version of the loss that works fine on single sample. But when I try vmap to get batched version, it throws an error.

Model class & Jacobian implementation

from typing import Callable

import mlx.nn as nn
import mlx.core as mx

class MLP(nn.Module):
    def __init__(self, d: int = 2):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Linear(d, 32), nn.ReLU(),
            nn.Linear(32, d)
        )
    def __call__(self, x):
        return self.layers(x)

def jacobian(f: Callable | nn.Module, x: mx.array):
    d, = x.shape
    I = mx.eye(d)
    jac = []
    for v in I:
        _, (jvp, ) = mx.jvp(f, [x, ], [v, ])
        jac.append(jvp)
    return mx.stack(jac, 0)

The loss implementation is given below

def _ism(x: mx.array, score: Callable):
    # Implicit Score Matching
    out = score(x)
    jac = jacobian(score, x)
    tr_jac = mx.trace(jac)
    half_norm = 0.5 * (mx.linalg.norm(out) ** 2)

    return half_norm - (- tr_jac)

batch_size, dim = 32, 2
X = mx.random.normal((batch_size, dim)) # data

mlp = MLP(dim) # simple MLP(2->32->2)

def S(x: mx.array): # score function S: R^2 -> R^2
    return mlp(x)

This one works fine

loss = _ism(X[0], S) # single sample (OK)

But the vmap-ed version throws error

loss = mx.vmap(_ism, in_axes=(0, None), out_axes=0)(X, S) # batched (ERROR)

Exception has occurred: ValueError       (note: full exception trace is shown but execution is paused at: _run_module_as_main)
[tree_flatten] The argument should contain only arrays
  File "/Users/ayandas/Codes/experiment/mlx/mlx-examples/score_matching/tmp.py", line 44, in <module>
    loss = mx.vmap(_ism, in_axes=(0, None), out_axes=0)(X, S) # batched (ERROR)
  File "/Users/ayandas/.local/share/uv/python/cpython-3.10.16-macos-aarch64-none/lib/python3.10/runpy.py", line 86, in _run_code
    exec(code, run_globals)
  File "/Users/ayandas/.local/share/uv/python/cpython-3.10.16-macos-aarch64-none/lib/python3.10/runpy.py", line 196, in _run_module_as_main (Current frame)
    return _run_code(code, main_globals, None,
ValueError: [tree_flatten] The argument should contain only arrays

I must be doing something stupid. Any help is appreciated.

[awni]

You can't vmap functions which take non-mx.array inputs. In your case that would be the callable S. You can instead use a closure or functools.partial. For example with a closure:

def _ism(x, score):
   ...

def vmappable_ism(x):
   return _ism(x, score)

mx.vmap(vmappable_ism, ...)

[dasayan05]

Ok thanks. It worked.

But may I ask why exactly is it that way ? Is there any reason it cannot take non-mx.array input. If I recall correctly from jax (which I didn't use for a while), they allow non-array inputs.

Also, mx.vmap's doc does not mention this. Someone coming from jax would probably expect this to work.

[dasayan05]

Is there a nicer way other than using closure or partial ?

[awni]

Also, mx.vmap's doc does not mention this. Someone coming from jax would probably expect this to work.

We should probably mention it in the doc and/or allow non-array inputs for None vmap axes. Both are pretty reasonable imo.

Is there a nicer way other than using closure or partial ?

I don't think so.. but I don't think either of those options is all that bad.