Operators

As part of the xlnstorch package, we provide a set of operations that are analogous to PyTorch’s built-in operations. These operations are registered with PyTorch’s internal dispatch mechanism, allowing them to be used in the same way as PyTorch’s built-in operations. These operations can be accessed via the traditional PyTorch mechanism, such as torch.add() or torch.matmul().

C++ Implementations

For operations that are computationally intensive, we provide C++ implementations that can be used to accelerate the computation. These implementations are available if the package is built with C++ extensions enabled and are enabled by default. You can toggle the use of C++ implementations with the following function:

from xlnstorch.operators import toggle_cpp_implementations

toggle_cpp_implementations(True) # enable C++ implementations
toggle_cpp_implementations(False) # disable C++ implementations

Currently, there are C++ implementations for the following operations:

• Addition

• Summation

• Matrix Multiplication

• Convolution functions

Custom SBDB functions

In LNS, numbers are represented by their logarithms. To implement addition and subtraction, we work with their logarithmic forms and a special function, commonly known as the “sum/difference in the log domain” (sbdb) function, which is closely related to Gaussian logarithms.

Suppose two numbers \(X\) and \(Y\) are represented as:

\[\begin{split}x &= \log_B \lvert X \rvert \\ y &= \log_B \lvert Y \rvert\end{split}\]

To compute \(\log_B \left( \vert X \rvert + \lvert Y \rvert \right)\) and \(\log_B \left( \vert X \rvert - \lvert Y \rvert \right)\), we use the identities:

\[\begin{split}\log_B \left( \lvert X \rvert + \lvert Y \rvert \right) &= x + s_B \left( y - x \right) \\ \log_B \left( \lvert X \rvert - \lvert Y \rvert \right) &= x + d_B \left( y - x \right)\end{split}\]

where the “sum” and “difference” helper functions are defined as:

\[\begin{split}s_B \left( z \right) &= \log_B \left( 1 + B ^ {z} \right) \\ d_B \left( z \right) &= \log_B \lvert 1 - B ^ {z} \rvert\end{split}\]

However, computing \(s_B(z)\) and \(d_B(z)\) directly is computationally expensive, especially on hardware that does not support efficient logarithm and exponentiation. For this reason, xlnstorch provides fast, approximate implementations of these functions to accelerate LNS addition and subtraction.

By using these approximations, we achieve a good trade-off between numerical accuracy and computational efficiency in LNS arithmetic within the package. See xlnsconf for more details on these implementations and the papers that describe them.

Tab

The ‘tab’ method provides fast LNS addition and subtraction by precomputing the values of \(s_B(z)\) and \(d_B(z)\) and storing them in lookup tables. During computation, these tables are used to quickly retrieve approximate results instead of calculating the logarithms and exponentials directly. This approach enables rapid and efficient evaluation of LNS arithmetic operations with minimal computational cost, at the expense of increased memory usage and a fixed precision determined by the table resolution.

To use the ‘tab’ method, you must first initialize the lookup table with the desired base or precision, and a filestem to store the table.

import xlnstorch as xltorch

xltorch.operators.implementations.tab.get_table("filestem", f=10)
xltorch.set_default_sbdb_implementation("tab")

a = xltorch.lnstensor([1.0, 2.0], f=10)
b = xltorch.lnstensor([3.0, 4.0], f=10)
c = torch.add(a, b) # uses the tab implementation

Utah-Tayco

The ‘utah_tayco’ method is an approximate implementation of the sbdb function that uses unpartitioned linear Taylor interpolation and/or cotransformation of the Gaussian logarithm.

To use the ‘utah_tayco’ method, there is no initialization required and can be used as follows:

import xlnstorch as xltorch

xltorch.set_default_sbdb_implementation("utah_tayco")

a = xltorch.lnstensor([1.0, 2.0], f=10)
b = xltorch.lnstensor([3.0, 4.0], f=10)
c = torch.add(a, b) # uses the utah_tayco implementation

Internal Operators

These are the internal operations performed on torch.Tensor internal representations of LNSTensor objects related to arithmetic operations. These operations are useful if you want to implement your own custom operations. For the most part, these internal operator functions wrap the apply_lns_op() function.

Arithmetic Operations

lns_add	See docs for torch.add() for parameter/return details.
lns_sub	See docs for torch.sub() for parameter/return details.
lns_mul	See docs for torch.mul() for parameter/return details.
lns_div	See docs for torch.div() for parameter/return details.
lns_neg	See docs for torch.neg() for parameter/return details.
lns_abs	See docs for torch.abs() for parameter/return details.
lns_sqrt	See docs for torch.sqrt() for parameter/return details.
lns_square	See docs for torch.square() for parameter/return details.
lns_pow	See docs for torch.pow() for parameter/return details.
lns_exp	See docs for torch.exp() for parameter/return details.
lns_log	See docs for torch.log() for parameter/return details.
lns_reciprocal	See docs for torch.reciprocal() for parameter/return details.
lns_sign	See docs for torch.sign() for parameter/return details.
lns_positive	See docs for torch.positive() for parameter/return details.
lns_sum	See docs for torch.sum() for parameter/return details.
lns_prod	See docs for torch.prod() for parameter/return details.
lns_mean	See docs for torch.mean() for parameter/return details.
lns_var	See docs for torch.var() for parameter/return details.
lns_matmul	See docs for torch.matmul() for parameter/return details.
lns_transpose	See docs for torch.transpose() for parameter/return details.

Comparison Operations

lns_equal	See docs for torch.equal() for parameter/return details.
lns_eq	See docs for torch.eq() for parameter/return details.
lns_ne	See docs for torch.ne() for parameter/return details.
lns_ge	See docs for torch.ge() for parameter/return details.
lns_gt	See docs for torch.gt() for parameter/return details.
lns_le	See docs for torch.le() for parameter/return details.
lns_lt	See docs for torch.lt() for parameter/return details.
lns_isclose	See docs for torch.isclose() for parameter/return details.
lns_allclose	See docs for torch.allclose() for parameter/return details.
lns_any	See docs for torch.any() for parameter/return details.
lns_all	See docs for torch.all() for parameter/return details.
lns_isin	See docs for torch.isin() for parameter/return details.
lns_sort	See docs for torch.sort() for parameter/return details.
lns_argsort	See docs for torch.argsort() for parameter/return details.
lns_kthvalue	See docs for torch.kthvalue() for parameter/return details.
lns_maximum	See docs for torch.maximum() for parameter/return details.
lns_minimum	See docs for torch.minimum() for parameter/return details.

Miscellaneous Operations

lns_broadcast_to	See docs for torch.broadcast_to() for parameter/return details.
lns_clone	See docs for torch.clone() for parameter/return details.
lns_squeeze	See docs for torch.squeeze() for parameter/return details.
lns_unsqueeze	See docs for torch.unsqueeze() for parameter/return details.
lns_stack	See docs for torch.stack() for parameter/return details.
lns_cat	See docs for torch.cat() for parameter/return details.
lns_chunk	See docs for torch.chunk() for parameter/return details.
lns_where	See docs for torch.where() for parameter/return details.
lns_pad	See docs for torch.nn.functional.pad() for parameter/return details.

Loss Operations

Note that in practice, you should use the standard PyTorch loss classes, such as torch.nn.MSELoss, which are already implemented to work with LNSTensor objects. However, if you want to implement your own custom loss functions, or want more control over the loss computation, you can use the following functions.

lns_mse_loss	See docs for torch.nn.functional.mse_loss() for parameter/return details.
lns_l1_loss	See docs for torch.nn.functional.l1_loss() for parameter/return details.
lns_binary_cross_entropy	See docs for torch.nn.functional.binary_cross_entropy() for parameter/return details.
lns_binary_cross_entropy_with_logits	See docs for torch.nn.functional.binary_cross_entropy_with_logits() for parameter/return details.
lns_nll_loss	See docs for torch.nn.functional.nll_loss() for parameter/return details.
lns_poisson_nll_loss	See docs for torch.nn.functional.poisson_nll_loss() for parameter/return details.
lns_hinge_embedding_loss	See docs for torch.nn.functional.hinge_embedding_loss() for parameter/return details.
lns_kl_div	See docs for torch.nn.functional.kl_div() for parameter/return details.
lns_margin_ranking_loss	See docs for torch.nn.functional.margin_ranking_loss() for parameter/return details.
lns_gaussian_nll_loss	See docs for torch.nn.functional.gaussian_nll_loss() for parameter/return details.
lns_huber_loss	See docs for torch.nn.functional.huber_loss() for parameter/return details.
lns_smooth_l1_loss	See docs for torch.nn.functional.smooth_l1_loss() for parameter/return details.
lns_cross_entropy	See docs for torch.nn.functional.cross_entropy() for parameter/return details.

Activation Operations

Again, in practice, you should use the standard PyTorch activation classes, such as torch.nn.ReLU, which are already implemented to work with LNSTensor objects.

lns_relu	See docs for torch.nn.functional.relu() for parameter/return details.
lns_relu_	See docs for torch.relu_() for parameter/return details.
lns_leaky_relu	See docs for torch.nn.functional.leaky_relu() for parameter/return details.
lns_leaky_relu_	See docs for torch._C._nn.leaky_relu_() for parameter/return details.
lns_threshold	See docs for torch.nn.functional._threshold() for parameter/return details.
lns_threshold_	See docs for torch.threshold_() for parameter/return details.
lns_tanh	See docs for torch.nn.functional.tanh() for parameter/return details.
lns_sigmoid	See docs for torch.nn.functional.sigmoid() for parameter/return details.
lns_logsigmoid	See docs for torch._C._nn.log_sigmoid() for parameter/return details.
lns_softmin	See docs for torch.nn.functional.softmin() for parameter/return details.
lns_softmax	See docs for torch.nn.functional.softmax() for parameter/return details.
lns_log_softmax	See docs for torch.nn.functional.log_softmax() for parameter/return details.
lns_hardtanh	See docs for torch.nn.functional.hardtanh() for parameter/return details.
lns_hardswish	See docs for torch.nn.functional.hardswish() for parameter/return details.
lns_elu	See docs for torch.nn.functional.elu() for parameter/return details.
lns_selu	See docs for torch.nn.functional.selu() for parameter/return details.
lns_celu	See docs for torch.nn.functional.celu() for parameter/return details.
lns_prelu	See docs for torch.prelu() for parameter/return details.
lns_rrelu	See docs for torch.nn.functional.rrelu() for parameter/return details.
lns_glu	See docs for torch.nn.functional.glu() for parameter/return details.
lns_hardshrink	See docs for torch.hardshrink() for parameter/return details.
lns_tanhshrink	See docs for torch.nn.functional.tanhshrink() for parameter/return details.
lns_softsign	See docs for torch.nn.functional.softsign() for parameter/return details.
lns_softplus	See docs for torch._C._nn.softplus() for parameter/return details.
lns_softshrink	See docs for torch._C._nn.softshrink() for parameter/return details.
lns_hardsigmoid	See docs for torch.nn.functional.hardsigmoid() for parameter/return details.
lns_silu	See docs for torch.nn.functional.silu() for parameter/return details.

Layer Operations

As per usual, you should use the standard PyTorch layer classes, such as torch.nn.Linear, which support LNSTensor objects.

lns_linear	See docs for torch._C._nn.linear() for parameter/return details.
lns_bilinear	See docs for torch.bilinear() for parameter/return details.
lns_dropout	See docs for torch.nn.functional.dropout() for parameter/return details.
lns_dropout1d	See docs for torch.nn.functional.dropout1d() for parameter/return details.
lns_dropout2d	See docs for torch.nn.functional.dropout2d() for parameter/return details.
lns_dropout3d	See docs for torch.nn.functional.dropout3d() for parameter/return details.
lns_conv1d	See docs for torch.conv1d() for parameter/return details.
lns_conv2d	See docs for torch.conv2d() for parameter/return details.
lns_conv3d	See docs for torch.conv3d() for parameter/return details.
lns_avg_pool1d	See docs for torch.avg_pool1d() for parameter/return details.
lns_avg_pool2d	See docs for torch._C._nn.avg_pool2d() for parameter/return details.
lns_avg_pool3d	See docs for torch._C._nn.avg_pool3d() for parameter/return details.
lns_adaptive_avg_pool1d	See docs for torch.adaptive_avg_pool1d() for parameter/return details.
lns_adaptive_avg_pool2d	See docs for torch.nn.functional.adaptive_avg_pool2d() for parameter/return details.
lns_adaptive_avg_pool3d	See docs for torch.nn.functional.adaptive_avg_pool3d() for parameter/return details.
lns_batch_norm	See docs for torch.nn.functional.batch_norm() for parameter/return details.
lns_layer_norm	See docs for torch.nn.functional.layer_norm() for parameter/return details.
lns_max_pool1d	See docs for torch.nn.functional.max_pool1d() for parameter/return details.
lns_max_pool2d	See docs for torch.nn.functional.max_pool2d() for parameter/return details.
lns_max_pool3d	See docs for torch.nn.functional.max_pool3d() for parameter/return details.