Skip to content

Operations¤

Connex models are built from operation objects. An operation is an equinox.Module that can own trainable parameters, declare which graph inputs it needs, transform values during a forward pass, and transfer parameters across topology edits.

The default pipeline is produced by connex.ops.default_ops():

  1. optional topology features, such as normalization or attention;
  2. an affine operation;
  3. hidden-node activation or adaptive activation;
  4. dropout;
  5. an output transform.
import connex as cnx
import jax

ops = cnx.ops.default_ops(
    affine="hybrid",
    activation=jax.nn.gelu,
    topo_norm=True,
    dropout=0.1,
)

model = cnx.NeuralDAG(spec, ops=ops, key=key)

Operation templates are initialized by NeuralDAG against the compiled topology. You normally create lightweight templates, pass them into the model, and let the model allocate arrays with the supplied key.

Forward Context¤

Batch-stage operations receive a ForwardContext. It contains the current compiled topological batch, the full model value buffer, target ids, predecessor ids, edge-position metadata, and any outputs produced by earlier operations.

The context exposes three gathered input views:

  • batch_inputs: one value per unique predecessor in the topological batch;
  • input_values: padded predecessor rows shaped like target nodes;
  • edge_input_values: one value per real edge.

Operations declare which views they need through the needs_batch_inputs, needs_padded_inputs, and needs_edge_inputs flags. Keeping these flags narrow matters for performance because NeuralDAG avoids preparing unused layouts.

Affine Backends¤

default_ops() uses affine="hybrid" unless configured otherwise. Hybrid selects an affine implementation per topological batch:

  • padded rows for compact predecessor layouts;
  • sparse edge accumulation when padding would dominate;
  • dense matmul for wide complete predecessor batches.

Force affine="padded", affine="sparse", or affine="matmul" when benchmarking a known graph family. The scan execution plan supports all built-in affine backends when the remaining stack is affine, activation, optional dropout, and output transform, or the equivalent fused default op.

ops = cnx.ops.default_ops(affine="sparse", activation=jax.nn.relu)

EdgeAffine stores padded rows and is usually best for compact batches. SparseEdgeAffine stores one weight per real edge and avoids wasted padded slots. DenseMatmulAffine stores a matrix over unique predecessors and is best when a batch resembles a dense layer. HybridEdgeAffine makes that choice per batch.

Feature Operations¤

Feature operations run before the affine transform and modify the gathered predecessor values:

  • TopoNorm normalizes the unique predecessor set for each topological batch.
  • TopoSelfAttention attends over each batch's unique predecessor values.
  • NeuronSelfAttention attends separately over each target node's input row.

Activation operations run after affine outputs are available:

  • Activation applies a fixed activation to hidden nodes only.
  • AdaptiveActivation learns per-node activation scales.
  • Dropout masks node values with explicit JAX random keys.
  • OutputTransform transforms the final ordered output array.

Fused Defaults¤

default_ops(fused=True) uses FusedDefaultOp when no feature operations or adaptive activations are requested. This combines affine, activation, and dropout into one operation and keeps the model eligible for the optimized scan execution plan on one-input chain regions.

ops = cnx.ops.default_ops(fused=True, affine="hybrid")

Custom Operations¤

Subclass Op when you need behavior not covered by the built-ins. A custom op can read the context, return an updated context, own trainable arrays, and define transfer behavior for topology edits.

from dataclasses import replace

import equinox as eqx
import jax
import jax.numpy as jnp


class BiasAfterAffine(cnx.ops.Op):
    bias: jax.Array
    name: str = eqx.field(static=True, default="bias_after_affine")
    needs_batch_inputs: bool = eqx.field(static=True, default=False)
    needs_padded_inputs: bool = eqx.field(static=True, default=False)
    needs_edge_inputs: bool = eqx.field(static=True, default=False)

    def apply(self, ctx, *, state=None, key=None):
        assert ctx.outputs is not None
        return replace(ctx, outputs=ctx.outputs + self.bias)


ops = (
    cnx.ops.EdgeAffine(),
    BiasAfterAffine(bias=jnp.asarray(0.1)),
    cnx.ops.Activation(activation=jax.nn.relu),
    cnx.ops.OutputTransform(),
)

Custom operations use the generic topological batch execution path. That path is the right default for arbitrary user code because Connex cannot infer the same scan and fusion guarantees it has for the built-in affine/activation/dropout stack.

Reference¤

connex.ops.Op ¤

Bases: Module

Base class for user-definable Connex operations.

Operations are small Equinox modules that participate in a NeuralDAG forward pass. They can own trainable arrays, expose static metadata, create runtime state such as dropout masks, and transfer parameters across topology edits.

The main extension point is apply(ctx, ...), which receives a ForwardContext for one topological batch and returns an updated context. Operations that need to touch input nodes before any hidden/output nodes are computed can override apply_inputs. Operations that transform the final model output can set stage="output" and override apply_output.

The needs_* flags tell NeuralDAG which gathered inputs to prepare. Set these narrowly for performance: a custom operation that only reads ctx.outputs should set all three to False.

name: str = eqx.field(static=True, default='op') class-attribute instance-attribute ¤
layout: ParamLayout = eqx.field(static=True, default=ParamLayout.GLOBAL) class-attribute instance-attribute ¤
stage: str = eqx.field(static=True, default='batch') class-attribute instance-attribute ¤
needs_batch_inputs: bool = eqx.field(static=True, default=True) class-attribute instance-attribute ¤
needs_padded_inputs: bool = eqx.field(static=True, default=True) class-attribute instance-attribute ¤
needs_edge_inputs: bool = eqx.field(static=True, default=True) class-attribute instance-attribute ¤
init(topology: CompiledTopology, *, key: Array) -> Op ¤
init_runtime(topology: CompiledTopology, *, key: Array | None) -> Any ¤
apply(ctx: ForwardContext, *, state: Any = None, key: Array | None = None) -> ForwardContext ¤
apply_inputs(topology: CompiledTopology, values: Array, *, state: Any = None, key: Array | None = None) -> Array ¤
apply_output(topology: CompiledTopology, y: Array, *, state: Any = None, key: Array | None = None) -> Array ¤
needs_batch_inputs_for_batch(batch_index: int) -> bool ¤
needs_padded_inputs_for_batch(batch_index: int) -> bool ¤
needs_edge_inputs_for_batch(batch_index: int) -> bool ¤
transfer(old_topology: CompiledTopology, new_topology: CompiledTopology, initialized: Op) -> Op ¤

connex.ops.EdgeAffine ¤

Bases: Op

Padded edge-wise affine transform.

Each topological batch stores weights shaped (batch.size, batch.max_inputs) plus one bias per target node. Missing padded predecessor slots are masked by the compiled topology. This layout is simple and fast for compact batches where every target has a similar number of predecessors.

weights: tuple[Array, ...] = () class-attribute instance-attribute ¤
biases: tuple[Array, ...] = () class-attribute instance-attribute ¤
weight_scale: float = eqx.field(static=True, default=0.1) class-attribute instance-attribute ¤
name: str = eqx.field(static=True, default='edge_affine') class-attribute instance-attribute ¤
layout: ParamLayout = eqx.field(static=True, default=ParamLayout.EDGE) class-attribute instance-attribute ¤
stage: str = eqx.field(static=True, default='batch') class-attribute instance-attribute ¤
needs_batch_inputs: bool = eqx.field(static=True, default=False) class-attribute instance-attribute ¤
needs_padded_inputs: bool = eqx.field(static=True, default=True) class-attribute instance-attribute ¤
needs_edge_inputs: bool = eqx.field(static=True, default=False) class-attribute instance-attribute ¤
init(topology: CompiledTopology, *, key: Array) -> EdgeAffine ¤
apply(ctx: ForwardContext, *, state: Any = None, key: Array | None = None) -> ForwardContext ¤
transfer(old_topology: CompiledTopology, new_topology: CompiledTopology, initialized: EdgeAffine) -> EdgeAffine ¤
init_runtime(topology: CompiledTopology, *, key: Array | None) -> Any ¤
apply_inputs(topology: CompiledTopology, values: Array, *, state: Any = None, key: Array | None = None) -> Array ¤
apply_output(topology: CompiledTopology, y: Array, *, state: Any = None, key: Array | None = None) -> Array ¤
needs_batch_inputs_for_batch(batch_index: int) -> bool ¤
needs_padded_inputs_for_batch(batch_index: int) -> bool ¤
needs_edge_inputs_for_batch(batch_index: int) -> bool ¤

connex.ops.SparseEdgeAffine ¤

Bases: Op

Sparse edge-wise affine transform.

Each topological batch stores one weight per real edge and scatters edge contributions into target-node outputs. This avoids padded storage and work when a batch has highly uneven fan-in.

weights: tuple[Array, ...] = () class-attribute instance-attribute ¤
biases: tuple[Array, ...] = () class-attribute instance-attribute ¤
weight_scale: float = eqx.field(static=True, default=0.1) class-attribute instance-attribute ¤
name: str = eqx.field(static=True, default='sparse_edge_affine') class-attribute instance-attribute ¤
layout: ParamLayout = eqx.field(static=True, default=ParamLayout.EDGE) class-attribute instance-attribute ¤
stage: str = eqx.field(static=True, default='batch') class-attribute instance-attribute ¤
needs_batch_inputs: bool = eqx.field(static=True, default=False) class-attribute instance-attribute ¤
needs_padded_inputs: bool = eqx.field(static=True, default=False) class-attribute instance-attribute ¤
needs_edge_inputs: bool = eqx.field(static=True, default=True) class-attribute instance-attribute ¤
init(topology: CompiledTopology, *, key: Array) -> SparseEdgeAffine ¤
apply(ctx: ForwardContext, *, state: Any = None, key: Array | None = None) -> ForwardContext ¤
transfer(old_topology: CompiledTopology, new_topology: CompiledTopology, initialized: SparseEdgeAffine) -> SparseEdgeAffine ¤
init_runtime(topology: CompiledTopology, *, key: Array | None) -> Any ¤
apply_inputs(topology: CompiledTopology, values: Array, *, state: Any = None, key: Array | None = None) -> Array ¤
apply_output(topology: CompiledTopology, y: Array, *, state: Any = None, key: Array | None = None) -> Array ¤
needs_batch_inputs_for_batch(batch_index: int) -> bool ¤
needs_padded_inputs_for_batch(batch_index: int) -> bool ¤
needs_edge_inputs_for_batch(batch_index: int) -> bool ¤

connex.ops.DenseMatmulAffine ¤

Bases: Op

Dense matrix affine transform over unique predecessor values.

Each topological batch stores a matrix shaped (batch.size, batch.unique_size) and multiplies it by the unique predecessor vector. This backend is useful for wide, dense predecessor layouts where the graph already resembles a layer-to-layer matrix multiply.

weights: tuple[Array, ...] = () class-attribute instance-attribute ¤
biases: tuple[Array, ...] = () class-attribute instance-attribute ¤
weight_scale: float = eqx.field(static=True, default=0.1) class-attribute instance-attribute ¤
name: str = eqx.field(static=True, default='dense_matmul_affine') class-attribute instance-attribute ¤
layout: ParamLayout = eqx.field(static=True, default=ParamLayout.EDGE) class-attribute instance-attribute ¤
stage: str = eqx.field(static=True, default='batch') class-attribute instance-attribute ¤
needs_batch_inputs: bool = eqx.field(static=True, default=True) class-attribute instance-attribute ¤
needs_padded_inputs: bool = eqx.field(static=True, default=False) class-attribute instance-attribute ¤
needs_edge_inputs: bool = eqx.field(static=True, default=False) class-attribute instance-attribute ¤
init(topology: CompiledTopology, *, key: Array) -> DenseMatmulAffine ¤
apply(ctx: ForwardContext, *, state: Any = None, key: Array | None = None) -> ForwardContext ¤
transfer(old_topology: CompiledTopology, new_topology: CompiledTopology, initialized: DenseMatmulAffine) -> DenseMatmulAffine ¤
init_runtime(topology: CompiledTopology, *, key: Array | None) -> Any ¤
apply_inputs(topology: CompiledTopology, values: Array, *, state: Any = None, key: Array | None = None) -> Array ¤
apply_output(topology: CompiledTopology, y: Array, *, state: Any = None, key: Array | None = None) -> Array ¤
needs_batch_inputs_for_batch(batch_index: int) -> bool ¤
needs_padded_inputs_for_batch(batch_index: int) -> bool ¤
needs_edge_inputs_for_batch(batch_index: int) -> bool ¤

connex.ops.HybridEdgeAffine ¤

Bases: Op

Per-batch affine backend selector.

Hybrid affine chooses padded, sparse, or dense matmul storage separately for each compiled topological batch. Sparse mode is selected when padded slots would dominate real edges. Matmul mode is selected for wide complete predecessor batches. Otherwise the padded layout is used.

weights: tuple[Array, ...] = () class-attribute instance-attribute ¤
biases: tuple[Array, ...] = () class-attribute instance-attribute ¤
sparse_batches: tuple[bool, ...] = eqx.field(static=True, default=()) class-attribute instance-attribute ¤
matmul_batches: tuple[bool, ...] = eqx.field(static=True, default=()) class-attribute instance-attribute ¤
padding_ratio_threshold: float = eqx.field(static=True, default=2.0) class-attribute instance-attribute ¤
weight_scale: float = eqx.field(static=True, default=0.1) class-attribute instance-attribute ¤
name: str = eqx.field(static=True, default='hybrid_edge_affine') class-attribute instance-attribute ¤
layout: ParamLayout = eqx.field(static=True, default=ParamLayout.EDGE) class-attribute instance-attribute ¤
stage: str = eqx.field(static=True, default='batch') class-attribute instance-attribute ¤
needs_batch_inputs: bool = eqx.field(static=True, default=False) class-attribute instance-attribute ¤
needs_padded_inputs: bool = eqx.field(static=True, default=True) class-attribute instance-attribute ¤
needs_edge_inputs: bool = eqx.field(static=True, default=True) class-attribute instance-attribute ¤
init(topology: CompiledTopology, *, key: Array) -> HybridEdgeAffine ¤
needs_batch_inputs_for_batch(batch_index: int) -> bool ¤
needs_padded_inputs_for_batch(batch_index: int) -> bool ¤
needs_edge_inputs_for_batch(batch_index: int) -> bool ¤
apply(ctx: ForwardContext, *, state: Any = None, key: Array | None = None) -> ForwardContext ¤
transfer(old_topology: CompiledTopology, new_topology: CompiledTopology, initialized: HybridEdgeAffine) -> HybridEdgeAffine ¤
init_runtime(topology: CompiledTopology, *, key: Array | None) -> Any ¤
apply_inputs(topology: CompiledTopology, values: Array, *, state: Any = None, key: Array | None = None) -> Array ¤
apply_output(topology: CompiledTopology, y: Array, *, state: Any = None, key: Array | None = None) -> Array ¤

connex.ops.TopoNorm ¤

Bases: Op

Normalize the unique inputs of each topological batch.

This is a graph/topology analogue of layer normalization. For every topological batch, Connex standardizes the batch's unique predecessor values and applies trainable elementwise affine parameters gamma and beta.

params: tuple[Array, ...] = () class-attribute instance-attribute ¤
name: str = eqx.field(static=True, default='topo_norm') class-attribute instance-attribute ¤
layout: ParamLayout = eqx.field(static=True, default=ParamLayout.TOPO_INPUT) class-attribute instance-attribute ¤
stage: str = eqx.field(static=True, default='batch') class-attribute instance-attribute ¤
needs_batch_inputs: bool = eqx.field(static=True, default=True) class-attribute instance-attribute ¤
needs_padded_inputs: bool = eqx.field(static=True, default=False) class-attribute instance-attribute ¤
needs_edge_inputs: bool = eqx.field(static=True, default=False) class-attribute instance-attribute ¤
init(topology: CompiledTopology, *, key: Array) -> TopoNorm ¤
apply(ctx: ForwardContext, *, state: Any = None, key: Array | None = None) -> ForwardContext ¤
transfer(old_topology: CompiledTopology, new_topology: CompiledTopology, initialized: TopoNorm) -> TopoNorm ¤
init_runtime(topology: CompiledTopology, *, key: Array | None) -> Any ¤
apply_inputs(topology: CompiledTopology, values: Array, *, state: Any = None, key: Array | None = None) -> Array ¤
apply_output(topology: CompiledTopology, y: Array, *, state: Any = None, key: Array | None = None) -> Array ¤
needs_batch_inputs_for_batch(batch_index: int) -> bool ¤
needs_padded_inputs_for_batch(batch_index: int) -> bool ¤
needs_edge_inputs_for_batch(batch_index: int) -> bool ¤

connex.ops.TopoSelfAttention ¤

Bases: Op

Single-headed self-attention over a batch's unique predecessor values.

This operation lets a topological batch mix information across all of its available inputs before the affine transform. It is cheaper than neuron-level attention because it attends once over the collective input set for the batch.

params: tuple[Array, ...] = () class-attribute instance-attribute ¤
weight_scale: float = eqx.field(static=True, default=0.1) class-attribute instance-attribute ¤
name: str = eqx.field(static=True, default='topo_self_attention') class-attribute instance-attribute ¤
layout: ParamLayout = eqx.field(static=True, default=ParamLayout.TOPO_INPUT) class-attribute instance-attribute ¤
stage: str = eqx.field(static=True, default='batch') class-attribute instance-attribute ¤
needs_batch_inputs: bool = eqx.field(static=True, default=True) class-attribute instance-attribute ¤
needs_padded_inputs: bool = eqx.field(static=True, default=False) class-attribute instance-attribute ¤
needs_edge_inputs: bool = eqx.field(static=True, default=False) class-attribute instance-attribute ¤
init(topology: CompiledTopology, *, key: Array) -> TopoSelfAttention ¤
apply(ctx: ForwardContext, *, state: Any = None, key: Array | None = None) -> ForwardContext ¤
transfer(old_topology: CompiledTopology, new_topology: CompiledTopology, initialized: TopoSelfAttention) -> TopoSelfAttention ¤
init_runtime(topology: CompiledTopology, *, key: Array | None) -> Any ¤
apply_inputs(topology: CompiledTopology, values: Array, *, state: Any = None, key: Array | None = None) -> Array ¤
apply_output(topology: CompiledTopology, y: Array, *, state: Any = None, key: Array | None = None) -> Array ¤
needs_batch_inputs_for_batch(batch_index: int) -> bool ¤
needs_padded_inputs_for_batch(batch_index: int) -> bool ¤
needs_edge_inputs_for_batch(batch_index: int) -> bool ¤

connex.ops.NeuronSelfAttention ¤

Bases: Op

Single-headed self-attention over each target neuron's input row.

Every target node attends over its own predecessor values before the affine transform. This is more expressive than TopoSelfAttention, but can use substantially more memory because attention parameters are stored per target node and padded input row.

params: tuple[Array, ...] = () class-attribute instance-attribute ¤
weight_scale: float = eqx.field(static=True, default=0.1) class-attribute instance-attribute ¤
name: str = eqx.field(static=True, default='neuron_self_attention') class-attribute instance-attribute ¤
layout: ParamLayout = eqx.field(static=True, default=ParamLayout.NODE) class-attribute instance-attribute ¤
stage: str = eqx.field(static=True, default='batch') class-attribute instance-attribute ¤
needs_batch_inputs: bool = eqx.field(static=True, default=False) class-attribute instance-attribute ¤
needs_padded_inputs: bool = eqx.field(static=True, default=True) class-attribute instance-attribute ¤
needs_edge_inputs: bool = eqx.field(static=True, default=False) class-attribute instance-attribute ¤
init(topology: CompiledTopology, *, key: Array) -> NeuronSelfAttention ¤
apply(ctx: ForwardContext, *, state: Any = None, key: Array | None = None) -> ForwardContext ¤
transfer(old_topology: CompiledTopology, new_topology: CompiledTopology, initialized: NeuronSelfAttention) -> NeuronSelfAttention ¤
init_runtime(topology: CompiledTopology, *, key: Array | None) -> Any ¤
apply_inputs(topology: CompiledTopology, values: Array, *, state: Any = None, key: Array | None = None) -> Array ¤
apply_output(topology: CompiledTopology, y: Array, *, state: Any = None, key: Array | None = None) -> Array ¤
needs_batch_inputs_for_batch(batch_index: int) -> bool ¤
needs_padded_inputs_for_batch(batch_index: int) -> bool ¤
needs_edge_inputs_for_batch(batch_index: int) -> bool ¤

connex.ops.Activation ¤

Bases: Op

Apply an activation to hidden-node outputs.

Output nodes are intentionally left unactivated here; use OutputTransform for group-wise output behavior such as softmax.

activation: Callable = _identity class-attribute instance-attribute ¤
name: str = eqx.field(static=True, default='activation') class-attribute instance-attribute ¤
layout: ParamLayout = eqx.field(static=True, default=ParamLayout.GLOBAL) class-attribute instance-attribute ¤
stage: str = eqx.field(static=True, default='batch') class-attribute instance-attribute ¤
needs_batch_inputs: bool = eqx.field(static=True, default=False) class-attribute instance-attribute ¤
needs_padded_inputs: bool = eqx.field(static=True, default=False) class-attribute instance-attribute ¤
needs_edge_inputs: bool = eqx.field(static=True, default=False) class-attribute instance-attribute ¤
apply(ctx: ForwardContext, *, state: Any = None, key: Array | None = None) -> ForwardContext ¤
transfer(old_topology: CompiledTopology, new_topology: CompiledTopology, initialized: Activation) -> Activation ¤
init(topology: CompiledTopology, *, key: Array) -> Op ¤
init_runtime(topology: CompiledTopology, *, key: Array | None) -> Any ¤
apply_inputs(topology: CompiledTopology, values: Array, *, state: Any = None, key: Array | None = None) -> Array ¤
apply_output(topology: CompiledTopology, y: Array, *, state: Any = None, key: Array | None = None) -> Array ¤
needs_batch_inputs_for_batch(batch_index: int) -> bool ¤
needs_padded_inputs_for_batch(batch_index: int) -> bool ¤
needs_edge_inputs_for_batch(batch_index: int) -> bool ¤

connex.ops.AdaptiveActivation ¤

Bases: Op

Trainable per-node activation scaling.

Hidden activations are transformed as a * activation(b * x), with trainable scalars a and b per target node. Output nodes are not transformed.

activation: Callable = _identity class-attribute instance-attribute ¤
params: tuple[Array, ...] = () class-attribute instance-attribute ¤
name: str = eqx.field(static=True, default='adaptive_activation') class-attribute instance-attribute ¤
layout: ParamLayout = eqx.field(static=True, default=ParamLayout.NODE) class-attribute instance-attribute ¤
stage: str = eqx.field(static=True, default='batch') class-attribute instance-attribute ¤
needs_batch_inputs: bool = eqx.field(static=True, default=False) class-attribute instance-attribute ¤
needs_padded_inputs: bool = eqx.field(static=True, default=False) class-attribute instance-attribute ¤
needs_edge_inputs: bool = eqx.field(static=True, default=False) class-attribute instance-attribute ¤
init(topology: CompiledTopology, *, key: Array) -> AdaptiveActivation ¤
apply(ctx: ForwardContext, *, state: Any = None, key: Array | None = None) -> ForwardContext ¤
transfer(old_topology: CompiledTopology, new_topology: CompiledTopology, initialized: AdaptiveActivation) -> AdaptiveActivation ¤
init_runtime(topology: CompiledTopology, *, key: Array | None) -> Any ¤
apply_inputs(topology: CompiledTopology, values: Array, *, state: Any = None, key: Array | None = None) -> Array ¤
apply_output(topology: CompiledTopology, y: Array, *, state: Any = None, key: Array | None = None) -> Array ¤
needs_batch_inputs_for_batch(batch_index: int) -> bool ¤
needs_padded_inputs_for_batch(batch_index: int) -> bool ¤
needs_edge_inputs_for_batch(batch_index: int) -> bool ¤

connex.ops.Dropout ¤

Bases: Op

Node-wise dropout with explicit runtime keys.

Dropout probabilities come from the op's own dropout field when supplied, otherwise from GraphSpec.dropout. If any probability is nonzero, forward calls must provide a JAX random key. This avoids hidden global randomness and makes stochastic behavior explicit under JIT and transformations.

dropout: Any = eqx.field(static=True, default=None) class-attribute instance-attribute ¤
dropout_probs: tuple[float, ...] = eqx.field(static=True, default=()) class-attribute instance-attribute ¤
active: bool = eqx.field(static=True, default=False) class-attribute instance-attribute ¤
name: str = eqx.field(static=True, default='dropout') class-attribute instance-attribute ¤
layout: ParamLayout = eqx.field(static=True, default=ParamLayout.GLOBAL) class-attribute instance-attribute ¤
stage: str = eqx.field(static=True, default='batch') class-attribute instance-attribute ¤
needs_batch_inputs: bool = eqx.field(static=True, default=False) class-attribute instance-attribute ¤
needs_padded_inputs: bool = eqx.field(static=True, default=False) class-attribute instance-attribute ¤
needs_edge_inputs: bool = eqx.field(static=True, default=False) class-attribute instance-attribute ¤
init(topology: CompiledTopology, *, key: Array) -> Dropout ¤
init_runtime(topology: CompiledTopology, *, key: Array | None) -> Array | None ¤
apply_inputs(topology: CompiledTopology, values: Array, *, state: Array | None = None, key: Array | None = None) -> Array ¤
apply(ctx: ForwardContext, *, state: Array | None = None, key: Array | None = None) -> ForwardContext ¤
with_dropout(dropout: DropoutLike) -> Dropout ¤
apply_output(topology: CompiledTopology, y: Array, *, state: Any = None, key: Array | None = None) -> Array ¤
needs_batch_inputs_for_batch(batch_index: int) -> bool ¤
needs_padded_inputs_for_batch(batch_index: int) -> bool ¤
needs_edge_inputs_for_batch(batch_index: int) -> bool ¤
transfer(old_topology: CompiledTopology, new_topology: CompiledTopology, initialized: Op) -> Op ¤

connex.ops.FusedDefaultOp ¤

Bases: Op

Fused affine, activation, and optional dropout operation.

This combines the common default stack into one batch-stage operation. It can reduce Python-level operation dispatch and is eligible for the same scan-chain execution plan as the equivalent unfused affine/activation/dropout stack when no extra feature operations are present.

affine: Op = eqx.field(default_factory=HybridEdgeAffine) class-attribute instance-attribute ¤
activation: Callable = _identity class-attribute instance-attribute ¤
dropout: Any = eqx.field(static=True, default=None) class-attribute instance-attribute ¤
dropout_probs: tuple[float, ...] = eqx.field(static=True, default=()) class-attribute instance-attribute ¤
active: bool = eqx.field(static=True, default=False) class-attribute instance-attribute ¤
name: str = eqx.field(static=True, default='fused_default') class-attribute instance-attribute ¤
layout: ParamLayout = eqx.field(static=True, default=ParamLayout.CUSTOM) class-attribute instance-attribute ¤
stage: str = eqx.field(static=True, default='batch') class-attribute instance-attribute ¤
needs_batch_inputs: bool = eqx.field(static=True, default=False) class-attribute instance-attribute ¤
needs_padded_inputs: bool = eqx.field(static=True, default=False) class-attribute instance-attribute ¤
needs_edge_inputs: bool = eqx.field(static=True, default=False) class-attribute instance-attribute ¤
init(topology: CompiledTopology, *, key: Array) -> FusedDefaultOp ¤
init_runtime(topology: CompiledTopology, *, key: Array | None) -> Array | None ¤
apply_inputs(topology: CompiledTopology, values: Array, *, state: Array | None = None, key: Array | None = None) -> Array ¤
needs_batch_inputs_for_batch(batch_index: int) -> bool ¤
needs_padded_inputs_for_batch(batch_index: int) -> bool ¤
needs_edge_inputs_for_batch(batch_index: int) -> bool ¤
apply(ctx: ForwardContext, *, state: Array | None = None, key: Array | None = None) -> ForwardContext ¤
transfer(old_topology: CompiledTopology, new_topology: CompiledTopology, initialized: FusedDefaultOp) -> FusedDefaultOp ¤
with_dropout(dropout: DropoutLike) -> FusedDefaultOp ¤
apply_output(topology: CompiledTopology, y: Array, *, state: Any = None, key: Array | None = None) -> Array ¤

connex.ops.OutputTransform ¤

Bases: Op

Apply a final transform to the model output array.

This operation runs after all graph nodes have been computed and receives only the ordered output array. Use it for group-wise transforms such as identity, sigmoid, or softmax over outputs.

transform: Callable = _identity class-attribute instance-attribute ¤
name: str = eqx.field(static=True, default='output_transform') class-attribute instance-attribute ¤
layout: ParamLayout = eqx.field(static=True, default=ParamLayout.GLOBAL) class-attribute instance-attribute ¤
stage: str = eqx.field(static=True, default='output') class-attribute instance-attribute ¤
needs_batch_inputs: bool = eqx.field(static=True, default=False) class-attribute instance-attribute ¤
needs_padded_inputs: bool = eqx.field(static=True, default=False) class-attribute instance-attribute ¤
needs_edge_inputs: bool = eqx.field(static=True, default=False) class-attribute instance-attribute ¤
apply_output(topology: CompiledTopology, y: Array, *, state: Any = None, key: Array | None = None) -> Array ¤
transfer(old_topology: CompiledTopology, new_topology: CompiledTopology, initialized: OutputTransform) -> OutputTransform ¤
init(topology: CompiledTopology, *, key: Array) -> Op ¤
init_runtime(topology: CompiledTopology, *, key: Array | None) -> Any ¤
apply(ctx: ForwardContext, *, state: Any = None, key: Array | None = None) -> ForwardContext ¤
apply_inputs(topology: CompiledTopology, values: Array, *, state: Any = None, key: Array | None = None) -> Array ¤
needs_batch_inputs_for_batch(batch_index: int) -> bool ¤
needs_padded_inputs_for_batch(batch_index: int) -> bool ¤
needs_edge_inputs_for_batch(batch_index: int) -> bool ¤

connex.ops.default_ops(*, affine: str = 'hybrid', activation: Callable = jnn.gelu, output_transform: Callable = _identity, dropout: DropoutLike | None = None, fused: bool = False, topo_norm: bool = False, topo_self_attention: bool = False, neuron_self_attention: bool = False, adaptive_activation: bool = False) -> tuple[Op, ...] ¤

Build the standard Connex operation pipeline.

Arguments:

  • affine: One of "padded", "sparse", "matmul", or "hybrid". Hybrid is the default and chooses a backend per topological batch.
  • activation: Elementwise hidden-node activation.
  • output_transform: Final transform applied to the ordered output array.
  • dropout: Optional scalar or mapping dropout configuration. If None, the model uses GraphSpec.dropout.
  • fused: If True, use FusedDefaultOp when no feature operations or adaptive activations are requested.
  • topo_norm: Insert TopoNorm before the affine op.
  • topo_self_attention: Insert TopoSelfAttention before the affine op.
  • neuron_self_attention: Insert NeuronSelfAttention before the affine op.
  • adaptive_activation: Use AdaptiveActivation instead of plain Activation.

Returns:

A tuple of initialized-operation templates suitable for NeuralDAG(spec, ops=...). The model initializes the returned operations against its compiled topology.