Error with bisection + vmap

[SNMS95]

Hey Patrick, thank you for the very nice package(s). I am extensively using equinox and optimistix in my current project for meta-learning.

When I try to take the gradient through the solver after applying vmap, I get the error Differentiation rule for 'unvmap_max' not implemented. I am trying to isolate the issue and get a minimum working example but it is taking a bit too long and so I am just posting it here to see if I am making some simple mistake. The pseudo code looks something like

def single_example_outer_loss(task_params, single_task_data):
    ....
    # call bisection to restrict the model output to have a particular mean
    return loss
    
    
vmapped_outer_loss = eqx.filter_vmap(single_example_outer_loss)

def average_meta_loss(task_params_batched, multi_task_data):
       return eqx.filter_value_and_grad(vmapped_outer_loss)(task_params_batched, multi_task_data)

I am trying to isolate the issue and provide a MWE

[patrick-kidger]

Hmm! That does seem a bit mysterious. I recognise unvmap_max as being a primitive we implement in Equinox, and one that does not have a differentiation rule. But I believe it's also only called with integer arrays, so it can't be differentiated anyway!

It's possible that the new 'stackless' JAX changes are biting you here. If the error goes away on JAX 0.4.35, but appears in JAX 0.4.36, then that'd be a likely culprit.

Anyway if you are able to provide a self-contained MWE then from that this is probably an easy one for me to debug :)

[SNMS95]

Thanks for the reply!
I am using jax 0.4.35 atm.
I will try to provide a MWE. This is a slippery one that is not being easy to isolate and reproduce!.

[SNMS95]

Hi @patrick-kidger. I managed to get a MWE (sort of minimal!).
And the issue seems to be in the config tag for disabling jit.

import equinox as eqx
import jax.numpy as jnp
import jax
import optimistix as optx

jax.config.update("jax_disable_jit", True)  # This is what causes the error
jax.config.update("jax_enable_x64", True)

def hoyer_volume_enforcing_filter(design, target_vf):
    """Enforces the mean of the design to be equal to the target volume fraction."""
    def transform_design(design, b):
        return jax.nn.sigmoid(design + b)

    def volume_violation_fn(y, args):
        return transform_design(args, y).mean() - target_vf

    b = optx.root_find(
        fn=volume_violation_fn,
        solver=optx.Bisection(atol=1e-5, rtol=1e-5),
        y0=100,
        args=design,
        options=dict(lower=-100, upper=100),
        max_steps=128,
        throw=True,
    )
    return transform_design(design, b.value)

def design_param_mock(input_coords, sample, params, static):
    """A NN + the enfocing filter."""
    model = eqx.combine(params, static)
    rho = jax.vmap(model)(input_coords)
    rho = hoyer_volume_enforcing_filter(
        rho, sample['target_vf'])  # Enforce volume
    return rho.reshape(32, 32)

X, Y = jnp.meshgrid(
    jnp.linspace(0, 1, 32),
    jnp.linspace(0, 1, 32),
    indexing="ij")
input_coords = jnp.stack([X.ravel(), Y.ravel()], axis=-1)
coords = input_coords

def outer_loss_fn(params, static, sample):
    """Outer loss function.
    """
    rho = design_param_mock(coords, sample, params, static)
    compliance = jnp.sum(rho**2)
    loss = compliance
    return loss, None


def get_direct_grad(adapted_task_params, task_static, validation_batch):

    vmapped_outer_loss = eqx.filter_vmap(
        outer_loss_fn, in_axes=(0, None, 0))

    def scalar_meta_loss(task_params):
        """Evaluates the meta-loss.
        """
        meta_losses, meta_metrics = vmapped_outer_loss(
            task_params, task_static, validation_batch)
        return jnp.sum(jnp.mean(meta_losses, axis=0)), None

    _ = jax.value_and_grad(scalar_meta_loss, has_aux=True,
                           argnums=(0))(adapted_task_params)
    return _
    
# TEST
sample = {
    'target_vf': 0.3,
    'id': 0}
samples = jax.tree.map(
    lambda x: jnp.expand_dims(x, axis=0), sample)

# Create a simple MLP model
model = eqx.nn.MLP(
    in_size=2,
    width_size=5,
    depth=1,
    out_size=1,
    activation=jnp.tanh,
    key=jax.random.PRNGKey(0),
)
# Initial model - without batch dimension
task_params, task_static = eqx.partition(
    model, eqx.is_array)
# adapted model - with batch dimension
adapted_params = jax.tree.map(lambda x: x[jnp.newaxis, ...], task_params)
# Loss
grad_single = get_direct_grad(
    adapted_params, task_static, samples)

[patrick-kidger]

Thanks for the MWE! This is exactly what I needed, and I've traced this to what looks like a bug in upstream JAX: jax-ml/jax#25724

(And as a workaround in the mean time you could try defining a dummy JVP rule for unvmap_max_p, which is defined in Equinox.)

[patrick-kidger]

Based on conversation in jax-ml/jax#25724 (comment) -- it looks like things are fixed (at least for your MWE) in JAX 0.5.0.

Can you give it a try on your actual use-case and see what happens?

[SNMS95]

Hey Patrick,

It does seem to be fixed now!
Thanks for helping out with this!