Behavior of custom_vmap with no batched dimensions

[justindomke]

Hello,

I'm interesting in writing custom vmap rules for some functions. I've been investigating the use of jax.custom_batching.custom_vmap. Everything seems to make sense, except I can't quite understand what's happening in the case where I don't provide any batched dimensions to vmap.

Here's a simple example:

@jax.custom_batching.custom_vmap
def myfun(x):
    print("MYFUN")
    return x+2

@myfun.def_vmap
def vmap_rule(axis_size, in_batched, xs):
    print("VMAP_RULE")
    return xs+2, True

batched_fun = jax.vmap(myfun,[0])

xs = jnp.array([0,1,2])
print(batched_fun(xs))
print(jax.make_jaxpr(batched_fun)(xs))

If I run this, I get this output:

MYFUN
VMAP_RULE
[2 3 4]
MYFUN
VMAP_RULE
{ lambda ; a:i32[3]. let b:i32[3] = add a 2 in (b,) }

No problem so far. The output is as expected, and this all seems to make sense. But consider running this code instead:

@jax.custom_batching.custom_vmap
def myfun(x):
    print("MYFUN")
    return x+2

@myfun.def_vmap
def vmap_rule(axis_size, in_batched, xs):
    print("VMAP_RULE")
    return xs+2, True

batched_fun = jax.vmap(myfun,[None],axis_size=3) # <---- no mapped axes

print(batched_fun(1))
print(jax.make_jaxpr(batched_fun)(1))

If I run this, I get this output:

MYFUN
[3 3 3]
MYFUN
{ lambda ; a:i32[]. let
    b:i32[] = custom_vmap_call[
      call={ lambda ; c:i32[]. let d:i32[] = add c 2 in (d,) }
      in_tree=PyTreeDef(([], (*,)))
      out_tree=PyTreeDef(*)
      rule=<function vmap_rule at 0x11e99e5c0>
    ] a
    e:i32[3] = broadcast_in_dim[broadcast_dimensions=() shape=(3,)] b
  in (e,) }

This gives the expected output. But my custom vmap rule was never called. From the jaxpr, I think that's because vmap thinks it's pointless to run the function multiple times—it can just run the function once and then make a copy of the result.

That's clearly a good optimization in most cases. But is there any way to avoid it? The problem for me is that I want to create "abstract" primitives that represent distributions. These don't themselves have any implementation. Instead, I have written program transformations like sample() and prob() that inject functionality (PRNGKeys, etc). But the injected functionality depends on the transformation, so I don't want to inject it when calling vmap, I only want to do it later.

For "normal" vmap calls where at least one dimension is batched, I can intercept the vmap call with a def_vmap rule that can bind a new "batched" abstract primitive, which I can then deal with in the program transformation. That all seems to work very well! But in the above case, the def_vmap rule gets skipped, and I seem to be out of luck—the jaxpr represents calling the abstract primitive once, and then making a bunch of copies of that. I don't want that because I want to eventually feed different PRNGKeys to each of the "branches".

Is there any solution to this? If it were possible, the ideal would just be a way to disable this "no batched dims optimization", so that my def_vmap rule always gets triggered by jax.vmap.

[jakevdp]

Thanks for the question! The issue is that when none of the axes are mapped, the batching tracer falls back to calling the primitive itself. The code is here: https://github.com/google/jax/blob/bfe8acb31e04a540daad3f568239ec0e5c3f0d0f/jax/_src/interpreters/batching.py#L393-L394

In the use-cases that vmap was designed for, this is always correct: a batched operation with no batched axes is just the original operation.

I think changing this would be a big task: if we were to call the batching rule even in cases where there was no batching, every single batching rule currently in use would have to be updated to handle this non-batched case. That would be a big project just for the batching rules defined in jax itself, let alone the batching rules defined by every downstream library that uses JAX. This would be a potentially large breakage of existing code, for the sake of changing something that already works correctly according to the contract of the vmap API. We'd have to have a very good motivation to make a change of that scope.

I wonder if there are other ways you could handle this in your case – for example, could you define a default implementation that handles the non-batched calling path for your primitive?

[jakevdp]

Thanks! I'm not going to be very responsive over the next several days; tagging @mattjj who loves this kind of thing and may be able to help if you run into further issues 😁

[justindomke]

OK, so I tried this. I think this solution might be more elegant than the previous way I was trying to make things vmappable. But I don't think it solves this particular problem, at all.

Here's my code from above, updated to make normal_p closed under batching.

import jax
from jax._src.util import safe_map
from jax import numpy as jnp
from jax._src.tree_util import tree_structure, tree_unflatten
from jax._src import abstract_arrays
from jax.interpreters import batching

def normal_sample(key, loc, scale):
    assert loc.shape == scale.shape
    if loc.shape == ():
        # scalar inputs
        return loc + scale * jax.random.normal(key)
    elif len(loc.shape)==1:
        # 1-D array inputs of same length
        keys = jax.random.split(key,loc.shape[0])
        return jax.vmap(normal_sample)(keys,loc,scale)
    else:
        raise NotImplementedError("")

base_key = jax.random.PRNGKey(0)

##########################################################################################
# Define an "abstract" normal distribution. It has an abstract eval, but no implementation
##########################################################################################

normal_p = jax.core.Primitive("normal")

def normal(loc, scale):
    return normal_p.bind(loc, scale)

@normal_p.def_abstract_eval
def abstract_eval(loc, scale):
    print("abstract")
    out_shape = jax.eval_shape(normal_sample, base_key, loc, scale)
    return abstract_arrays.ShapedArray(out_shape.shape, out_shape.dtype)

def batch_eval(vector_arg_values, batch_axes):
    print("batch")
    assert batch_axes == (0,0)
    loc_vec, scale_vec = vector_arg_values
    res = normal_p.bind(loc_vec,scale_vec)
    return res, batch_axes[0]
batching.primitive_batchers[normal_p] = batch_eval

##########################################################################################
# Take a jaxpr that includes `normal` primitives, and add keys wherever needed
##########################################################################################

def sample_eval_jaxpr(closed_jaxpr, args_flat, key):
    env = {}

    jaxpr = closed_jaxpr.jaxpr
    consts = closed_jaxpr.literals

    def read(var):
        if type(var) is jax.core.Literal:
            return var.val
        return env[var]

    def write(var, val):
        env[var] = val

    safe_map(write, jaxpr.invars, args_flat)
    safe_map(write, jaxpr.constvars, consts)

    for eqn in jaxpr.eqns:
        invals = safe_map(read, eqn.invars)
        if eqn.primitive == normal_p:
            loc, scale = invals
            key, subkey = jax.random.split(key)
            outvals = normal_sample(subkey, loc, scale)
        # extra elif blocks to handle cond_p, scan_p, etc. removed for simplicity
        else:
            outvals = eqn.primitive.bind(*invals, **eqn.params)
        if not eqn.primitive.multiple_results:
            outvals = [outvals]
        assert len(eqn.outvars) == len(outvals)
        safe_map(write, eqn.outvars, outvals)

    return safe_map(read, jaxpr.outvars)


##########################################################################################
# program transformation: take a function that contains normals (can't be called)
# and get a new function that splits PRNGKeys as needed and does sampling
##########################################################################################

def sample(fun):
    def newfun(key, *args):
        args_flat, args_tree = jax.tree_util.tree_flatten(args)
        def fun_flat(*args_flat):
            args = jax.tree_util.tree_unflatten(args_tree, args_flat)
            return fun(*args)

        closed_jaxpr, out_shape = jax.make_jaxpr(fun_flat, return_shape=True)(*args_flat)
        out_flat = sample_eval_jaxpr(closed_jaxpr, args_flat, key)
        out_tree = tree_structure(out_shape)
        out = tree_unflatten(out_tree, out_flat)
        return out

    return newfun

If I call this with a standard vmap everything works nicely. If I run

def diag_normal(locs,scales):
    return jax.vmap(normal)(locs,scales)

locs = jnp.array([-1.0, 2.0, 3.0])
scales = jnp.array([2.0,3.0,4.0])
print(jax.make_jaxpr(diag_normal)(locs,scales))

key = jax.random.PRNGKey(0)
print(sample(diag_normal)(key,locs,scales))

Then I get:

batch
abstract
{ lambda ; a:f32[3] b:f32[3]. let c:f32[3] = normal a b in (c,) }
batch
abstract
[0.571116  9.120268  3.3160863]

But if I call it with no batched dimensions? I get the same problem as before. If I run this code:

def iid_normal(loc,scale):
    return jax.vmap(normal,[None,None],axis_size=3)(loc,scale)

loc = jnp.array(-1.0)
scale = jnp.array(3.0)
print(jax.make_jaxpr(iid_normal)(loc,scale))

key = jax.random.PRNGKey(0)
print(sample(iid_normal)(key,loc,scale))

then I get this output:

abstract
{ lambda ; a:f32[] b:f32[]. let
    c:f32[] = normal a b
    d:f32[3] = broadcast_in_dim[broadcast_dimensions=() shape=(3,)] c
  in (d,) }
abstract
[-4.7546167 -4.7546167 -4.7546167]

My batching rule is never used, and so I get a single sample copied 3 times instead of multiple samples.

In principle, I guess one option would be to have my custom interpreter handle broadcast_in_dim primitives by checking if they are being run on an input that produced by a normal_p primitive, and if so, calling special logic there? Not sure how robust that would be...

[jakevdp]

Sure, it's expected that the batching rule is never used here (the batching dynamics we initially discussed still hold). My suggestion was to move the relevant logic from the batching rule into the abstract eval rule, and make the batching rule essentially a no-op.

[justindomke]

Sorry for being dense, but could you elaborate a bit on how this would work? It seems to me that regardless of what's in the abstract eval rule, if I call vmap with in_axes=[None]*n and axis_shape=3, then abstract_eval rule is only called with non-batched input shapes, meaning there's no way for it to "know" about batching.

Here's a minimal example of what I mean:

import jax
from jax import numpy as jnp
from jax._src import abstract_arrays
from jax.interpreters import batching

base_key = jax.random.PRNGKey(0)

def normal_sample(key, loc, scale):
    return loc + scale*jax.random.normal(key)

normal_p = jax.core.Primitive("normal")

def normal(loc, scale):
    return normal_p.bind(loc, scale)

@normal_p.def_abstract_eval
def abstract_eval(loc, scale):
    print(f"abstract eval: {loc=} {scale=}")
    out_shape = jax.eval_shape(normal_sample, base_key, loc, scale)
    return abstract_arrays.ShapedArray(out_shape.shape, out_shape.dtype)

def batch_eval(vector_arg_values, batch_axes):
    raise NotImplementedError("doesn't matter")
batching.primitive_batchers[normal_p] = batch_eval

jaxpr = jax.make_jaxpr(jax.vmap(normal,[None,None],axis_size=3))(1.0,2.0)
print(f'{jaxpr=}')

This produces this output:

abstract eval: loc=ShapedArray(float32[], weak_type=True) scale=ShapedArray(float32[], weak_type=True)
jaxpr={ lambda ; a:f32[] b:f32[]. let
    c:f32[] = normal a b
    d:f32[3] = broadcast_in_dim[broadcast_dimensions=() shape=(3,)] c
  in (d,) }

The abstract eval is only called with scalar inputs, so I don't understand how batching logic could live there. Basically, JAX seems to decide on its own to call normal once and make 3 copies of it, and it does that without sending any signals to my code that a vmap even happened.

[jakevdp]

Agreed – my hope was that you could treat the abstract eval rule as the default un-batched case. Will that not work for you?