Visualization#

The xlnstorch.viz submodule provides tools for visualizing and analyzing LNS operations and their numerical properties. This module includes functionality for creating error heatmaps, generating autograd graphs, and analyzing the precision characteristics of LNS arithmetic operations.

Quick Start#

Here’s a simple example of visualizing the error characteristics of an LNS operation:

import torch
import xlnstorch as xltorch
from xlnstorch.viz import make_error_grid, plot_error_heatmap
import matplotlib.pyplot as plt

# Create error grid for multiplication operation
xs, ys, errors = make_error_grid(
    torch.mul,
    x_range=(-2.0, 2.0),
    y_range=(-2.0, 2.0),
    steps=100,
    f=8,
)

# Plot the error heatmap
fig, ax = plt.subplots(figsize=(8, 6))
plot_error_heatmap(errors, xs, ys, ax=ax)
plt.show()

For visualizing autograd graphs with LNS tensors:

from xlnstorch.viz.graph import make_autograd_graph

# Create some LNS tensors with gradients
x = xltorch.randn(3, 3, requires_grad=True)
y = xltorch.randn(3, 3, requires_grad=True)

# Perform operations
z = x * y + x.pow(2)
loss = z.sum()

# Generate autograd graph
graph = make_autograd_graph(loss, params={'x': x, 'y': y})
graph.render('autograd_graph', format='svg')

Overview#

The visualization module is organized into these main areas:

Error Analysis Functions#

The error analysis tools help understand the numerical behavior of LNS operations compared to exact arithmetic.

make_error_grid

Generate a uniformly-sampled grid of differences between an xlnstorch operation op and an exact reference computed with Decimal.

plot_error_heatmap

Visualise the error tensor returned by make_error_grid.

plot_staircase

Plot staircase functions for different LNS precisions on a logarithmic x-axis.

plot_spacing_heatmap

Create a heatmap showing a heatmap of the spacing differences for various LNS precisions.

Precision Comparison Functions#

The precision comparison tools allow you to analyze how different LNS configurations affect numerical accuracy for the same operation.

precision_sweep_analysis

Analyze how different precision levels affect operation accuracy.

plot_precision_comparison

Plot how error metrics change with precision level.

plot_precision_heatmap_grid

Create a grid of heatmaps showing error patterns at different precisions.

Autograd visualization#

The autograd visualization tools provide insight into the computational graph structure when using LNS tensors.

graph.make_autograd_graph

Build (and return) a graphviz.Digraph object that visualizes the PyTorch autograd graph for the given variables.

Dependencies#

The visualization module has optional dependencies:

  • matplotlib: Required for the plotting functions.

  • graphviz: Required for make_autograd_graph

Install with:

pip install matplotlib graphviz

Note that for graphviz, you may also need to install the system graphviz libraries:

# On Ubuntu/Debian
sudo apt-get install graphviz

# On macOS with Homebrew
brew install graphviz

# On Windows
# Download from https://graphviz.org/download/