Memory issues with grad and lax.scan

[vboussange]

Hey there,
I am looking to implement a differentiable least cost path algorithm, working with large graphs. I have been considering the Bellman-Ford algorithm, as I am especially interested in having one to all vertex distances. Here is a custom implementation, that seems to work beautifully:

import jax
import jax.numpy as jnp
from jax import lax, ops
import networkx as nx
from jax.experimental.sparse import BCOO
from equinox import filter_jit, filter_grad, filter_checkpoint

def bellman_ford(W_data, W_indices, source, nb_nodes):
    D = jnp.full(nb_nodes, jnp.inf).at[source].set(0.0)
    
    @filter_checkpoint
    def body_fun(D, _):
        D_u_plus_w = D[W_indices[:, 0]] + W_data
        D_v_min = ops.segment_min(D_u_plus_w, W_indices[:, 1], num_segments=nb_nodes)
        return jnp.minimum(D, D_v_min), None

    D, _ = lax.scan(body_fun, D, None, length=nb_nodes - 1)
    return D

N = 500
G = nx.grid_2d_graph(N, N, create_using=nx.DiGraph)
adj = BCOO.from_scipy_sparse(nx.adjacency_matrix(G))
nb_nodes = adj.shape[0]

device = jax.devices("gpu")[0]
W_data, W_indices = jax.device_put(adj.data.astype(jnp.float32), device), jax.device_put(adj.indices, device)

jit_bellman_ford = filter_jit(bellman_ford)
dist_to_node = jit_bellman_ford(W_data, W_indices, 0, nb_nodes)  # Forward pass works

However, when trying to differentiate bellman_ford with respect to edge weights, I encounter a RESOURCE_EXHAUSTED: Out of memory error.

grad_bellman = equinox.filter_grad(lambda *args: jnp.sum(bellman_ford(*args))) 
W_data_grad = grad_bellman(W_data, W_indices, 0, nb_nodes) # This fails with: XlaRuntimeError: RESOURCE_EXHAUSTED: Out of memory while ...

It seems that the @jax.checkpoint helps reduce memory usage, but it’s not sufficient to prevent memory build-up. I am actually suspecting a memory leak.

Am I doing something wrong here? Would there be a better implementation to avoid memory exhaustion?

Side note: I’m considering creating a JAX-based package for graph utilities (all JIT-compatible and differentiable). If anyone is interested, let me know!

Thanks in advance for your help!