Out Axes Specification when returning NamedTuples

[juliandwain]

Hi!
I am trying to specify a function which should construct a graph using jraph's GraphsTuple class.
My input data is a python dictionary containing information about cells and features, i.e., it looks approx. like this

inputs = {
  "cells": Array of shape (*, 3),
  "mesh_pos": Array of shape (*,2),
  "node_type": Array of shape (*,1),
  "feature1": Array of shape (batch_size, tsteps, *, 2),
  "feature2": Array of shape (batch_size, tsteps, *, 2)
}

The cells, mesh_pos, and node_type arrays are "static", i.e., the do not change over batch_size and tsteps.
I have a function which acts on feature*'s last dimensions, i.e., for each batch and each tsteps.

I wanted to use JAX' vmap function for this. Now my code looks approx. like this:

def _build(inputs: dict):
  # some fancy stuff happens
  # build a graph tuple
  graph = jraph.GraphsTuple(
    nodes=nodes,
    edges=edges,
    receivers=receivers,
    senders=senders,
    globals=None,
    n_node=n_node,
    n_edge=n_edge

  return graph

in_axes = {
  "cells": None,
  "mesh_pos": None,
  "node_type": None,
  "feature1": 0,
  "feature2": 0,

build = jax.vmap(jax.vmap(_build, in_axes=in_axes), in_axes=in_axes)

The above code does exactly what I want, i.e., batching only over feature1 and feature2.
Yet, I want the output to be a GraphTuple in which only nodes and edges are mapped over the batch and tsteps dimension, as the other parameters do not change.

Is this somehow possible to specify using out_axes?
I tried to pass out_axes = (0, None, None, None, None, None, None) and out_axes = (0, 0, None, None, None, None, None, None) but neither of them worked.
I searched on the internet for solutions on how to use NamedTuple outputs within out_axes but I was not able to find a thing.

I am happy for every help I get and appreciate every single contribution!

Thanks in advance!

Best, J

[jakevdp]

The out axes specification should have the same tree structure as the outputs. This means that if you're returning a GraphsTuple, then the out_axes specification should also be a GraphsTuple, so something like this:

out_axes = jraph.GraphsTuple(0, None, None, None, None, None, None)

This is sometimes problematic for custom pytree implementations that do some sort of validation on their inputs; I don't know the best way to get around that issue.

[juliandwain]

Thank you very much, that helped my issue!