boundlab.zono.softmax_handler#

boundlab.zono.softmax_handler(x, dim=-1, dtype=None)[source]#

Zonotope softmax transformer using the DeepT decomposition.

Softmax is rewritten as:

\[\mathrm{softmax}(x)_i = \frac{\exp(x_i)}{\sum_j \exp(x_j)} = \frac{1}{\sum_j \exp(x_j - x_i)}\]

so only subtraction, exp, reduce-sum, and reciprocal are required (no bilinear element-wise product).

Parameters:

  • x (Expr) – Input expression.

  • dim (int) – Dimension along which to apply softmax (default: -1).

  • dtype – Ignored (for API compatibility with torch.softmax).

Returns:

An expression over-approximating torch.softmax(x, dim=dim).

Return type:

Expr

Examples

>>> import torch
>>> import boundlab.expr as expr
>>> from boundlab.zono.softmax import softmax_handler
>>> x = expr.ConstVal(torch.zeros(2, 3)) + 0.1 * expr.LpEpsilon([2, 3])
>>> y = softmax_handler(x, dim=1)
>>> y.shape
torch.Size([2, 3])