vmap with static (optional) keyword-arguments

[adtzlr]

Hi,

I created a custom wrapper for vmap which treats all trailing, non-specified arguments as static, or in other words, non-mapped (keyword + optional keyword) arguments. This is sometimes asked for, see e.g. #13836, #13090, https://stackoverflow.com/questions/74969744/jax-vmap-in-axes-doesnt-work-if-keyword-argument-is-passed.

The wrapper is based on the answer given in #13836, but more general. I hope I'm not missing something here. Do you think such helpers should be avoided in general?

Here's the code-block:

Decorator my_vmap()

from functools import wraps
import inspect
import jax

def my_vmap(fun, in_axes=0, out_axes=0):
    @wraps(fun)
    def vmap_with_static_kwargs(*args, **kwargs):
        # sorted list of all parameter keys, including kwargs with default values
        sig = inspect.signature(fun)
        keys = [
            key
            for key, value in sig.parameters.items()
            if not (key in ["args", "kwargs"] and value.default == inspect._empty)
        ]
        
        # check if unexpected keyword-argument is given
        for key in kwargs.keys():
            if key not in keys:
                raise TypeError(
                    f"{fun.__name__}() got an unexpected keyword argument '{key}'"
                )

        # dict with default values for all parameters
        parameters = dict(
            [(key, value.default) for key, value in sig.parameters.items()]
        )

        # merge dict of default values with custom keyword arguments
        items = {**parameters, **kwargs}

        # create sorted list of values of keyword-arguments, including default kwargs
        keyword_args = [items[key] for key in keys[len(args) :]]

        # don't map non-given arguments and keyword-arguments
        if not hasattr(in_axes, "__len__"):
            in_axes_tuple = (in_axes,)
        else:
            in_axes_tuple = in_axes

        static_argnums = len(args) + len(keyword_args) - len(in_axes_tuple)
        in_axes_new = (*in_axes_tuple, *([None] * static_argnums))

        return jax.vmap(fun, in_axes=in_axes_new, out_axes=out_axes)(
            *args, *keyword_args
        )

    return vmap_with_static_kwargs

Example Script

import jax.numpy as jnp
import numpy as np


def f(x, y, a=1):
    return jnp.array([jnp.linalg.det(x), jnp.linalg.det(y)])


x = np.random.rand(3, 3, 1000)
y = np.random.rand(3, 3, 1000)
h = my_vmap(f, in_axes=(-1, -1), out_axes=-1)

assert np.allclose(h(x, y), h(x, y, a=1.0))
assert np.allclose(h(x, y), h(x=x, y=y))
assert np.allclose(h(x, y), h(a=1.0, y=y, x=x))
assert np.allclose(h(x, y), h(x, y, 1.0))

[jakevdp]

That's a nice approach! It does a good job working around what I think is the fundamental issue of vmap: that it returns a wrapper that treats positional and keyword arguments differently. The "real" fix would be to modernize the vmap API (similar to the work a while ago with jit to handle both static_argnames and static_argnums in a uniform way), but that's not something anybody has thought deeply about yet I think.

[adtzlr]

Thanks @jakevdp for the details! I just wasn't sure if there are any reasons why such a decorator should not be used.

I'm going to add JAX as AD-backend in 🔍 FElupe for hyperelastic material formulations. Right now, I'm using a self-written forward-mode AD package tensortrax - but like autograd, there are limits without a (JIT) compiler. The classic simple hyperelastic material models are all okay, but more advanced material models are getting slow. I really like JAX because of its NumPy-like API and the vmap/jit-features are great 🥇.