Skip to content

Networks

Five built-in architectures, dispatched by NetworkConfig.type:

type Module Use case
mlp mlp.MLP, MultiHeadMLP, ResMLP Most models; default
lstm lstm.LSTMPolicy History-dependent policies, history_len > 1
transformer transformer.TransformerPolicy Same; multi-head attention over history window
linear_plus_mlp linear_plus_mlp.LinearPlusMLP policy = linear(state) + mlp(state); init at the BK linearization
kf_anchored_mlp kf_anchored_mlp CMR-class disaster: K/F outputs pinned to BK anchor

All factories take (n_states, n_policies, hidden_sizes, ..., key) and return an Equinox eqx.Module whose __call__(state) -> policy works on both [n_states] and [batch, n_states] inputs (via jax.vmap of a per-sample helper).

Output bounds are enforced at the network output, per-dimension:

  • Finite policy_upper[i] → sigmoid scaled to [lower, upper].
  • policy_upper[i] = jnp.inf → softplus(x) + lower.

Shared utilities live in networks/common.py: _normalize_input (input shift/scale stop_gradient), _apply_bounds (the sigmoid/softplus dispatch), INIT_FNS (the init-name → init-fn table).

For adding a new network type, see Adding a network.

deqn_jax.networks.mlp

MLP policy network using Equinox.

MLP

Bases: Module

Multi-layer perceptron for policy approximation.

Outputs are optionally bounded using sigmoid scaling

output = lower + (upper - lower) * sigmoid(raw_output)

Attributes:

Name Type Description
layers list

List of linear layers
activations tuple

Per-layer activation functions (one per hidden layer)
output_lower Optional[tuple]

Lower bounds for outputs [n_outputs]
output_upper Optional[tuple]

Upper bounds for outputs [n_outputs]
input_shift Optional[tuple]

Input normalization shift (subtracted) [n_inputs]
input_scale Optional[tuple]

Input normalization scale (divided) [n_inputs]
Source code in src/deqn_jax/networks/mlp.py 
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
class MLP(eqx.Module):
    """Multi-layer perceptron for policy approximation.

    Outputs are optionally bounded using sigmoid scaling:
        output = lower + (upper - lower) * sigmoid(raw_output)

    Attributes:
        layers: List of linear layers
        activations: Per-layer activation functions (one per hidden layer)
        output_lower: Lower bounds for outputs [n_outputs]
        output_upper: Upper bounds for outputs [n_outputs]
        input_shift: Input normalization shift (subtracted) [n_inputs]
        input_scale: Input normalization scale (divided) [n_inputs]
    """

    layers: list
    activations: tuple = eqx.field(static=True)
    # Static fields below: stored as tuple-of-floats so the optimizer
    # can't write to them via Adam-family second-moment updates. See
    # networks/common.py for rationale.
    output_lower: Optional[tuple] = eqx.field(static=True)
    output_upper: Optional[tuple] = eqx.field(
        static=True
    )  # inf replaced with safe finite values
    _has_upper: Optional[tuple] = eqx.field(static=True)  # per-output sigmoid mask
    input_shift: Optional[tuple] = eqx.field(static=True)
    input_scale: Optional[tuple] = eqx.field(static=True)

    def __init__(
        self,
        in_features: int,
        out_features: int,
        hidden_sizes: Sequence[int] = (64, 64),
        activations: Sequence[Callable] = (jax.nn.tanh, jax.nn.tanh),
        output_lower: Optional[Array] = None,
        output_upper: Optional[Array] = None,
        input_shift: Optional[Array] = None,
        input_scale: Optional[Array] = None,
        init: str = "default",
        *,
        key: Array,
    ):
        self.activations = tuple(activations)
        self.output_lower = _to_tuple(output_lower)
        safe_upper, mask = _sanitize_upper(output_upper, output_lower)
        self.output_upper = safe_upper
        self._has_upper = mask
        self.input_shift = _to_tuple(input_shift)
        self.input_scale = _to_tuple(input_scale)

        # Build layers
        sizes = [in_features] + list(hidden_sizes) + [out_features]
        n_layers = len(sizes) - 1
        use_custom_init = init != "default" and init in INIT_FNS

        if use_custom_init:
            # Need extra keys for re-initialization
            all_keys = jax.random.split(key, 2 * n_layers)
            layer_keys = all_keys[:n_layers]
            init_keys = all_keys[n_layers:]
        else:
            layer_keys = jax.random.split(key, n_layers)

        self.layers = []
        for i, (in_size, out_size) in enumerate(zip(sizes[:-1], sizes[1:])):
            layer = eqx.nn.Linear(in_size, out_size, key=layer_keys[i])
            if use_custom_init:
                layer = _apply_init(layer, INIT_FNS[init], init_keys[i])  # pyright: ignore[reportArgumentType]  # ty: ignore[invalid-argument-type]
            self.layers.append(layer)

    def _forward_single(self, x: Array) -> Array:
        """Forward pass for single input [in_features]."""
        x = _normalize_input(x, self.input_shift, self.input_scale)

        # Forward through hidden layers with per-layer activation
        for i, layer in enumerate(self.layers[:-1]):
            x = self.activations[i](layer(x))

        # Output layer (no activation before bounds)
        x = self.layers[-1](x)

        x = _apply_bounds(x, self.output_lower, self.output_upper, self._has_upper)

        return x

    def __call__(self, x: Array) -> Array:
        """Forward pass.

        Args:
            x: Input tensor [batch, in_features] or [in_features]

        Returns:
            Output tensor [batch, out_features] or [out_features]
        """
        if x.ndim == 1:
            return self._forward_single(x)
        else:
            return jax.vmap(self._forward_single)(x)

__call__ 

__call__(x: Array) -> Array

Forward pass.

Parameters:

Name Type Description Default
x Array

Input tensor [batch, in_features] or [in_features]

 required

Returns:

Type Description
Array

Output tensor [batch, out_features] or [out_features]
Source code in src/deqn_jax/networks/mlp.py 
106
107
108
109
110
111
112
113
114
115
116
117
118
def __call__(self, x: Array) -> Array:
    """Forward pass.

    Args:
        x: Input tensor [batch, in_features] or [in_features]

    Returns:
        Output tensor [batch, out_features] or [out_features]
    """
    if x.ndim == 1:
        return self._forward_single(x)
    else:
        return jax.vmap(self._forward_single)(x)

ResMLP

Bases: Module

MLP with residual (skip) connections between hidden layers.

Each hidden layer computes: h_{i+1} = act(W_i @ h_i + b_i) + proj(h_i) where proj is identity if sizes match, or a learned linear projection if hidden sizes differ.

This improves gradient flow and lets the network learn corrections rather than full mappings — helpful for multi-equation PINNs.

Source code in src/deqn_jax/networks/mlp.py 
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
class ResMLP(eqx.Module):
    """MLP with residual (skip) connections between hidden layers.

    Each hidden layer computes: h_{i+1} = act(W_i @ h_i + b_i) + proj(h_i)
    where proj is identity if sizes match, or a learned linear projection
    if hidden sizes differ.

    This improves gradient flow and lets the network learn corrections
    rather than full mappings — helpful for multi-equation PINNs.
    """

    layers: list
    skip_projs: list  # Linear projections for size mismatches (or None)
    activations: tuple = eqx.field(static=True)
    output_lower: Optional[tuple] = eqx.field(static=True)
    output_upper: Optional[tuple] = eqx.field(static=True)
    _has_upper: Optional[tuple] = eqx.field(static=True)
    input_shift: Optional[tuple] = eqx.field(static=True)
    input_scale: Optional[tuple] = eqx.field(static=True)

    def __init__(
        self,
        in_features: int,
        out_features: int,
        hidden_sizes: Sequence[int] = (64, 64),
        activations: Sequence[Callable] = (jax.nn.tanh, jax.nn.tanh),
        output_lower: Optional[Array] = None,
        output_upper: Optional[Array] = None,
        input_shift: Optional[Array] = None,
        input_scale: Optional[Array] = None,
        init: str = "default",
        *,
        key: Array,
    ):
        self.activations = tuple(activations)
        self.output_lower = _to_tuple(output_lower)
        safe_upper, mask = _sanitize_upper(output_upper, output_lower)
        self.output_upper = safe_upper
        self._has_upper = mask
        self.input_shift = _to_tuple(input_shift)
        self.input_scale = _to_tuple(input_scale)

        sizes = [in_features] + list(hidden_sizes) + [out_features]
        n_layers = len(sizes) - 1
        use_custom_init = init != "default" and init in INIT_FNS

        # Keys: layers + skip projections (separate split to not change MLP PRNG)
        layer_keys = jax.random.split(key, n_layers)
        skip_key = jax.random.fold_in(key, 999)
        skip_keys = jax.random.split(skip_key, n_layers)

        if use_custom_init:
            init_key = jax.random.fold_in(key, 998)
            init_keys = jax.random.split(init_key, n_layers)

        self.layers = []
        self.skip_projs = []
        for i, (in_size, out_size) in enumerate(zip(sizes[:-1], sizes[1:])):
            layer = eqx.nn.Linear(in_size, out_size, key=layer_keys[i])
            if use_custom_init:
                layer = _apply_init(layer, INIT_FNS[init], init_keys[i])  # pyright: ignore[reportArgumentType]  # ty: ignore[invalid-argument-type]
            self.layers.append(layer)

            # Skip connection for hidden layers (not the output layer)
            if i < n_layers - 1:
                if in_size == out_size:
                    self.skip_projs.append(None)  # identity
                else:
                    self.skip_projs.append(
                        eqx.nn.Linear(
                            in_size, out_size, use_bias=False, key=skip_keys[i]
                        )
                    )
            else:
                self.skip_projs.append(None)  # no skip for output layer

    def _forward_single(self, x: Array) -> Array:
        if self.input_shift is not None:
            x = (x - jax.lax.stop_gradient(self.input_shift)) / jax.lax.stop_gradient(
                self.input_scale
            )

        for i, layer in enumerate(self.layers[:-1]):
            residual = x
            x = self.activations[i](layer(x))
            # Add skip connection
            proj = self.skip_projs[i]
            if proj is not None:
                x = x + proj(residual)
            else:
                x = x + residual

        # Output layer (no skip, no activation before bounds)
        x = self.layers[-1](x)

        x = _apply_bounds(x, self.output_lower, self.output_upper, self._has_upper)

        return x

    def __call__(self, x: Array) -> Array:
        if x.ndim == 1:
            return self._forward_single(x)
        else:
            return jax.vmap(self._forward_single)(x)

MultiHeadMLP

Bases: Module

MLP with separate output heads per policy variable.

Shared trunk → per-policy linear heads. Gives each policy its own output parameters, reducing gradient interference between equations that depend on different policies.

Source code in src/deqn_jax/networks/mlp.py 
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
class MultiHeadMLP(eqx.Module):
    """MLP with separate output heads per policy variable.

    Shared trunk → per-policy linear heads. Gives each policy its own
    output parameters, reducing gradient interference between equations
    that depend on different policies.
    """

    trunk_layers: list
    heads: list  # list of eqx.nn.Linear, one per output
    activations: tuple = eqx.field(static=True)
    output_lower: Optional[tuple] = eqx.field(static=True)
    output_upper: Optional[tuple] = eqx.field(static=True)
    _has_upper: Optional[tuple] = eqx.field(static=True)
    input_shift: Optional[tuple] = eqx.field(static=True)
    input_scale: Optional[tuple] = eqx.field(static=True)

    def __init__(
        self,
        in_features: int,
        out_features: int,
        hidden_sizes: Sequence[int] = (64, 64),
        activations: Sequence[Callable] = (jax.nn.tanh, jax.nn.tanh),
        output_lower: Optional[Array] = None,
        output_upper: Optional[Array] = None,
        input_shift: Optional[Array] = None,
        input_scale: Optional[Array] = None,
        init: str = "default",
        *,
        key: Array,
    ):
        self.activations = tuple(activations)
        self.output_lower = _to_tuple(output_lower)
        safe_upper, mask = _sanitize_upper(output_upper, output_lower)
        self.output_upper = safe_upper
        self._has_upper = mask
        self.input_shift = _to_tuple(input_shift)
        self.input_scale = _to_tuple(input_scale)

        # Build trunk (hidden layers only, no output layer)
        sizes = [in_features] + list(hidden_sizes)
        n_trunk = len(sizes) - 1
        use_custom_init = init != "default" and init in INIT_FNS

        key, *trunk_keys = jax.random.split(key, n_trunk + 1)
        if use_custom_init:
            key, *init_keys = jax.random.split(key, n_trunk + 1)

        self.trunk_layers = []
        for i, (in_size, out_size) in enumerate(zip(sizes[:-1], sizes[1:])):
            layer = eqx.nn.Linear(in_size, out_size, key=trunk_keys[i])
            if use_custom_init:
                layer = _apply_init(layer, INIT_FNS[init], init_keys[i])  # pyright: ignore[reportArgumentType]  # ty: ignore[invalid-argument-type]
            self.trunk_layers.append(layer)

        # Build per-policy output heads: each hidden_sizes[-1] → 1
        head_keys = jax.random.split(key, out_features)
        self.heads = [
            eqx.nn.Linear(hidden_sizes[-1], 1, key=head_keys[i])
            for i in range(out_features)
        ]

    def _forward_single(self, x: Array) -> Array:
        if self.input_shift is not None:
            x = (x - jax.lax.stop_gradient(self.input_shift)) / jax.lax.stop_gradient(
                self.input_scale
            )

        for i, layer in enumerate(self.trunk_layers):
            x = self.activations[i](layer(x))

        # Each head produces one scalar, concat into [out_features]
        raw = jnp.concatenate([head(x) for head in self.heads], axis=-1)

        raw = _apply_bounds(raw, self.output_lower, self.output_upper, self._has_upper)

        return raw

    def __call__(self, x: Array) -> Array:
        if x.ndim == 1:
            return self._forward_single(x)
        else:
            return jax.vmap(self._forward_single)(x)

create_mlp 

create_mlp(
    n_states: int,
    n_policies: int,
    hidden_sizes: Sequence[int] = (64, 64),
    activation: str = "tanh",
    activations: Optional[Sequence[str]] = None,
    init: str = "default",
    policy_lower: Optional[Array] = None,
    policy_upper: Optional[Array] = None,
    multi_head: bool = False,
    skip_connections: bool = False,
    input_shift: Optional[Array] = None,
    input_scale: Optional[Array] = None,
    *,
    key: Array,
) -> eqx.Module

Factory function to create MLP with common configurations.

Parameters:

Name Type Description Default
n_states int

Number of state variables (input dimension)

 required
n_policies int

Number of policy variables (output dimension)

 required
hidden_sizes Sequence[int]

Tuple of hidden layer sizes

 (64, 64)
activation str

Activation name for all hidden layers (default: "tanh")

 'tanh'
activations Optional[Sequence[str]]

Per-layer activation names (overrides activation if set)

 None
init str

Weight initialization ("xavier_normal", "xavier_uniform", "he_normal", "he_uniform", "lecun_normal", "default")

 'default'
policy_lower Optional[Array]

Lower bounds for policy outputs

 None
policy_upper Optional[Array]

Upper bounds for policy outputs

 None
skip_connections bool

Use residual connections between hidden layers

 False
key Array

JAX PRNG key

 required

Returns:

Type Description
Module

Initialized MLP model
Source code in src/deqn_jax/networks/mlp.py 
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
def create_mlp(
    n_states: int,
    n_policies: int,
    hidden_sizes: Sequence[int] = (64, 64),
    activation: str = "tanh",
    activations: Optional[Sequence[str]] = None,
    init: str = "default",
    policy_lower: Optional[Array] = None,
    policy_upper: Optional[Array] = None,
    multi_head: bool = False,
    skip_connections: bool = False,
    input_shift: Optional[Array] = None,
    input_scale: Optional[Array] = None,
    *,
    key: Array,
) -> eqx.Module:
    """Factory function to create MLP with common configurations.

    Args:
        n_states: Number of state variables (input dimension)
        n_policies: Number of policy variables (output dimension)
        hidden_sizes: Tuple of hidden layer sizes
        activation: Activation name for all hidden layers (default: "tanh")
        activations: Per-layer activation names (overrides activation if set)
        init: Weight initialization ("xavier_normal", "xavier_uniform",
              "he_normal", "he_uniform", "lecun_normal", "default")
        policy_lower: Lower bounds for policy outputs
        policy_upper: Upper bounds for policy outputs
        skip_connections: Use residual connections between hidden layers
        key: JAX PRNG key

    Returns:
        Initialized MLP model
    """
    n_hidden = len(hidden_sizes)

    # Resolve per-layer activations
    if activations is not None:
        if len(activations) != n_hidden:
            raise ValueError(
                f"activations length ({len(activations)}) must match "
                f"hidden_sizes length ({n_hidden})"
            )
        act_fns = tuple(_resolve_activation(a) for a in activations)
    else:
        act_fn = _resolve_activation(activation)
        act_fns = tuple(act_fn for _ in range(n_hidden))

    if multi_head:
        cls = MultiHeadMLP
    elif skip_connections:
        cls = ResMLP
    else:
        cls = MLP
    return cls(
        in_features=n_states,
        out_features=n_policies,
        hidden_sizes=hidden_sizes,
        activations=act_fns,
        output_lower=policy_lower,
        output_upper=policy_upper,
        input_shift=input_shift,
        input_scale=input_scale,
        init=init,
        key=key,
    )

deqn_jax.networks.lstm

Multi-layer LSTM policy network using Equinox.

Takes a history window of states [B, H, D] and outputs policy [B, P] for the current (last) timestep. Supports: - Multi-layer stacked LSTM with configurable depth - Input normalization (frozen via stop_gradient) - Mixed softplus/sigmoid output bounding (via common._apply_bounds)

LSTMPolicy

Bases: Module

Multi-layer LSTM for history-dependent policy approximation.

Input: [batch, history_len, n_states] or [history_len, n_states] Output: [batch, n_policies] or [n_policies]

Architecture

input [D] -> Linear -> hidden_size -> LSTMCell layer 1 -> ... -> LSTMCell layer L -> final hidden state -> Linear -> n_policies -> bounds

Source code in src/deqn_jax/networks/lstm.py 
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
class LSTMPolicy(eqx.Module):
    """Multi-layer LSTM for history-dependent policy approximation.

    Input: [batch, history_len, n_states] or [history_len, n_states]
    Output: [batch, n_policies] or [n_policies]

    Architecture:
        input [D] -> Linear -> hidden_size
        -> LSTMCell layer 1 -> ... -> LSTMCell layer L
        -> final hidden state -> Linear -> n_policies -> bounds
    """

    input_proj: eqx.nn.Linear
    cells: list  # list of eqx.nn.LSTMCell
    output_proj: eqx.nn.Linear
    hidden_size: int = eqx.field(static=True)
    n_layers: int = eqx.field(static=True)
    history_len: int = eqx.field(static=True)
    output_lower: Optional[tuple] = eqx.field(static=True)
    output_upper: Optional[tuple] = eqx.field(static=True)
    _has_upper: Optional[tuple] = eqx.field(static=True)
    input_shift: Optional[tuple] = eqx.field(static=True)
    input_scale: Optional[tuple] = eqx.field(static=True)

    def __init__(
        self,
        in_features: int,
        out_features: int,
        hidden_sizes: Sequence[int] = (64,),
        history_len: int = 10,
        output_lower: Optional[Array] = None,
        output_upper: Optional[Array] = None,
        input_shift: Optional[Array] = None,
        input_scale: Optional[Array] = None,
        *,
        key: Array,
    ):
        self.history_len = history_len
        self.output_lower = _to_tuple(output_lower)
        safe_upper, mask = _sanitize_upper(output_upper, output_lower)
        self.output_upper = safe_upper
        self._has_upper = mask
        self.input_shift = _to_tuple(input_shift)
        self.input_scale = _to_tuple(input_scale)

        if isinstance(hidden_sizes, int):
            hidden_sizes = [hidden_sizes]
        else:
            hidden_sizes = list(hidden_sizes)
        self.hidden_size = hidden_sizes[0]
        self.n_layers = len(hidden_sizes)

        # Split keys: input_proj, each LSTM cell, output_proj
        keys = jax.random.split(key, self.n_layers + 2)

        # Input projection: n_states -> hidden_size
        self.input_proj = eqx.nn.Linear(in_features, hidden_sizes[0], key=keys[0])  # pyright: ignore[reportAssignmentType]  # ty: ignore[invalid-assignment]

        # Stacked LSTM cells
        self.cells = []
        for i in range(self.n_layers):
            cell_in = hidden_sizes[i]  # first cell takes from input_proj
            cell_hidden = hidden_sizes[i]
            if i > 0:
                cell_in = hidden_sizes[i - 1]
            # LSTMCell(input_size, hidden_size)
            # For layer 0: input is hidden_sizes[0] (from input_proj)
            # For layer i>0: input is hidden_sizes[i-1] (from previous layer's h)
            self.cells.append(eqx.nn.LSTMCell(cell_in, cell_hidden, key=keys[i + 1]))

        # Output projection: last hidden_size -> n_policies
        self.output_proj = eqx.nn.Linear(hidden_sizes[-1], out_features, key=keys[-1])  # pyright: ignore[reportAssignmentType]  # ty: ignore[invalid-assignment]

    def _forward_single(self, x: Array) -> Array:
        """Forward pass for single sequence [history_len, n_states]."""

        # Normalize each timestep's input
        def norm_step(x_t):
            return _normalize_input(x_t, self.input_shift, self.input_scale)

        x = jax.vmap(norm_step)(x)  # [H, D]

        # Project input: [H, D] -> [H, hidden_size]
        x = jax.vmap(lambda x_t: jax.nn.tanh(self.input_proj(x_t)))(x)

        # Scan through each layer sequentially
        layer_input = x  # [H, hidden_sizes[0]]
        for layer_idx in range(self.n_layers):
            cell = self.cells[layer_idx]
            h_size = cell.hidden_size

            init_state = (jnp.zeros(h_size), jnp.zeros(h_size))

            def scan_fn(carry, x_t, cell=cell):
                (h, c) = cell(x_t, carry)
                return (h, c), h

            (final_h, _), all_h = jax.lax.scan(scan_fn, init_state, layer_input)

            # Next layer takes this layer's hidden states as input
            layer_input = all_h  # [H, hidden_sizes[layer_idx]]

        # Use final hidden state from last layer
        raw = self.output_proj(final_h)
        return _apply_bounds(raw, self.output_lower, self.output_upper, self._has_upper)

    def __call__(self, x: Array) -> Array:
        """Forward pass.

        Args:
            x: [batch, history_len, n_states] or [history_len, n_states]

        Returns:
            [batch, n_policies] or [n_policies]
        """
        if x.ndim == 2:
            return self._forward_single(x)
        else:
            return jax.vmap(self._forward_single)(x)

__call__ 

__call__(x: Array) -> Array

Forward pass.

Parameters:

Name Type Description Default
x Array

[batch, history_len, n_states] or [history_len, n_states]

 required

Returns:

Type Description
Array

[batch, n_policies] or [n_policies]
Source code in src/deqn_jax/networks/lstm.py 
131
132
133
134
135
136
137
138
139
140
141
142
143
def __call__(self, x: Array) -> Array:
    """Forward pass.

    Args:
        x: [batch, history_len, n_states] or [history_len, n_states]

    Returns:
        [batch, n_policies] or [n_policies]
    """
    if x.ndim == 2:
        return self._forward_single(x)
    else:
        return jax.vmap(self._forward_single)(x)

create_lstm 

create_lstm(
    n_states: int,
    n_policies: int,
    hidden_sizes: Sequence[int] = (64,),
    history_len: int = 10,
    policy_lower: Optional[Array] = None,
    policy_upper: Optional[Array] = None,
    input_shift: Optional[Array] = None,
    input_scale: Optional[Array] = None,
    *,
    key: Array,
) -> LSTMPolicy

Factory function to create multi-layer LSTM policy network.

Parameters:

Name Type Description Default
n_states int

Number of state variables

 required
n_policies int

Number of policy variables

 required
hidden_sizes Sequence[int]

Tuple of hidden sizes (one per LSTM layer)

 (64,)
history_len int

Length of history window

 10
policy_lower Optional[Array]

Lower bounds for outputs

 None
policy_upper Optional[Array]

Upper bounds for outputs

 None
input_shift Optional[Array]

Input normalization shift

 None
input_scale Optional[Array]

Input normalization scale

 None
key Array

PRNG key

 required

Returns:

Type Description
LSTMPolicy

Initialized LSTMPolicy
Source code in src/deqn_jax/networks/lstm.py 
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
def create_lstm(
    n_states: int,
    n_policies: int,
    hidden_sizes: Sequence[int] = (64,),
    history_len: int = 10,
    policy_lower: Optional[Array] = None,
    policy_upper: Optional[Array] = None,
    input_shift: Optional[Array] = None,
    input_scale: Optional[Array] = None,
    *,
    key: Array,
) -> LSTMPolicy:
    """Factory function to create multi-layer LSTM policy network.

    Args:
        n_states: Number of state variables
        n_policies: Number of policy variables
        hidden_sizes: Tuple of hidden sizes (one per LSTM layer)
        history_len: Length of history window
        policy_lower: Lower bounds for outputs
        policy_upper: Upper bounds for outputs
        input_shift: Input normalization shift
        input_scale: Input normalization scale
        key: PRNG key

    Returns:
        Initialized LSTMPolicy
    """
    return LSTMPolicy(
        in_features=n_states,
        out_features=n_policies,
        hidden_sizes=hidden_sizes,
        history_len=history_len,
        output_lower=policy_lower,
        output_upper=policy_upper,
        input_shift=input_shift,
        input_scale=input_scale,
        key=key,
    )

deqn_jax.networks.transformer

Transformer policy network using Equinox.

Takes a history window of states [B, H, D] and outputs policy [B, P] for the current (last) timestep. Supports: - Learned positional embeddings - Pre-LN transformer blocks (LayerNorm before attention/FFN) - Mixed softplus/sigmoid output bounding (via common._apply_bounds)

MultiHeadSelfAttention

Bases: Module

Simple multi-head self-attention without dropout.

Avoids eqx.nn.MultiheadAttention whose Dropout module has a non-static inference boolean that breaks jax.jit tracing.

Source code in src/deqn_jax/networks/transformer.py 
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
class MultiHeadSelfAttention(eqx.Module):
    """Simple multi-head self-attention without dropout.

    Avoids eqx.nn.MultiheadAttention whose Dropout module has a
    non-static `inference` boolean that breaks jax.jit tracing.
    """

    q_proj: eqx.nn.Linear
    k_proj: eqx.nn.Linear
    v_proj: eqx.nn.Linear
    o_proj: eqx.nn.Linear
    num_heads: int = eqx.field(static=True)
    head_dim: int = eqx.field(static=True)

    def __init__(self, hidden_dim: int, num_heads: int, *, key: Array):
        k1, k2, k3, k4 = jax.random.split(key, 4)
        self.num_heads = num_heads
        self.head_dim = hidden_dim // num_heads
        # eqx.nn.Linear's typestub returns the Module base; the runtime
        # value IS Linear and the field annotations are correct. Suppress
        # the four init-site assignments.
        self.q_proj = eqx.nn.Linear(hidden_dim, hidden_dim, key=k1)  # pyright: ignore[reportAssignmentType]  # ty: ignore[invalid-assignment]
        self.k_proj = eqx.nn.Linear(hidden_dim, hidden_dim, key=k2)  # pyright: ignore[reportAssignmentType]  # ty: ignore[invalid-assignment]
        self.v_proj = eqx.nn.Linear(hidden_dim, hidden_dim, key=k3)  # pyright: ignore[reportAssignmentType]  # ty: ignore[invalid-assignment]
        self.o_proj = eqx.nn.Linear(hidden_dim, hidden_dim, key=k4)  # pyright: ignore[reportAssignmentType]  # ty: ignore[invalid-assignment]

    def __call__(self, x: Array) -> Array:
        """Self-attention on [seq_len, hidden_dim]."""
        seq_len = x.shape[0]
        # Project Q, K, V: [S, D] -> [S, D]
        q = jax.vmap(self.q_proj)(x)
        k = jax.vmap(self.k_proj)(x)
        v = jax.vmap(self.v_proj)(x)

        # Reshape to [S, H, head_dim] then transpose to [H, S, head_dim]
        q = q.reshape(seq_len, self.num_heads, self.head_dim).transpose(1, 0, 2)
        k = k.reshape(seq_len, self.num_heads, self.head_dim).transpose(1, 0, 2)
        v = v.reshape(seq_len, self.num_heads, self.head_dim).transpose(1, 0, 2)

        # Scaled dot-product attention: [H, S, S]
        scale = jnp.sqrt(jnp.array(self.head_dim, dtype=q.dtype))
        attn_weights = jnp.matmul(q, k.transpose(0, 2, 1)) / scale
        attn_weights = jax.nn.softmax(attn_weights, axis=-1)

        # Weighted values: [H, S, head_dim]
        attn_out = jnp.matmul(attn_weights, v)

        # Reshape back: [H, S, head_dim] -> [S, D]
        attn_out = attn_out.transpose(1, 0, 2).reshape(seq_len, -1)

        # Output projection
        return jax.vmap(self.o_proj)(attn_out)

__call__ 

__call__(x: Array) -> Array

Self-attention on [seq_len, hidden_dim].

Source code in src/deqn_jax/networks/transformer.py 
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
def __call__(self, x: Array) -> Array:
    """Self-attention on [seq_len, hidden_dim]."""
    seq_len = x.shape[0]
    # Project Q, K, V: [S, D] -> [S, D]
    q = jax.vmap(self.q_proj)(x)
    k = jax.vmap(self.k_proj)(x)
    v = jax.vmap(self.v_proj)(x)

    # Reshape to [S, H, head_dim] then transpose to [H, S, head_dim]
    q = q.reshape(seq_len, self.num_heads, self.head_dim).transpose(1, 0, 2)
    k = k.reshape(seq_len, self.num_heads, self.head_dim).transpose(1, 0, 2)
    v = v.reshape(seq_len, self.num_heads, self.head_dim).transpose(1, 0, 2)

    # Scaled dot-product attention: [H, S, S]
    scale = jnp.sqrt(jnp.array(self.head_dim, dtype=q.dtype))
    attn_weights = jnp.matmul(q, k.transpose(0, 2, 1)) / scale
    attn_weights = jax.nn.softmax(attn_weights, axis=-1)

    # Weighted values: [H, S, head_dim]
    attn_out = jnp.matmul(attn_weights, v)

    # Reshape back: [H, S, head_dim] -> [S, D]
    attn_out = attn_out.transpose(1, 0, 2).reshape(seq_len, -1)

    # Output projection
    return jax.vmap(self.o_proj)(attn_out)

TransformerBlock

Bases: Module

Pre-LN Transformer block: LayerNorm -> Attention -> Residual -> LayerNorm -> FFN -> Residual.

Source code in src/deqn_jax/networks/transformer.py 
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
class TransformerBlock(eqx.Module):
    """Pre-LN Transformer block: LayerNorm -> Attention -> Residual -> LayerNorm -> FFN -> Residual."""

    attn: MultiHeadSelfAttention
    ln1: eqx.nn.LayerNorm
    ln2: eqx.nn.LayerNorm
    ffn_up: eqx.nn.Linear
    ffn_down: eqx.nn.Linear

    def __init__(self, hidden_dim: int, num_heads: int, *, key: Array):
        k1, k2, k3 = jax.random.split(key, 3)
        # See MultiHeadSelfAttention.__init__ for the eqx-typing rationale.
        self.attn = MultiHeadSelfAttention(hidden_dim, num_heads, key=k1)  # pyright: ignore[reportAssignmentType]  # ty: ignore[invalid-assignment]
        self.ln1 = eqx.nn.LayerNorm(hidden_dim)  # pyright: ignore[reportAssignmentType]  # ty: ignore[invalid-assignment]
        self.ln2 = eqx.nn.LayerNorm(hidden_dim)  # pyright: ignore[reportAssignmentType]  # ty: ignore[invalid-assignment]
        ffn_dim = hidden_dim * 4
        self.ffn_up = eqx.nn.Linear(hidden_dim, ffn_dim, key=k2)  # pyright: ignore[reportAssignmentType]  # ty: ignore[invalid-assignment]
        self.ffn_down = eqx.nn.Linear(ffn_dim, hidden_dim, key=k3)  # pyright: ignore[reportAssignmentType]  # ty: ignore[invalid-assignment]

    def __call__(self, x: Array) -> Array:
        """Forward pass for single sequence [seq_len, hidden_dim]."""
        # Pre-LN self-attention + residual
        normed = jax.vmap(self.ln1)(x)
        attn_out = self.attn(normed)
        x = x + attn_out

        # Pre-LN FFN + residual
        normed = jax.vmap(self.ln2)(x)
        ffn_out = jax.vmap(lambda h: self.ffn_down(jax.nn.gelu(self.ffn_up(h))))(normed)
        x = x + ffn_out

        return x

__call__ 

__call__(x: Array) -> Array

Forward pass for single sequence [seq_len, hidden_dim].

Source code in src/deqn_jax/networks/transformer.py 
 98
 99
100
101
102
103
104
105
106
107
108
109
110
def __call__(self, x: Array) -> Array:
    """Forward pass for single sequence [seq_len, hidden_dim]."""
    # Pre-LN self-attention + residual
    normed = jax.vmap(self.ln1)(x)
    attn_out = self.attn(normed)
    x = x + attn_out

    # Pre-LN FFN + residual
    normed = jax.vmap(self.ln2)(x)
    ffn_out = jax.vmap(lambda h: self.ffn_down(jax.nn.gelu(self.ffn_up(h))))(normed)
    x = x + ffn_out

    return x

TransformerPolicy

Bases: Module

Transformer for history-dependent policy approximation.

Input: [batch, history_len, n_states] or [history_len, n_states] Output: [batch, n_policies] or [n_policies]

Architecture

input [H, D] -> normalize -> Linear -> [H, hidden_dim] + learned positional embeddings -> N TransformerBlocks -> extract last timestep [hidden_dim] -> Linear -> [n_policies] -> bounds

Source code in src/deqn_jax/networks/transformer.py 
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
class TransformerPolicy(eqx.Module):
    """Transformer for history-dependent policy approximation.

    Input: [batch, history_len, n_states] or [history_len, n_states]
    Output: [batch, n_policies] or [n_policies]

    Architecture:
        input [H, D] -> normalize -> Linear -> [H, hidden_dim]
        + learned positional embeddings
        -> N TransformerBlocks
        -> extract last timestep [hidden_dim]
        -> Linear -> [n_policies] -> bounds
    """

    input_proj: eqx.nn.Linear
    pos_embed: Array  # [max_history, hidden_dim]
    blocks: list  # list of TransformerBlock
    final_ln: eqx.nn.LayerNorm
    output_proj: eqx.nn.Linear
    hidden_dim: int = eqx.field(static=True)
    n_layers: int = eqx.field(static=True)
    history_len: int = eqx.field(static=True)
    output_lower: Optional[tuple] = eqx.field(static=True)
    output_upper: Optional[tuple] = eqx.field(static=True)
    _has_upper: Optional[tuple] = eqx.field(static=True)
    input_shift: Optional[tuple] = eqx.field(static=True)
    input_scale: Optional[tuple] = eqx.field(static=True)

    def __init__(
        self,
        in_features: int,
        out_features: int,
        hidden_dim: int = 64,
        n_layers: int = 2,
        num_heads: int = 4,
        history_len: int = 10,
        output_lower: Optional[Array] = None,
        output_upper: Optional[Array] = None,
        input_shift: Optional[Array] = None,
        input_scale: Optional[Array] = None,
        *,
        key: Array,
    ):
        self.hidden_dim = hidden_dim
        self.n_layers = n_layers
        self.history_len = history_len
        self.output_lower = _to_tuple(output_lower)
        safe_upper, mask = _sanitize_upper(output_upper, output_lower)
        self.output_upper = safe_upper
        self._has_upper = mask
        self.input_shift = _to_tuple(input_shift)
        self.input_scale = _to_tuple(input_scale)

        # Split keys: input_proj, pos_embed, each block, final_ln, output_proj
        keys = jax.random.split(key, n_layers + 3)

        self.input_proj = eqx.nn.Linear(in_features, hidden_dim, key=keys[0])  # pyright: ignore[reportAssignmentType]  # ty: ignore[invalid-assignment]
        self.pos_embed = jax.random.normal(keys[1], (history_len, hidden_dim)) * 0.02

        self.blocks = [
            TransformerBlock(hidden_dim, num_heads, key=keys[i + 2])
            for i in range(n_layers)
        ]

        self.final_ln = eqx.nn.LayerNorm(hidden_dim)  # pyright: ignore[reportAssignmentType]  # ty: ignore[invalid-assignment]
        self.output_proj = eqx.nn.Linear(hidden_dim, out_features, key=keys[-1])  # pyright: ignore[reportAssignmentType]  # ty: ignore[invalid-assignment]

    def _forward_single(self, x: Array) -> Array:
        """Forward pass for single sequence [history_len, n_states]."""
        # Normalize each timestep
        x = jax.vmap(
            lambda x_t: _normalize_input(x_t, self.input_shift, self.input_scale)
        )(x)

        # Project to hidden dim: [H, D] -> [H, hidden_dim]
        x = jax.vmap(self.input_proj)(x)

        # Add positional embeddings
        x = x + self.pos_embed[: x.shape[0]]

        # Transformer blocks
        for block in self.blocks:
            x = block(x)

        # Final layer norm on last timestep
        last = self.final_ln(x[-1])

        # Output projection + bounds
        raw = self.output_proj(last)
        return _apply_bounds(raw, self.output_lower, self.output_upper, self._has_upper)

    def __call__(self, x: Array) -> Array:
        """Forward pass.

        Args:
            x: [batch, history_len, n_states] or [history_len, n_states]

        Returns:
            [batch, n_policies] or [n_policies]
        """
        if x.ndim == 2:
            return self._forward_single(x)
        else:
            return jax.vmap(self._forward_single)(x)

__call__ 

__call__(x: Array) -> Array

Forward pass.

Parameters:

Name Type Description Default
x Array

[batch, history_len, n_states] or [history_len, n_states]

 required

Returns:

Type Description
Array

[batch, n_policies] or [n_policies]
Source code in src/deqn_jax/networks/transformer.py 
204
205
206
207
208
209
210
211
212
213
214
215
216
def __call__(self, x: Array) -> Array:
    """Forward pass.

    Args:
        x: [batch, history_len, n_states] or [history_len, n_states]

    Returns:
        [batch, n_policies] or [n_policies]
    """
    if x.ndim == 2:
        return self._forward_single(x)
    else:
        return jax.vmap(self._forward_single)(x)

create_transformer 

create_transformer(
    n_states: int,
    n_policies: int,
    hidden_dim: int = 64,
    n_layers: int = 2,
    num_heads: int = 4,
    history_len: int = 10,
    policy_lower: Optional[Array] = None,
    policy_upper: Optional[Array] = None,
    input_shift: Optional[Array] = None,
    input_scale: Optional[Array] = None,
    *,
    key: Array,
) -> TransformerPolicy

Factory function to create Transformer policy network.

Parameters:

Name Type Description Default
n_states int

Number of state variables

 required
n_policies int

Number of policy variables

 required
hidden_dim int

Transformer hidden dimension (must be divisible by num_heads)

 64
n_layers int

Number of transformer blocks

 2
num_heads int

Number of attention heads

 4
history_len int

Length of history window

 10
policy_lower Optional[Array]

Lower bounds for outputs

 None
policy_upper Optional[Array]

Upper bounds for outputs

 None
input_shift Optional[Array]

Input normalization shift

 None
input_scale Optional[Array]

Input normalization scale

 None
key Array

PRNG key

 required

Returns:

Type Description
TransformerPolicy

Initialized TransformerPolicy
Source code in src/deqn_jax/networks/transformer.py 
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
def create_transformer(
    n_states: int,
    n_policies: int,
    hidden_dim: int = 64,
    n_layers: int = 2,
    num_heads: int = 4,
    history_len: int = 10,
    policy_lower: Optional[Array] = None,
    policy_upper: Optional[Array] = None,
    input_shift: Optional[Array] = None,
    input_scale: Optional[Array] = None,
    *,
    key: Array,
) -> TransformerPolicy:
    """Factory function to create Transformer policy network.

    Args:
        n_states: Number of state variables
        n_policies: Number of policy variables
        hidden_dim: Transformer hidden dimension (must be divisible by num_heads)
        n_layers: Number of transformer blocks
        num_heads: Number of attention heads
        history_len: Length of history window
        policy_lower: Lower bounds for outputs
        policy_upper: Upper bounds for outputs
        input_shift: Input normalization shift
        input_scale: Input normalization scale
        key: PRNG key

    Returns:
        Initialized TransformerPolicy
    """
    return TransformerPolicy(
        in_features=n_states,
        out_features=n_policies,
        hidden_dim=hidden_dim,
        n_layers=n_layers,
        num_heads=num_heads,
        history_len=history_len,
        output_lower=policy_lower,
        output_upper=policy_upper,
        input_shift=input_shift,
        input_scale=input_scale,
        key=key,
    )

deqn_jax.networks.linear_plus_mlp

Residual policy: analytically-added linear policy + MLP correction.

Two output parameterizations are supported, selectable per-policy via output_links:

  • "linear" (default, original behavior): policy_i(s) = ss_i + P_i^level @ (s - ss_state) + delta_mlp_i(s)

  • "log" (positive-multiplicative around SS): policy_i(s) = ss_i * exp(P_i^log @ (s - ss_state) + delta_mlp_i(s))

    where P_i^log = P_i^level / ss_i (delta-method conversion of the BK matrix). At init (init_scale=0) and at s=ss_state, both forms reduce to ss_i exactly. The log form bakes in positivity and is the natural parameterization for log-deviations-from-SS, which is the standard DSGE convention (cf. Dynare's log-linearized solutions).

The linear part is the Dynare-linearized (Blanchard-Kahn) solution, which solves the model to first order and is correct by construction near SS. The MLP starts at zero (via a scaled final layer) so at initialization the full policy IS the linear policy (in level or log space depending on the per-policy link). Training only learns corrections on top.

This architecture solves a specific PINN pathology observed when networks are trained directly against equilibrium residuals: a bare MLP can converge to a degenerate fixed point where policies collapse to low values, even when that produces dynamics far from the true SS. A residual parameterization inherits the linear policy's correctness as a floor — the network can only help, never hurt.

The MLP is UNBOUNDED (no sigmoid/softplus on output). Instead, the linear baseline + small MLP corrections keep policies in the valid region. Hard clipping to policy bounds happens at the very end to prevent catastrophic policy outputs during early training.

This module is model-agnostic. Per-model shape priors (e.g. K/F gauge masking, ELB feature augmentation, output-space reparameterizations to encode equation-specific curvature) live in the model's own network.py module, not here. See models/disaster/network.py for the disaster-specific shape priors layered on top of this core.

LinearPlusMLP

Bases: Module

Generic residual ansatz: policy = clip(π_BK + δ_θ, lower, upper).

Attributes:

Name Type Description
mlp MLP

Unbounded MLP that outputs the correction (delta) in output space.
P Array

Policy rule matrix from linearization [n_policies, n_states].
ss_state Array

Steady-state state vector [n_states].
ss_policy Array

Steady-state policy vector [n_policies].
policy_lower / policy_upper

Hard clipping bounds (safety fence; the linear + delta should normally keep policies well inside these).

The forward pass is intentionally minimal: linearization, MLP delta, sum, hard clip. Per-model shape priors should be implemented in a model-specific subclass or wrapper that lives next to the model definition, not bolted onto this generic core.

Source code in src/deqn_jax/networks/linear_plus_mlp.py 
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
class LinearPlusMLP(eqx.Module):
    """Generic residual ansatz: ``policy = clip(π_BK + δ_θ, lower, upper)``.

    Attributes:
        mlp: Unbounded MLP that outputs the correction (delta) in output space.
        P: Policy rule matrix from linearization [n_policies, n_states].
        ss_state: Steady-state state vector [n_states].
        ss_policy: Steady-state policy vector [n_policies].
        policy_lower / policy_upper: Hard clipping bounds (safety fence; the
            linear + delta should normally keep policies well inside these).

    The forward pass is intentionally minimal: linearization, MLP delta,
    sum, hard clip. Per-model shape priors should be implemented in a
    model-specific subclass or wrapper that lives next to the model
    definition, not bolted onto this generic core.
    """

    mlp: MLP
    # P, ss_state, ss_policy are fixed throughout training but participate
    # in the forward pass as JAX arrays (not static): marking them static
    # would cache them as pytree metadata and require stop_gradient.
    # P is stored *per-policy in the appropriate space*: row i is in level
    # units for linear-output policies, in log units for log-output policies
    # (i.e. log-row i = level-row i / ss_i). Conversion happens in the
    # factory so _forward_single doesn't branch on link type for the BK row.
    P: Array
    ss_state: Array
    ss_policy: Array
    # output_links: tuple of ints [n_policies], 0=linear, 1=log. Stored as
    # Python tuple for true staticness (mirrors policy_lower/upper pattern);
    # converted to a JAX array in _forward_single.
    output_links: tuple = eqx.field(static=True)
    policy_lower: Optional[tuple] = eqx.field(static=True)
    policy_upper: Optional[tuple] = eqx.field(static=True)

    def __init__(
        self,
        n_states: int,
        n_policies: int,
        hidden_sizes: Sequence[int],
        activation: str,
        P: Array,
        ss_state: Array,
        ss_policy: Array,
        output_links: Optional[Sequence[str]] = None,
        policy_lower: Optional[Array] = None,
        policy_upper: Optional[Array] = None,
        init: str = "default",
        init_scale: float = 0.01,
        input_shift: Optional[Array] = None,
        input_scale: Optional[Array] = None,
        *,
        key: Array,
    ):
        # UNBOUNDED MLP for the correction. The linear component provides
        # the "raw" output, so we don't want sigmoid/softplus on top.
        # eqx-typing: MLP() is typed as returning Module (the base class);
        # the runtime is MLP and the field annotation is correct.
        act_fn = _resolve_activation(activation)
        self.mlp = MLP(  # pyright: ignore[reportAssignmentType]  # ty: ignore[invalid-assignment]
            in_features=n_states,
            out_features=n_policies,
            hidden_sizes=hidden_sizes,
            activations=[act_fn] * len(hidden_sizes),
            output_lower=None,
            output_upper=None,
            input_shift=input_shift,
            input_scale=input_scale,
            init=init,
            key=key,
        )

        # Shrink final layer so delta ≈ 0 at initialization: the policy IS
        # the linear policy at init, so training only moves away from it
        # to the extent that doing so reduces residuals.
        last = self.mlp.layers[-1]
        scaled_w = last.weight * init_scale
        zero_b = jnp.zeros_like(last.bias)
        new_last = eqx.tree_at(
            lambda l: (l.weight, l.bias),
            last,
            (scaled_w, zero_b),
        )
        self.mlp = eqx.tree_at(lambda m: m.layers[-1], self.mlp, new_last)

        self.P = jnp.asarray(P)
        self.ss_state = jnp.asarray(ss_state)
        self.ss_policy = jnp.asarray(ss_policy)
        if output_links is None:
            link_codes_tuple: tuple = (0,) * n_policies
        else:
            if len(output_links) != n_policies:
                raise ValueError(
                    f"output_links length {len(output_links)} != n_policies {n_policies}"
                )
            mapping = {"linear": 0, "log": 1}
            try:
                link_codes_tuple = tuple(mapping[link] for link in output_links)
            except KeyError as e:
                raise ValueError(
                    f"unknown output_link {e.args[0]!r}; "
                    f"valid choices: {list(mapping.keys())}"
                ) from None
            # Log-link requires positive ss_policy at that slot
            ss_arr = jnp.asarray(ss_policy)
            bad_idxs = [
                i
                for i, code in enumerate(link_codes_tuple)
                if code == 1 and float(ss_arr[i]) <= 0.0
            ]
            if bad_idxs:
                raise ValueError(
                    f"output_link='log' requires ss_policy > 0; non-positive at "
                    f"indices {bad_idxs} (ss values {[float(ss_arr[i]) for i in bad_idxs]})"
                )
        self.output_links = link_codes_tuple
        self.policy_lower = _to_tuple(policy_lower)
        self.policy_upper = _to_tuple(policy_upper)

    def _forward_single(self, state: Array) -> Array:
        # stop_gradient on the linearization constants: they are frozen
        # parameters of the architecture, not trainable. Matches the
        # pattern used for input_shift/scale and output bounds in MLP.
        ss_state = jax.lax.stop_gradient(self.ss_state)
        ss_policy = jax.lax.stop_gradient(self.ss_policy)
        P = jax.lax.stop_gradient(self.P)
        bk_corr = P @ (state - ss_state)  # in per-row natural space (level or log)

        delta = self.mlp(state)
        # Per-policy: linear policies use additive, log policies use
        # multiplicative. Resolve at Python time from the static link tuple:
        # if all-linear (common case) skip exp(); if all-log skip where().
        # Mixed case computes both and selects via jnp.where.
        if all(code == 0 for code in self.output_links):
            raw = ss_policy + bk_corr + delta
        elif all(code == 1 for code in self.output_links):
            raw = ss_policy * jnp.exp(bk_corr + delta)
        else:
            # Mixed links: apply exp() ONLY to log-linked outputs. output_links
            # is static (Python tuple), so unroll and never feed exp() the
            # linear-linked exponents — a large linear (bk_corr+delta) would
            # otherwise overflow exp() in the unselected branch and poison the
            # reverse pass with NaN even though the forward correctly selects
            # the linear value (audit JAX-SILENT-05).
            add = bk_corr + delta
            raw = jnp.stack(
                [
                    ss_policy[i] * jnp.exp(add[i])
                    if code == 1
                    else ss_policy[i] + add[i]
                    for i, code in enumerate(self.output_links)
                ]
            )

        if self.policy_lower is not None:
            lower = jax.lax.stop_gradient(_to_array(self.policy_lower))
            raw = jnp.maximum(raw, lower)
        if self.policy_upper is not None:
            upper = jax.lax.stop_gradient(_to_array(self.policy_upper))
            safe_upper = jnp.where(jnp.isinf(upper), jnp.array(1e10), upper)
            raw = jnp.minimum(raw, safe_upper)
        return raw

    def __call__(self, x: Array) -> Array:
        if x.ndim == 1:
            return self._forward_single(x)
        return jax.vmap(self._forward_single)(x)

create_linear_plus_mlp 

create_linear_plus_mlp(
    model,
    hidden_sizes: Sequence[int] = (128, 128),
    activation: str = "tanh",
    init: str = "default",
    init_scale: float = 0.01,
    input_shift: Optional[Array] = None,
    input_scale: Optional[Array] = None,
    output_links: Optional[Sequence[str]] = None,
    *,
    key: Array,
) -> LinearPlusMLP

Factory: build a generic LinearPlusMLP using the model's linearization.

Parameters:

Name Type Description Default
output_links Optional[Sequence[str]]

per-policy parameterization. None defaults to all-linear (legacy behavior). Each entry must be "linear" or "log". Log-output policies require ss_policy[i] > 0.

 None

This factory has NO model-specific knobs. For disaster-specific shape priors (K/F gauge mask, ELB feature, q-as-M reparameterization), use create_disaster_policy_net from deqn_jax.models.disaster.network.

Source code in src/deqn_jax/networks/linear_plus_mlp.py 
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
def create_linear_plus_mlp(
    model,
    hidden_sizes: Sequence[int] = (128, 128),
    activation: str = "tanh",
    init: str = "default",
    init_scale: float = 0.01,
    input_shift: Optional[Array] = None,
    input_scale: Optional[Array] = None,
    output_links: Optional[Sequence[str]] = None,
    *,
    key: Array,
) -> LinearPlusMLP:
    """Factory: build a generic LinearPlusMLP using the model's linearization.

    Args:
        output_links: per-policy parameterization. None defaults to all-linear
            (legacy behavior). Each entry must be ``"linear"`` or ``"log"``.
            Log-output policies require ``ss_policy[i] > 0``.

    This factory has NO model-specific knobs. For disaster-specific shape
    priors (K/F gauge mask, ELB feature, q-as-M reparameterization), use
    ``create_disaster_policy_net`` from ``deqn_jax.models.disaster.network``.
    """
    from deqn_jax.training.linearize import linearize_model

    P_level, _Q = linearize_model(model, verbose=False)
    ss_state, ss_policy = model.steady_state_fn(model.constants)

    if output_links is None:
        # Backward-compat: pass P unchanged, link defaults to all-linear inside __init__.
        P = P_level
    else:
        P = _convert_p_per_link(P_level, ss_policy, output_links)

    return LinearPlusMLP(
        n_states=model.n_states,
        n_policies=model.n_policies,
        hidden_sizes=hidden_sizes,
        activation=activation,
        P=P,
        ss_state=ss_state,
        ss_policy=ss_policy,
        output_links=output_links,
        policy_lower=model.policy_lower,
        policy_upper=model.policy_upper,
        init=init,
        init_scale=init_scale,
        input_shift=input_shift,
        input_scale=input_scale,
        key=key,
    )