Sketching API

Sketching API#

The panther.sketch module provides sketching operators for dimensionality reduction and fast linear algebra.

Sketching Operators#

panther.sketch.dense_sketch_operator(m, n, distribution, device=None, dtype=None)[source]

Creates a dense sketch operator matrix with entries sampled from a specified distribution. This function generates a dense random matrix of shape (m, n) where each entry is drawn independently from the given distribution. Such sketch operators are commonly used in randomized linear algebra, dimensionality reduction, and compressed sensing to project high-dimensional data into a lower-dimensional space while preserving certain geometric properties.

Parameters:

  • m – int The number of rows of the sketch operator (i.e., the target dimension after sketching).

  • n – int The number of columns of the sketch operator (i.e., the original dimension of the data).

  • distribution – DistributionFamily The distribution family from which to sample the entries of the sketch operator. This could be, for example, a standard normal distribution, a uniform distribution, or any other supported distribution.

  • device – Optional[torch.device], default=None The device on which to allocate the resulting tensor (e.g., ‘cpu’ or ‘cuda’). If None, defaults to the current device.

  • dtype – Optional[torch.dtype], default=None The desired data type of the returned tensor. If None, defaults to the default dtype of the current torch device.

Returns:

torch.Tensor

A dense tensor of shape (m, n) with entries sampled from the specified distribution.

Example

>>> import torch
>>> from panther.sketch import dense_sketch_operator
>>> from panther.sketch import DistributionFamily
>>> m, n = 100, 500
>>> sketch = dense_sketch_operator(m, n, DistributionFamily.GAUSSIAN)
>>> print(sketch.shape)
torch.Size([100, 500])

Notes

  • The choice of distribution affects the properties of the sketch operator. For example, using a standard normal distribution yields a Johnson-Lindenstrauss transform.

  • For large matrices, consider using sparse sketch operators for improved efficiency.

panther.sketch.sparse_sketch_operator(m, n, vec_nnz, axis, device=None, dtype=None, seed=None)[source]

Creates a sparse sketch operator matrix with specified number of non-zero entries per vector.

This function generates a sparse random matrix of shape (m, n) where each row or column (depending on the axis) has exactly vec_nnz non-zero entries. This is commonly used in randomized linear algebra and sketching algorithms to project high-dimensional data into a lower-dimensional space while preserving certain geometric properties.

Parameters:

  • m – int The number of rows of the sketch operator (i.e., the target dimension after sketching).

  • n – int The number of columns of the sketch operator (i.e., the original dimension of the data).

  • vec_nnz – int The number of non-zero entries per vector (row or column) in the sketch operator.

  • axis – Axis The axis along which to create the sparse sketch operator. Can be either Axis.Short or Axis.Long.

  • device – Optional[torch.device], default=None The device on which to allocate the resulting tensor (e.g., ‘cpu’ or ‘cuda’). If None, defaults to the current device.

  • dtype – Optional[torch.dtype], default=None The desired data type of the returned tensor. If None, defaults to the default dtype of the current torch device.

  • seed – Optional[int], default=None Integer seed for the internal C++ random number generator used to construct the sparse pattern. When None (default) the seed is drawn from torch’s current RNG, so calling torch.manual_seed(...) before this function produces a reproducible sketch. Pass an explicit integer to bypass PyTorch’s RNG entirely.

Returns:

torch.Tensor

A sparse COO tensor of shape (m, n) with exactly vec_nnz non-zero entries per vector along the specified axis.

Example

>>> import torch
>>> from panther.sketch import sparse_sketch_operator, Axis
>>> m, n = 100, 500
>>> vec_nnz = 5
>>> # Reproducible via torch seed
>>> torch.manual_seed(42)
>>> sketch = sparse_sketch_operator(m, n, vec_nnz, Axis.Short)
>>> print(sketch.shape)
torch.Size([100, 500])
>>> print(sketch._nnz())  # Number of non-zero entries in the sparse tensor
panther.sketch.sketch_tensor(input, axis, new_size, distribution=None, sketch_matrix=None, device=None, dtype=None)[source]

Sketches a given input tensor along a specified axis using a randomized projection technique.

This function reduces the dimensionality of the input tensor along the specified axis to new_size by applying a sketching (random projection) operation. The type of random projection is determined by the distribution_or_sketch_matrix parameter, which can be either a distribution family or a precomputed sketching matrix. If a distribution family is provided, the function returns a tuple (sketched_tensor, sketch_matrix). If a sketching matrix (torch.Tensor) is provided, only the sketched tensor is returned.

Parameters:

  • input (torch.Tensor) – The input tensor to be sketched. Can be of any shape, but the dimension along axis will be reduced.

  • axis (int) – The axis along which to apply the sketching operation. Must be a valid axis for input.

  • new_size (int) – The target size for the specified axis after sketching. Must be less than or equal to the original size.

  • distribution_or_sketch_matrix (Union[DistributionFamily, torch.Tensor]) – Either the distribution family to use for generating the sketching matrix (e.g., Gaussian, Rademacher), or a precomputed sketching matrix (torch.Tensor) to use directly.

  • device (Optional[torch.device], default=None) – The device on which to perform the computation and allocate the sketching matrix. If None, uses the device of input.

  • dtype (Optional[torch.dtype], default=None) – The desired data type of the output tensor and sketching matrix. If None, uses the dtype of input.

Returns:

If a distribution family is provided, returns a tuple (sketched_tensor, sketch_matrix). If a sketching matrix is provided, returns only the sketched tensor.

Return type:

Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]

Example

>>> import torch
>>> from panther.sketch.core import sketch_tensor, DistributionFamily
>>> x = torch.randn(10, 100)
>>> # Using a distribution family
>>> sketched_x, sketch_matrix = sketch_tensor(
...     input=x,
...     axis=1,
...     new_size=20,
...     distribution_or_sketch_matrix=DistributionFamily.GAUSSIAN
... )
>>> print(sketched_x.shape)
torch.Size([10, 20])
>>> print(sketch_matrix.shape)
torch.Size([20, 100])
>>> # Using a precomputed sketching matrix
>>> from panther.sketch import dense_sketch_operator
>>> S = dense_sketch_operator(100, 20, DistributionFamily.GAUSSIAN)
>>> sketched_x = sketch_tensor(
...     input=x,
...     axis=1,
...     new_size=20,
...     distribution_or_sketch_matrix=S
... )
>>> print(sketched_x.shape)
torch.Size([10, 20])

Notes

  • The sketching operation is commonly used for dimensionality reduction, speeding up computations, and preserving essential structure in large-scale machine learning and signal processing tasks.

  • The choice of distribution affects the quality and properties of the sketch.

  • The returned sketching matrix can be reused for consistent projections or analysis.

panther.sketch.scaled_sign_sketch(m, n, device=None, dtype=None)[source]

Generates a scaled sign sketch matrix as a PyTorch tensor.

A scaled sign sketch is a random projection matrix where each entry is independently set to +1 or -1, scaled by a normalization factor. This is commonly used in randomized linear algebra and sketching algorithms.

Parameters:

  • m (int) – Number of rows in the output tensor (sketch dimension).

  • n (int) – Number of columns in the output tensor (original dimension).

  • device (Optional[torch.device], optional) – The device on which to create the tensor. Defaults to None (uses current device).

  • dtype (Optional[torch.dtype], optional) – The desired data type of returned tensor. Defaults to None (uses default dtype).

Returns:

A tensor of shape (m, n) containing the scaled sign sketch matrix.

Return type:

torch.Tensor

Example

>>> import torch
>>> from panther.sketch import scaled_sign_sketch
>>> m, n = 32, 128
>>> S = scaled_sign_sketch(m, n)
>>> print(S.shape)
torch.Size([32, 128])
>>> # Each entry is either +1/sqrt(m) or -1/sqrt(m)
>>> print(torch.unique(S))
tensor([-0.1768,  0.1768])  # For m=32, 1/sqrt(32) ≈ 0.1768

Specific Sketching Methods#

panther.sketch.gaussian_skop(m, d, device=None, dtype=None)[source]

Generates a Gaussian random projection matrix of shape (m, d).

Each entry in the sketching matrix is an independent sample from the normal distribution N(0, 1/m). This type of projection approximately preserves Euclidean distances and inner products with high probability.

Parameters:

  • m (int) – Number of rows in the sketching matrix (target dimension).

  • d (int) – Number of columns in the sketching matrix (original dimension).

Returns:

A dense sketching matrix of shape (m, d) with entries ~ N(0, 1/m).

Return type:

torch.Tensor

Example

>>> S = gaussian_skop(100, 1000)
>>> X = torch.randn(64, 1000)
>>> X_sketched = X @ S.t()
>>> X_sketched.shape
torch.Size([64, 100])
panther.sketch.srht(x, m)[source]

Subsampled Randomized Hadamard Transform (SRHT).

This function computes a Subsampled Randomized Hadamard Transform of a 1D input tensor x of length n (where n must be a power of 2). It performs the following steps:

  1. Multiply the input by a random diagonal matrix with ±1 entries (random sign flipping).

  2. Apply the Fast Walsh-Hadamard Transform (FWHT).

  3. Uniformly subsample m rows from the result.

This is commonly used in randomized numerical linear algebra and compressed sensing to reduce dimensionality while approximately preserving distances.

Parameters:

  • x (torch.Tensor) – A 1D input tensor of shape (n,) where n is a power of 2. The tensor must reside on the CPU and be of floating point type.

  • m (int) – The number of rows to subsample from the Hadamard-transformed output. Must satisfy 0 < m <= n.

Returns:

A 1D tensor of shape (m,) containing the subsampled rows from the Hadamard-transformed vector.

Return type:

torch.Tensor

Raises:

RuntimeError – If x is not on CPU, is not 1-dimensional, or n is not a power of 2. If m is greater than n.

Example

>>> import torch
>>> from panther.sketch import srht
>>> x = torch.randn(8)
>>> y = srht(x, 4)
>>> print(y.shape)
torch.Size([4])
panther.sketch.sjlt_skop(m, d, sparsity=2, device=None, dtype=None)[source]

Generates a sparse Johnson-Lindenstrauss Transform (SJLT) sketching matrix.

Each column of the matrix has exactly sparsity non-zero entries, chosen uniformly at random among the rows, with values ±1. The matrix is scaled by 1/√sparsity. This sketch is ideal for fast projections with low memory usage.

Parameters:

  • m (int) – Number of rows in the sketching matrix (target dimension).

  • d (int) – Number of columns in the sketching matrix (original dimension).

  • sparsity (int, optional) – Number of non-zero entries per column. Defaults to 2.

Returns:

A sparse sketching matrix of shape (m, d), entries ∈ {±1/√sparsity}.

Return type:

torch.Tensor

Example

>>> S = sjlt_skop(200, 1000, sparsity=3)
>>> X = torch.randn(32, 1000)
>>> X_sketched = X @ S.t()
>>> X_sketched.shape
torch.Size([32, 200])
panther.sketch.count_skop(m, d, device=None, dtype=None)[source]

Generates a CountSketch sketching matrix of shape (m, d).

Each column of the sketching matrix has exactly one non-zero entry, which is randomly assigned to a row via a simple hash (uniform random), and assigned a value of ±1 with equal probability. The result is a sparse sketching operator used to project high-dimensional data into a lower-dimensional space while approximately preserving inner products.

This is commonly used in streaming and randomized linear algebra algorithms.

Parameters:

  • m (int) – Number of rows in the sketching matrix (target dimension).

  • d (int) – Number of columns in the sketching matrix (original dimension).

Returns:

A sparse sketching matrix of shape (m, d) with exactly one

non-zero entry per column (either +1 or -1).

Return type:

torch.Tensor

Example

>>> import torch
>>> from panther.sketch import count_sketch_operator
>>> m, d = 100, 1000
>>> S = count_skop(m, d)
>>> X = torch.randn(50, d)  # Input data
>>> X_sketched = X @ S.t()
>>> X_sketched.shape
torch.Size([50, 100])

Examples#

Basic Sketching

import torch
import panther as pr

# Create a random matrix to sketch
A = torch.randn(1000, 500)

# Create Gaussian sketching matrix (100 x 1000)
S = pr.sketch.dense_sketch_operator(
    m=100,                                    # output dimension
    n=1000,                                   # input dimension
    distribution=pr.sketch.DistributionFamily.Gaussian
)

# Apply sketching: SA has shape (100, 500)
SA = S @ A
print(f"Original: {A.shape}, Sketched: {SA.shape}")

Sparse Sketching for Memory Efficiency

# Sparse sketching matrix (much less memory)
# Reproducible: torch.manual_seed() controls the sketch when seed=None (default)
torch.manual_seed(42)
S_sparse = pr.sketch.sparse_sketch_operator(
    m=200,                              # output dimension
    n=5000,                             # input dimension
    vec_nnz=3,                          # non-zeros per column
    axis=pr.sketch.Axis.Short
)

# Or pass an explicit seed to bypass PyTorch's RNG
S_fixed = pr.sketch.sparse_sketch_operator(
    m=200, n=5000, vec_nnz=3,
    axis=pr.sketch.Axis.Short,
    seed=123
)

# Apply to large matrix
A_large = torch.randn(5000, 2000)
SA_sparse = S_sparse @ A_large
print(f"Sparse sketched: {SA_sparse.shape}")

Direct Tensor Sketching

# Sketch along different axes
A = torch.randn(1000, 800)

# Sketch rows (reduce first dimension)
sketched_rows = pr.sketch.sketch_tensor(
    input=A,
    axis=0,                             # sketch along rows
    new_size=200,                       # new row count
    distribution=pr.sketch.DistributionFamily.Gaussian
)
print(f"Row sketched: {sketched_rows[0].shape}")  # (200, 800)

# Sketch columns (reduce second dimension)
sketched_cols = pr.sketch.sketch_tensor(
    input=A,
    axis=1,                             # sketch along columns
    new_size=300,                       # new column count
    distribution=pr.sketch.DistributionFamily.Uniform
)
print(f"Column sketched: {sketched_cols[0].shape}")  # (1000, 300)

Subsampled Randomized Hadamard Transform (SRHT)

# SRHT is very fast for power-of-2 dimensions
# SRHT operates on 1D tensors, not matrices
x = torch.randn(1024)  # dimension must be power of 2

# Apply SRHT to reduce dimensionality
x_sketched = pr.sketch.srht(
    x=x,                                # input 1D tensor
    m=256                               # output dimension
)

print(f"SRHT sketched: {x_sketched.shape}")  # (256,)

Scaled Sign Sketching

# Generate scaled sign (Rademacher-like) sketching matrix
sign_matrix = pr.sketch.scaled_sign_sketch(
    m=64,                               # sketch dimension (rows)
    n=256                               # original dimension (columns)
)

print(f"Sign matrix: {sign_matrix.shape}")  # (64, 256)
print(f"Values are ±1/√m: {sign_matrix[0, :5]}")

Distribution Families#

Gaussian Sketching

  • Pros: Best theoretical guarantees, universally applicable

  • Cons: Requires more memory (dense matrices)

  • Use when: Maximum accuracy is needed

gaussian_sketch = pr.sketch.dense_sketch_operator(
    100, 500, pr.sketch.DistributionFamily.Gaussian
)

Uniform Sketching

  • Pros: Faster to generate, good performance

  • Cons: Slightly different properties than Gaussian

  • Use when: Speed is more important than optimal theoretical constants

uniform_sketch = pr.sketch.dense_sketch_operator(
    100, 500, pr.sketch.DistributionFamily.Uniform
)

Sparse Sketching

  • Pros: Very memory efficient, can handle huge matrices

  • Cons: May require larger sketch size for same accuracy

  • Use when: Working with very large matrices

sparse_sketch = pr.sketch.sparse_sketch_operator(
    100, 500, vec_nnz=3, axis=pr.sketch.Axis.Rows
)

SRHT (Subsampled Randomized Hadamard Transform)

  • Pros: Extremely fast, good for structured matrices

  • Cons: Requires power-of-2 dimensions, operates on 1D tensors only

  • Use when: Input dimension is power of 2

x = torch.randn(1024)  # 1024 = 2^10
x_srht = pr.sketch.srht(x, 128)

Performance Comparison#

Sketching matrix generation time (for 1000×5000 matrix):

Method

Time (ms)

Memory (MB)

Quality

Gaussian

45.2

19.1

Excellent

Uniform

15.2

19.1

Very Good

Sparse (nnz=3)

8.4

0.6

Good

SRHT

2.1

0.1

Very Good

GPU Acceleration#

All sketching operations support GPU:

device = torch.device('cuda')

# Create sketching matrix on GPU
S = pr.sketch.dense_sketch_operator(
    200, 1000,
    pr.sketch.DistributionFamily.Gaussian,
    device=device
)

# Apply to GPU tensor
A = torch.randn(1000, 800, device=device)
SA = S @ A  # Computed on GPU

print(f"Result device: {SA.device}")

Best Practices#

1. Choose sketch size appropriately

# Rule of thumb: sketch_size ≈ 2-4 × target_rank
import numpy as np
target_rank = 50
sketch_size = 4 * target_rank  # 200

# For ε-accuracy: sketch_size ≥ k + log(1/ε)
epsilon = 0.01
sketch_size_theory = target_rank + int(np.log(1/epsilon))

2. Use sparse sketching for very large matrices

# For matrices larger than available memory
very_large_sketch = pr.sketch.sparse_sketch_operator(
    m=1000,
    n=1000000,    # 1M dimensions
    vec_nnz=4,    # Only 4 non-zeros per column
    axis=pr.sketch.Axis.Long
)