compose.utils

class ApplyFn(*args, **kwargs)

GNN apply function protocol.

batch_apply_fn(_apply_fn)

Combine vmapped graphs into disconnected supergraphs.

A mapped parameter leaf cannot be shared by one supergraph evaluation. In that case the affected map level is evaluated sequentially, so a model primitive is never asked to accept a parameter batch dimension.

The returned function supports arbitrary nesting of vmap and reverse-mode differentiation. JAX does not support forward-mode differentiation of custom_vjp functions, so applying jax.jvp to this wrapper raises TypeError.

Return type:

ApplyFn

flatten_graph(axis_size, in_batched, senders, receivers, edge_features, node_features)

Flatten one mapped graph axis into a disconnected supergraph.