boundlab.zono.softmax_handler#

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

Zonotope softmax transformer built from primitive handlers.

Softmax is decomposed as:

\[\mathrm{softmax}(x)_j = \frac{\exp(x_j)}{\sum_k \exp(x_k)}\]

The implementation applies: exp -> reduce-sum -> reciprocal -> element-wise product. For stability, it first shifts by the center maximum along the softmax dimension. Currently, only 2D inputs with dim == 1 are supported.

Parameters:

  • x (Expr) – Input expression with shape (m, n).

  • 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])