omnibias.torch.activations.nqs
PyTorch activation derivatives, trainable operator layers, network jets, and direct differential-operator contractions.
import omnibias.torch.activations.nqsModule contract & mathematical conventions
NQS-friendly activations (lattice work, mirrored from JAX).
These activations are the canonical building blocks of neural-quantum-state ansatze on lattice systems:
================== ========================================= ===================
Activation K=2 collapse / role Max fastpath order
================== ========================================= ===================
log_cosh log(cosh(z)): smooth |z|; the standard n in {0, 1, 2, 3} log-amplitude in spin-NQS (Carleo-Troyer RBM family). 1st derivative is tanh.
softabs sqrt(z^2 + eps^2) - eps: smooth |z| n in {0, 1, 2} used in Jastrow factors and complex amplitude magnitudes.
smooth_sign tanh(z / T): temperature-controlled all orders smooth sign; used in variational annealing schedules.
mish z * tanh(softplus(z)): self-gated all orders residual activation, transformer FFN.
================== ========================================= ===================
smooth_sign is tanh rescaled, so it inherits the whole Riccati tower;
mish is the analytic product z * g(z), so Leibniz over the (t, s)
recurrence gives every order. Both were once capped and are no longer.
The PyTorch implementations here mirror omnibias.jax.activations so
both backends register the same set of names; the parity test in
tests/test_jax_parity.py then verifies bit-stable agreement.
Shared formulas do not imply universal floating-point identity. Verify dtype, shape, derivative-order support, and the stated numerical assumptions for the path you use.
LOG_COSH
omnibias.torch.activations.nqs.LOG_COSHSource ↗LOG_COSH = register_activation(ActivationSpec(name='log_cosh', forward=_log_cosh_forward, derivative=_log_cosh_derivative, fastpath=_log_cosh_fastpath, riccati_polynomial=None, noise_model='laplace_smoothed', operator_role='Smooth-|z| activation: log(cosh(z)) ~ |z| for large |z|, ~ z^2/2 for small. Standard log-amplitude in spin-NQS; first derivative is tanh, second is 1 - tanh^2.', aliases=('logcosh',)))This exported constant is defined by the expression above. Its exact value and annotation come from the pinned source; use it with the module contract.
Read the implementation
LOG_COSH = register_activation(
ActivationSpec(
name="log_cosh",
forward=_log_cosh_forward,
derivative=_log_cosh_derivative,
fastpath=_log_cosh_fastpath,
riccati_polynomial=None,
noise_model="laplace_smoothed",
operator_role=(
"Smooth-|z| activation: log(cosh(z)) ~ |z| for large |z|, "
"~ z^2/2 for small. Standard log-amplitude in spin-NQS; "
"first derivative is tanh, second is 1 - tanh^2."
),
aliases=("logcosh",),
)
)MISH
omnibias.torch.activations.nqs.MISHSource ↗MISH = register_activation(ActivationSpec(name='mish', forward=_mish_forward, derivative=_mish_derivative, fastpath=_mish_fastpath, riccati_polynomial=None, noise_model='none', operator_role='Mish: z * tanh(softplus(z)). Self-gated residual activation.'))This exported constant is defined by the expression above. Its exact value and annotation come from the pinned source; use it with the module contract.
Read the implementation
MISH = register_activation(
ActivationSpec(
name="mish",
forward=_mish_forward,
derivative=_mish_derivative,
fastpath=_mish_fastpath,
riccati_polynomial=None,
noise_model="none",
operator_role=("Mish: z * tanh(softplus(z)). Self-gated residual activation."),
)
)SMOOTH_SIGN
omnibias.torch.activations.nqs.SMOOTH_SIGNSource ↗SMOOTH_SIGN = register_activation(ActivationSpec(name='smooth_sign', forward=_smooth_sign_forward, derivative=_smooth_sign_derivative, fastpath=_smooth_sign_fastpath, integral=_smooth_sign_integral, riccati_polynomial=None, noise_model='symmetric_bernoulli_tempered', operator_role='Temperature-controlled smooth sign: tanh(z/T) -> sign(z) as T -> 0. The beta-tempered sign surrogate; used in variational annealing schedules.', aliases=('soft_sign',), limit_pos_inf=1.0, limit_neg_inf=-1.0))This exported constant is defined by the expression above. Its exact value and annotation come from the pinned source; use it with the module contract.
Read the implementation
SMOOTH_SIGN = register_activation(
ActivationSpec(
name="smooth_sign",
forward=_smooth_sign_forward,
derivative=_smooth_sign_derivative,
fastpath=_smooth_sign_fastpath,
integral=_smooth_sign_integral,
riccati_polynomial=None,
noise_model="symmetric_bernoulli_tempered",
operator_role=(
"Temperature-controlled smooth sign: tanh(z/T) -> sign(z) as "
"T -> 0. The beta-tempered sign surrogate; used in variational "
"annealing schedules."
),
aliases=("soft_sign",),
limit_pos_inf=1.0,
limit_neg_inf=-1.0,
)
)SOFTABS
omnibias.torch.activations.nqs.SOFTABSSource ↗SOFTABS = register_activation(ActivationSpec(name='softabs', forward=_softabs_forward, derivative=_softabs_derivative, fastpath=_softabs_fastpath, integral=_softabs_integral, riccati_polynomial=None, noise_model='huber_smoothed', operator_role='Smooth absolute value: sqrt(z^2 + eps^2) - eps. C^2 everywhere; the eps-tempered abs surrogate (-> |z| as eps -> 0); useful in Jastrow factors and complex-amplitude magnitudes.', aliases=('soft_abs',)))This exported constant is defined by the expression above. Its exact value and annotation come from the pinned source; use it with the module contract.
Read the implementation
SOFTABS = register_activation(
ActivationSpec(
name="softabs",
forward=_softabs_forward,
derivative=_softabs_derivative,
fastpath=_softabs_fastpath,
integral=_softabs_integral,
riccati_polynomial=None,
noise_model="huber_smoothed",
operator_role=(
"Smooth absolute value: sqrt(z^2 + eps^2) - eps. C^2 everywhere; "
"the eps-tempered abs surrogate (-> |z| as eps -> 0); useful in "
"Jastrow factors and complex-amplitude magnitudes."
),
aliases=("soft_abs",),
)
)