`vmap` doesn't handle non-JAX-types in non-vmap'd outputs

Not sure if this counts as a bug or a feature request.

import jax

def make1(x):
    return [x, 0]

def make2(x):
    return [x, object()]

jax.vmap(make1, out_axes=[0, None])(jax.numpy.zeros(2))
# [DeviceArray([0., 0.], dtype=float32), 0]
jax.vmap(make2, out_axes=[0, None])(jax.numpy.zeros(2))
# TypeError: Output from batched function <object object at 0x7f372dbf9860> with type
# <class 'object'> is not a valid JAX type

If we’re not vmap’ing the value anyway, it’d be nice if non-JAX-types can be valid outputs from vmap’d functions.

In contrast the in-axis equivalent does work:

jax.vmap(lambda x, y: x, in_axes=(0, None))(jax.numpy.zeros(2), object())
# DeviceArray([0., 0.], dtype=float32)

[jakevdp]

I think this is by design: functions passed to vmap are traced, and therefore their outputs are required to be valid JAX types, which can be a normal numeric scalar or array, or a custom type that you have registered.

The in_args version works because the tracing in vmap is only with respect to batched arguments; note however that if you wrap the function in jit (which traces with respect to all arguments) you’ll get the same error there as well.

[patrick-kidger]

That makes a lot of sense. Thanks for the response!