boundlab.diff.op.DiffLinear

class boundlab.diff.op.DiffLinear

Bases: Module

Two parallel linear layers whose outputs are paired via diff_pair().

At the concrete-tensor level this is equivalent to running fc1(x) (fc2’s output is discarded at runtime via the diff_pair no-op). When the model is exported and interpreted by the differential interpreter (e.g. boundlab.diff.zono3.interpret), the diff_pair node is lifted into a DiffExpr2 that tracks both branches simultaneously.

Parameters:

• fc1 – First linear layer.

• fc2 – Second linear layer; must have the same in_features, out_features, and dtype as fc1.

Examples

>>> import torch
>>> from torch import nn
>>> from boundlab.diff.op import DiffLinear
>>> fc1 = nn.Linear(4, 3)
>>> fc2 = nn.Linear(4, 3)
>>> model = DiffLinear(fc1, fc2)
>>> out = model(torch.zeros(4))
>>> out.shape
torch.Size([3])

Methods

__init__
forward

__init__(fc1, fc2)

forward(x)

__call__(*args, **kwargs)

Call self as a function.

static __new__(cls, *args, **kwargs)