Alternatives to conditional recursion (for graph traversal)?

[ttt733]

Hi there! I'm trying to use Jax (and Haiku) to write a network that can navigate a very large (di)graph. It starts at the root node, looks at features of the connected edges, and assigns a probability to each edge that it wants to travel along it. My first attempt at this was to do it recursively.

## Graph-related stuff; not specifically relevant to my issue with JAX (but might help you understand the situation)
# graph = really big ndarray, padded to max size, containing the edges connected to each node. (dimensions = max_nodes x max_edges_per_node x edge_features + 1)
# search_mixer = haiku network that takes in edge features and outputs a softmaxed probability for traveling along each edge

def search_nodes(current_node_idx, paths_remaining):

   ## TL;DR: transform the graph's edges at this node into a set of target nodes and the number of paths to send to them.
   ## Question-relevant code is down in the iteration
   current_node = graph[current_node_idx][:,1:] # Position 0 stores the target node's index, hence the +1 in the graph dimensions
   edge_probabilities = search_mixer(current_node, train=is_training)
   path_distribution = edge_probabilities * paths
   path_distribution = path_distribution.astype('int32')
   # Transform the path distribution and connect the target node indexes for easier iteration
   # i.e. we want indexed_paths = (target_index, path_count)
   indexed_paths = jnp.stack((graph[current_node_idx][:,0].astype(int), path_distribution))
   indexed_paths = jnp.transpose(indexed_paths)

   ## --- Relevant stuff here ---
   for p in indexed_paths:
        target_index = p[0]
        path_count = p[1]
        # If 1 or more paths have been assigned to the edge, recurse at its target node
        r = jax.lax.cond(
            path_count > 0,
            lambda _: search_nodes(top_k, target_index, path_count, search_keys),
            lambda _: jnp.zeros(1),
            operand=None,
        )
       # ... Do some more stuff to compile the results of our search and then return

(There's more to my search algorithm than that, but I tried to simplify it as much as possible.) This basic approach didn't work, though, because I pretty quickly ran into this: RecursionError: maximum recursion depth exceeded while calling a Python object. I searched a bit and found this comment explaining it: #4858 (comment)

I thought about opening a feature request for what's discussed there, but the fact that it's not been talked about again since then tells me that I'm on the wrong track anyhow, and that maybe there's a way to do this already. But I'm kind of at a loss at how to continue. When I'm writing "normal" code, I know I can rewrite recursion as iteration by maintaining my own stack. I'm not sure how I could translate that to JAX, though. Has anyone run into anything similar, or do I just need to open a feature request and try a different library in the meantime? (Making that kind of contribution myself is way over my head right now, unfortunately.)

[ttt733]

Stared at this for a while and realized that there was a fairly simple solution that skirts the issue - just impose a maximum depth on the recursion (and the search.) Instead of calling the same function recursively, I'm now doing something like this:

def build_unrolled_search_fn(graph, child_fn):
    def search_fn(node_idx, search_keys):
        if child_fn is None:
           return 0 # Return whatever you're searching for
        # Execute the search in here to get next node index
        next_node_idx = 0
        result = child_fn(node_idx, search_keys)
        # Do whatever result transformations are necessary, then...
        return result
   return search_fn

# Build up to the highest-level search function by passing children along
max_depth = 4
child_search_fn = build_unrolled_search_fn(graph, None)
search_fn = None
for i in range(max_depth):
    search_fn = build_unrolled_search_fn(child_search_fn)
    child_search_fn = search_fn

# Now we can call search_fn(...) to execute the search!

Now, whether or not this is the best approach or will even play nicely with other JAX stuff like pmap, I haven't had a chance to see yet - but at least at low maximum depths, it seems to work. Still open to hearing any other ideas on the subject, but figured I would leave a note about this as a possibility in case anyone else stumbled across this while working on the same issue.