DenseMLP

connex.nn.DenseMLP builds a densely connected layered DAG. Each layer connects to every later layer, so later hidden nodes and outputs can reuse all earlier representations instead of only the immediately preceding layer.

import connex as cnx
import jax
import jax.random as jr

model = cnx.nn.DenseMLP(
    input_size=4,
    output_size=2,
    width=16,
    depth=3,
    activation=jax.nn.silu,
    key=jr.key(0),
)

Like MLP, this is a regular NeuralDAG subclass with generated integer node labels:

• inputs come first;
• hidden layers follow in order;
• outputs are the final nodes;
• every node in a non-final layer connects to all nodes in later layers.

This graph family is useful when you want skip connectivity without writing the graph by hand. It increases the number of edges compared with MLP, so it is a good candidate for the hybrid affine backend: compact regions use padded rows, sparse regions use edge accumulation, and wide complete regions can use dense matmul.

model = cnx.nn.DenseMLP(
    4,
    2,
    width=16,
    depth=3,
    ops=cnx.ops.default_ops(affine="hybrid", fused=True),
    key=jr.key(1),
)

You can still edit the generated topology. For example, remove a skip edge, add a new hidden node, or export the learned weights for graph analysis:

model = cnx.edit(model).remove_edges([(0, model.spec.outputs[0])]).build(key=jr.key(2))
weighted = model.to_networkx_weighted_digraph()

Reference

connex.nn.DenseMLP

Bases: NeuralDAG

A densely connected MLP represented as a Connex DAG.

DenseMLP is similar to MLP, but each layer connects to every later layer rather than only to the next layer. This gives later nodes access to all earlier representations, in the spirit of DenseNet-style skip connectivity, while still using the normal NeuralDAG runtime.

__init__(input_size: int, output_size: int, width: int, depth: int, *, activation: Callable = jnn.gelu, output_transform: Callable = _identity, dropout: DropoutLike = 0.0, ops: Sequence[cnx_ops.Op] | None = None, key: Array | None = None)

Create a densely connected layered graph.

Arguments:

• input_size: Number of input nodes.
• output_size: Number of output nodes.
• width: Number of nodes in each hidden layer.
• depth: Number of hidden layers.
• activation: Elementwise hidden-node activation used by the default op stack.
• output_transform: Final transform applied to ordered outputs.
• dropout: Scalar or mapping dropout configuration.
• ops: Optional custom operation sequence. If supplied, activation and output_transform are ignored unless your ops use them.
• key: Random key for parameter initialization.