jax.vmap on function operating on a jraph graph tuple

[hsharsh]

So, I am trying to write a function to implement a Top-k pooling layer in Jraph, but I am running into issues with vmap. It definitely has to do how Masking+Jit does not really work with dynamic arrays, but I was wondering if there is any workaround to do the same thing. Here's a minimal example of what I want to do.

Starting from this graph:

I would like to truncate the graph based on the nodes, I want to keep. Let's say I would like to keep (0, 1, 3), then the output graph would be like this:

Here's the code to do it for a graph:

import jax
import jraph
import jax.numpy as jnp

def filter_graph(graph, selected_indices):
    senders = graph.senders
    receivers = graph.receivers
    num_nodes = graph.nodes.shape[0]

    # Compute which edges need to be removed
    new_node_index = jnp.arange(selected_indices.shape[0])
    mask = -1 * jnp.ones((num_nodes,))
    mask = mask.at[selected_indices].set(new_node_index)
    new_senders, new_receivers = mask[senders], mask[receivers]

    # These two lines are the problem
    mask_indices = jnp.nonzero((new_senders >= 0) & (new_receivers >= 0))[0]
    new_senders, new_receivers = jnp.take(new_senders, mask_indices), jnp.take(new_receivers, mask_indices)

    # Set the new nodes as the selected nodes
    new_nodes = graph.nodes[selected_indices]

    # Return the filtered graph
    outgraph = jraph.GraphsTuple(
        nodes = new_nodes,
        edges = None,
        senders = new_senders,
        receivers = new_receivers,
        n_node = jnp.array([new_nodes.shape[0]]),
        n_edge = jnp.array([new_senders.shape[0]]),
        globals = None
    )
    
    return outgraph

graph = jraph.GraphsTuple(
    nodes = jnp.array([[0.], [2.], [4.], [6.]]),
    edges = None,
    senders = jnp.array([0, 1, 2, 3, 0]),
    receivers = jnp.array([1, 2, 0, 0, 3]),
    n_node = jnp.array([4]),
    n_edge = jnp.array([5]),
    globals = None
)

selected_indices = jnp.array([0, 1, 3])

new_graph = filter_graph(graph, selected_indices)

This works fine but the problem is when I try vmap this to work for a batch of graphs, this won't work:

filter_graph_vmap = jax.vmap(filter_graph, in_axes = 0, out_axes = 0)

batched_graph = jraph.GraphsTuple(
    nodes = jnp.array([[[0.], [2.], [4.], [6.]],[[0.], [2.], [4.], [6.]]]),
    edges = None,
    senders = jnp.array([[0, 1, 2, 3, 0],[0, 1, 2, 3, 0]]),
    receivers = jnp.array([[1, 2, 0, 0, 3],[1, 2, 0, 0, 3]]),
    n_node = jnp.array([[4],[4]]),
    n_edge = jnp.array([[5],[5]]),
    globals = None
)

selected_indices = jnp.array([[0, 1, 3],[0, 1, 3]])

filter_graph_vmap(batched_graph, selected_indices)

Any ideas on how to work around this?

[jakevdp]

The issue here is that you're attempting to create a dynamically-shaped array (see JAX sharp bits: dynamic shapes).

The semantics of jnp.nonzero are that it returns an array whose length is the number of nonzero elements in its argument; this length cannot be known at compile time, and so using it in this way is not compatible with transformations like jit and vmap which require all arrays to be statically-shaped.

The solution would be to re-express your algorithm in a way that uses statically-shaped arrays. For example, if the number of nonzero entries is known statically, you can pass that to the size argument of jnp.nonzero. I'm not sure whether this is relevant in your case: I'm not familiar with jraph or how it represents graphs. But hopefully this is enough of a hint that you can arrive at your own solution.

[hsharsh]

Yes, I completely agree. The problem here is that there is no way to define size statically in this scenario. The algorithm (and any reformulation of the same) are dynamic by design.
I was just wondering if this is any more progress on this front which would be very beneficial to what I want to do. Otherwise, it's all good.

[jakevdp]

I would check with the jraph support forums to see how folks address this. There is probably a way to pad graph representations to keep them statically-shaped.