Map PyTorch: 'aten:bitwise_and', 'aten:bitwise_cpu' to NNEF.
Map aten:bitwise_left_shift / << op to NNEF -> tract_shl.
Tract's shift ops live under the tract_core extension despite
the bare tract_shl / tract_shr names in the registry.
Map PyTorch: 'aten:bitwise_not', 'aten:bitwise_not_cpu' to NNEF.
On bool inputs, PyTorch's ~ is semantically a logical not, so we emit
the standard NNEF not op (keeps the graph portable and self-documenting,
rather than relying on tract's bitnot happening to do the right thing on
bool). For integer inputs, emit tract_core_bitnot for true bitwise
inversion.
Map PyTorch: 'aten:bitwise_or' to NNEF.
Map PyTorch: 'aten:bitwise_right_shift' / >> op -> tract_shr.
Map PyTorch: 'aten:bitwise_xor' to NNEF.
Map PyTorch: 'aten:cdist' to NNEF via a fragment.
cdist(a, b, p) computes the pairwise distance matrix between
rows of a (shape (..., M, D)) and b (shape (..., N, D)):
out[..., i, j] = (sum(|a[..., i, :] - b[..., j, :]|^p))^(1/p).
The fragment broadcasts via unsqueeze (one new axis on each
input) and reduces along the trailing feature axis. Pure stdlib.
Map PyTorch: 'aten:cosine_similarity' to NNEF via a fragment.
The fragment lives in op/fragment/cosine_similarity.nnef and is
composed only of NNEF stdlib ops, so no tract-side change is needed.
Negative dim is normalized to a non-negative index via
pick_axis before reaching the fragment: under dynamic-axes
mode tract's reduce path crashes (index out of bounds at
core/src/ops/nn/reduce.rs) when handed a negative axis against a
symbolic-rank tensor.
Map PyTorch: 'aten:cross' to NNEF via a fragment.
cross(a, b, dim) is the 3-D vector cross product along dim.
The fragment slices each input along dim into its three
components and computes the standard (a1*b2 - a2*b1, a2*b0 -
a0*b2, a0*b1 - a1*b0) triplet. Pure stdlib.
Map PyTorch: aten::cummax(self, dim) -> (values, indices).
Map PyTorch: aten::cummin(self, dim) -> (values, indices).
Map PyTorch: 'aten:cumprod' to NNEF using a scan fragment.
Mirror of cumsum with a mul scan body and an init of 1 (built
pointwise via mul(first, 0) + 1 to keep init shape-matching).
Map PyTorch: 'aten:cumsum' to NNEF using a scan fragment (Tract).
- Implemented via
tract_core_scaninsidefragment cumsum(axis=0). - For arbitrary dim, transpose input to bring that axis to 0, apply fragment, then transpose back.
Map aten::diff(input, n, dim, prepend?, append?) to NNEF.
n-th order finite differences along dim: each step replaces
x with x[1:] - x[:-1] along dim. prepend / append are
not supported (raise T2NErrorNotImplemented).
aten::digamma -> psi(x) = (d/dx) log(Gamma(x)).
Asymptotic series after shifting x up by 6 via the recurrence
psi(x+1) = psi(x) + 1/x. Valid for x > 0; negative inputs would
need the reflection formula (left as a follow-up).
Map PyTorch: aten::dist(input, other, p) to NNEF via fragment.
Scalar (sum_{all axes} |a - b|^p)^(1/p) between two broadcastable
tensors. p must be finite and > 0; the inf / -inf / 0 norms would
need max_reduce / min_reduce / count_nonzero branches (raise
T2NErrorNotImplemented for those).
Map PyTorch: 'aten::floor_divide' / 'aten::floordiv' to NNEF.
JIT records aten::floordiv for Python //; upstream
normalize_ops.cpp does not bridge it to floor_divide, so we
alias it on our side.
Map aten::frexp(input) -> (mantissa, exponent) to NNEF.
The exponent output is int32; the fragment uses tract_core_cast
so this handler is tract-only.
Map PyTorch: 'aten:hardshrink' to the hardshrink fragment.
aten::heaviside -> heaviside fragment (step with tie-breaker).
aten::i0 / aten::special_i0 -> Bessel I_0(x).
Abramowitz & Stegun polynomial approximation; two branches at
|x| = 3.75.
aten::i1 / aten::special_i1 -> Bessel I_1(x).
Abramowitz & Stegun 9.8.3 / 9.8.4 polynomial branches; same
structure as i0 with the sign carried explicitly so the function
remains odd.
aten::isclose -> |a - b| <= atol + rtol * |b|.
Reads optional rtol / atol from the trace (defaults match
torch's 1e-5 / 1e-8 via the fragment defaults). equal_nan=True
is not yet handled -- rare in real traces; raises if set.
aten::lgamma -> log-Gamma via Lanczos.
Numerical Recipes Lanczos approximation (g = 5, N = 6). Valid for
x > 0.5; smaller arguments need a reflection branch (left as a
follow-up).
Map PyTorch: 'aten:log_sigmoid' to the log_sigmoid fragment.
aten::logaddexp -> numerically-stable log(exp a + exp b).
Map PyTorch: aten::logcumsumexp(self, dim).
Numerically-stable scan with two state vars (running_max,
running_sum_shifted). Init: finfo.min and 0 so the first
step naturally produces out[0] = input[0].
Map PyTorch: 'aten:logical_xor' to NNEF.
Map PyTorch: 'aten:logsumexp' to the logsumexp fragment.
PyTorch signature: logsumexp(input, dim, keepdim=False).
The fragment always reduces the named axis (no keepdim); when the
user asks for keepdim=True we follow up with an unsqueeze on
the same axis to reinstate it.
aten::mvlgamma -> multivariate log-Gamma.
mvlgamma(x, p) = (p*(p-1)/4) * log(pi)
+ sum_{i=1..p} lgamma(x + (1-i)/2).
p is a static int from the trace; we unroll the sum into p
lgamma fragment calls plus a single constant offset. Inherits
lgamma's domain restriction (each shifted argument must stay
> 0.5).
Map PyTorch: 'aten:nan_to_num' to NNEF.
Decomposed to pure NNEF stdlib (ne for NaN via the IEEE-754
NaN != NaN invariant, gt/lt against the dtype's finite
range for +/-inf, plus select). No tract extension needed.
Defaults match torch: NaN -> 0; +inf -> finfo.max; -inf -> finfo.min.
Map PyTorch: 'aten:outer' to NNEF.
torch.outer(a, b) over 1-D inputs is a[:, None] * b[None, :].
Lower to two unsqueezes and a broadcasting mul.
Axes are kept positive. Tract's NNEF unsqueeze deserializer
(tract_core::ops::change_axes::AxisOp::change_shape) does not
normalize negative axes and panics with smallvec: index exceeds
length; verified across tract 0.20.22 through 0.23.0-dev.5. This
matches the wider t2n convention: the dedicated unsqueeze op
handler also normalizes via pick_axis.
Map PyTorch: 'aten:pairwise_distance' to NNEF via a fragment.
pairwise_distance(a, b, p, eps, keepdim) computes
(sum(|a - b + eps|^p, axis=-1))^(1/p). Torch keeps the reduced
last axis only when keepdim=True; the fragment squeezes it
unconditionally and we re-unsqueeze after the call when needed.
Pure NNEF stdlib (sub / abs / pow / sum_reduce / squeeze).
Map PyTorch: aten::pdist(self, p) to NNEF.
torch.pdist(x, p) returns a 1-D tensor of length N*(N-1)/2
holding the pairwise p-norm distances between rows of the
rank-2 input x (shape (N, D)). We reuse the existing cdist
fragment to build the (N, N) distance matrix against x itself,
flatten it, and gather the strict upper-triangle entries
i*N + j for i < j. Static N only (the upper-triangle index
constant is baked at export time).
Map PyTorch: 'aten:remainder' to NNEF.
aten::signbit -> x < 0 (bool).
NNEF can't see the IEEE-754 sign bit, so signbit(-0.0) returns
False here (vs PyTorch's True); same caveat as copysign /
atan2.
Map PyTorch: 'aten:softshrink' to the softshrink fragment.
aten::special_entr -> -x * log(x) (0 at x=0).
See special_entr.nnef for the eps-clamped formulation that keeps
log(0) from poisoning the discarded select branch.
aten::special_i0e -> exp(-|x|) * I_0(x).
Same Abramowitz & Stegun polynomial branches as i0 but the
large-x branch drops exp(|x|) so the result stays finite for
arbitrarily large |x|.
aten::special_i1e -> exp(-|x|) * I_1(x).
Same polynomial branches as i1; the large-|x| branch drops
exp(|x|) so the result stays finite for arbitrarily large |x|.
aten::special_xlog1py -> x * log(1 + y) (0 at x=0).
See special_xlog1py.nnef for the eps-clamped log argument.
Map PyTorch: 'aten:square' to NNEF.
x.square() is x * x; the dedicated NNEF sqr op exists in some
runtimes but the simplest portable form is a mul with the same
tensor on both sides. square preserves dtype, so no int->float
promotion is needed.
Map PyTorch: 'aten:std_mean' (returns (std, mean)) to NNEF.
Map PyTorch: 'aten:tensordot' to NNEF via tract_core_einsum.
tensordot(a, b, dims_a, dims_b) contracts the paired axes (same
size on each side) and produces an output of rank
a.rank + b.rank - 2 * len(dims_a). The einsum expression is
built so each contracted axis-pair shares a label.
Map aten::trapezoid(y, dx_or_x, dim) (alias trapz) to NNEF.
Uniform-dx case only: the second arg must be a scalar float
constant (the dx overload). The tensor-x overload is rejected
for now (would need a (x[1:] - x[:-1]) multiply that broadcasts
against y[1:] + y[:-1]).
Map PyTorch: 'aten:var' to NNEF.
Centered second moment along dim with arbitrary correction
(denominator = N - correction). Lowered to mean_reduce + sub + sqr
+ (sum_reduce / mean_reduce) so any correction value works without
relying on NNEF's var fragment (which always squeezes and so
can't honor keepdim=True).
Map PyTorch: 'aten:var_mean' (returns (var, mean)) to NNEF.