Skip to content

Model¤

Connex separates a model into two pieces:

  • GraphSpec is the immutable graph description: the directed acyclic graph, ordered input nodes, ordered output nodes, topological order, and dropout configuration.
  • NeuralDAG is the trainable Equinox module compiled from that specification and an operation pipeline.

This split is intentional. Graph structure is static metadata from JAX's point of view, while trainable arrays live inside operation objects. The result works with eqx.filter_jit, eqx.filter_value_and_grad, eqx.apply_updates, and ordinary PyTree filtering.

Graph Specifications¤

GraphSpec accepts a networkx.DiGraph or anything accepted by networkx.DiGraph(...), including adjacency dictionaries and edge lists.

import connex as cnx
import networkx as nx

graph = nx.DiGraph()
graph.add_edges_from(
    [
        ("x0", "h0"),
        ("x1", "h0"),
        ("h0", "y"),
    ]
)

spec = cnx.GraphSpec(
    graph,
    inputs=["x0", "x1"],
    outputs=["y"],
    topo_sort=["x0", "x1", "h0", "y"],
    dropout={"h0": 0.1},
)

The input and output orders are semantic. A forward call places x[i] at spec.inputs[i], and returns outputs in spec.outputs order.

Connex validates the structure before parameters are initialized:

  • the graph must be acyclic;
  • input nodes must exist and have no incoming edges;
  • output nodes must exist and have no outgoing edges;
  • no node may be both an input and an output;
  • an explicit topo_sort, if supplied, must contain exactly the graph nodes and respect every edge;
  • dropout probabilities must be in [0, 1].

Supplying topo_sort is optional. It is useful when you want deterministic placement of isolated hidden nodes or when you already have a topological order from another graph construction step.

Runtime Model¤

NeuralDAG compiles the graph into integer node ids, edge-position tables, and topological batches. A topological batch is a group of target nodes whose predecessor values are already available. During a forward pass, Connex gathers only the input layouts requested by the active operations, runs the operation pipeline for each batch, and writes the new values back into a dense value buffer.

import jax
import jax.numpy as jnp
import jax.random as jr

model = cnx.NeuralDAG(
    spec,
    ops=cnx.ops.default_ops(activation=jax.nn.relu),
    key=jr.key(0),
)

y = model(jnp.array([1.0, -1.0]))

If ops is omitted, Connex uses connex.ops.default_ops(): a hybrid affine operation, hidden-node activation, dropout, and output transform. Custom operation sequences use the same NeuralDAG runtime.

Batching¤

Use model.batched(x) when every item in a batch can share the same runtime state. This is the native batched path and avoids mapping over the whole model.

x = jnp.ones((32, 2))
y = model.batched(x)

For independent stochastic state per example, split keys and map over both inputs and keys:

keys = jr.split(jr.key(1), x.shape[0])
y = jax.vmap(lambda x_i, key_i: model(x_i, key=key_i))(x, keys)

This distinction matters for dropout. model.batched(x, key=key) uses one runtime dropout state for the batch; jax.vmap(...) with split keys gives each example its own mask.

Execution Planning¤

The runtime squeezes structure out of the compiled topology before execution:

  • ordinary graph regions run one compiled topological batch at a time;
  • compatible one-input chain regions run as jax.lax.scan segments;
  • scan segments scatter only live-out values, meaning graph outputs and nodes consumed outside the segment;
  • operation input layouts are prepared only when some operation declares that it needs them;
  • the hybrid affine backend chooses padded rows, sparse edge accumulation, or dense matmul per topological batch.

Custom operations, attention features, normalization, and adaptive activation use the generic batch path. That path is more general because Connex cannot assume the local semantics required to fuse them into scan segments.

Exporting Weights¤

to_networkx_weighted_digraph() returns a copy of the original graph with learned edge weights attached when a built-in affine operation is present.

weighted = model.to_networkx_weighted_digraph()
weight = weighted["x0"]["h0"]["weight"]

This is useful for downstream NetworkX analysis, visualization, pruning experiments, or debugging topology edits.

Reference¤

connex.GraphSpec dataclass ¤

Validated, immutable graph definition for a Connex model.

A GraphSpec is the non-trainable description of a model: the graph, the ordered input and output nodes, the canonical topological order, and the dropout configuration. NeuralDAG compiles this object into static topology metadata and initializes trainable operation parameters from it.

graph may be a networkx.DiGraph or any object accepted by networkx.DiGraph(graph), including adjacency dictionaries and edge lists. Connex validates that the graph is a DAG, that input nodes have no incoming edges, that output nodes have no outgoing edges, and that no node appears in both inputs and outputs.

The order of inputs and outputs is semantically meaningful. Forward calls read x[i] into inputs[i] and return output values in outputs order.

__init__(graph: nx.DiGraph | Any, inputs: Sequence[Any], outputs: Sequence[Any], *, topo_sort: Sequence[Any] | None = None, dropout: DropoutLike = 0.0) ¤

Create a validated graph specification.

Arguments:

  • graph: A directed graph or graph-like object accepted by networkx.DiGraph. The resulting graph must be acyclic.
  • inputs: Ordered input node labels. These nodes must exist in the graph and must not receive incoming edges.
  • outputs: Ordered output node labels. These nodes must exist in the graph and must not have outgoing edges.
  • topo_sort: Optional topological order. Supplying this avoids a NetworkX topological sort during construction and gives stable ordering for isolated hidden nodes. The sequence must contain exactly the graph nodes and must respect every edge.
  • dropout: Either a scalar probability or a mapping from node label to probability. Scalar dropout applies to hidden nodes only. Mapping dropout defaults unspecified nodes to zero and can target inputs, hidden nodes, and outputs.

Raises:

  • ValueError: If the graph is cyclic, missing required nodes, has invalid input/output edge structure, has overlapping inputs/outputs, has an invalid topological order, or contains dropout probabilities outside [0, 1].
  • TypeError: If dropout is neither a scalar nor a mapping.
dropout_by_node() -> dict[Any, float] ¤

Return dropout probabilities keyed by graph node label.

Scalar dropout expands to every hidden node and leaves inputs and outputs at probability zero. Mapping dropout is returned with missing nodes filled in as zero. The output dictionary follows topo_sort ordering.

with_graph(graph: nx.DiGraph, *, inputs: Sequence[Any] | None = None, outputs: Sequence[Any] | None = None, topo_sort: Sequence[Any] | None = None, dropout: DropoutLike | None = None) -> GraphSpec ¤

Return a new GraphSpec with replacement graph metadata.

Any argument left as None is copied from the current specification. The returned object is fully revalidated, so this method is useful for small immutable graph transformations outside the editor API.

with_dropout(dropout: DropoutLike) -> GraphSpec ¤

Return a new GraphSpec with the same graph and new dropout.

connex.NeuralDAG ¤

Bases: Module

A trainable neural network runtime compiled from a DAG specification.

NeuralDAG is the core Connex model type. It is an equinox.Module whose trainable leaves live inside operation objects, while graph structure and compiled topology metadata are static. This makes models compatible with eqx.filter_jit, eqx.filter_value_and_grad, eqx.apply_updates, and the usual JAX transformations.

Forward evaluation proceeds in compiled topological order. Connex groups nodes into batches whose predecessors are already available, gathers the data requested by each operation, and writes newly computed node values into a value buffer. Compatible one-input chain regions are represented as jax.lax.scan execution-plan segments; unsupported operation stacks and custom operations use the generic topological batch path.

__init__(spec: GraphSpec, *, ops: Sequence[Any] | None = None, key: Array | None = None) ¤

Initialize a model from a graph specification.

Arguments:

  • spec: Validated graph, input/output ordering, topological order, and dropout configuration.
  • ops: Operation pipeline. If None, Connex uses connex.ops.default_ops(), which is a hybrid affine op, activation, dropout, and output transform.
  • key: JAX random key used to initialize operation parameters. If omitted, jax.random.key(0) is used.

The operation sequence is initialized once against the compiled topology. Each operation may create trainable arrays, static metadata, or no state at all.

__call__(x: Array, *, key: Array | None = None) -> Array ¤

Evaluate one input example.

Arguments:

  • x: Array whose trailing dimension matches len(spec.inputs). Values are assigned to input nodes in specification order.
  • key: Optional JAX random key used by runtime stochastic operations such as dropout. If any active dropout operation is present, a key is required.

Returns:

Output values in spec.outputs order after the output-stage operations have been applied.

batched(x: Array, *, key: Array | None = None) -> Array ¤

Evaluate a batch of examples with one shared runtime state.

This is the native batched path and avoids compiling jax.vmap(model) over the whole DAG. It is appropriate for deterministic evaluation or when a shared dropout mask is desired. For independent stochastic state per example, split keys and use jax.vmap(lambda x_i, key_i: model(x_i, key=key_i))(x, keys).

Arguments:

  • x: Array of shape (batch, len(spec.inputs)).
  • key: Optional runtime key shared by the whole batch.

Returns:

Array of shape (batch, len(spec.outputs)).

rebuild(spec: GraphSpec, *, key: Array | None = None, ops: Sequence[Any] | None = None) -> NeuralDAG ¤

Recompile a model against a new graph specification.

The new operation pipeline is initialized for spec, then each old operation is asked to transfer compatible parameters into the initialized replacement. Built-in affine operations preserve parameters for nodes and edges that still exist by label.

Arguments:

  • spec: New graph specification.
  • key: Random key for initializing newly created parameters.
  • ops: Optional replacement operation sequence. If omitted, the current operations are reused.

Returns:

A new NeuralDAG; the original model is unchanged.

to_networkx_weighted_digraph() -> nx.DiGraph ¤

Export the model graph with learned edge weights.

The returned networkx.DiGraph copies self.spec.graph and annotates each edge with a "weight" attribute when a built-in affine operation is present. Sparse, matmul, padded, and hybrid affine layouts are all mapped back to the original graph edge labels.