Quantization module

This module provides near-optimal quantization of rank-one matrices and butterfly matrices improving bit/accuracy tradeoffs over naive rounding, as illustrated by the following figure:

legend

(left) Accuracy gain achieved by the optimal quantization over the naive strategy that rounds to the nearest (see [1] for more details). We randomly draw \(100\) Butterfly decomposition with dtype = torch.float64 of size \(N=128\), \(256\), \(512\) and \(1024\) and quantize to two different targets torch.bfloat16 and torch.float8_e4m3fn. code to reproduce (right) The quantized operator \(L_q\) offers much better performance thanks to the fact that the computation of \(L_qX\) uses target dtype = torch.bfloat16 instead of base dtype = torch.float64. The first dimension of \(X\) is equal to \(2048\). We use a NVIDIA RTX A6000 with 49GB to run the benchmark. code to reproduce

The following table shows the possible quantization given a real base dtype and a real target dtype.

	target
	torch only		NumPy/CuPy and torch
base	float8_e4m3fn	float8_e5m2	bfloat16	float16	float32	float64
float8_e4m3fn
float8_e5m2
bfloat16
float16
float32
float64

The following table shows the possible quantization given a complex base dtype and a complex target dtype.

	target
	torch only	NumPy/CuPy and torch
base	complex32	complex64	complex128
complex32
complex64
complex128

legend

• If the number of bits in mantissa of the target dtype is greater than the number of bits in mantissa of the base dtype quantization simply returns a casted copy of the inputs.

• If base dtype is a NumPy/CuPy dtype and target is a torch.dtype, the quantization functions (qrank_one(), qbutterfly(), qmonarch()) first convert the input arrays to torch tensors. It is particularly interesting to quantize a NumPy/CuPy array to torch.bfloat16 (10 bits in mantissa) or torch.float8_* format; two formats that do not exist in NumPy/CuPy.

• The quantization from torch.float8_e4m3fn to torch.float8_e5m2 and the reverse are not implemented yet.

As an illustration we show how to quantize the Butterfly decomposition of matrix A from dtype=torch.float32 to torch.bfloat16:

>>> import torch
>>> N = 2 ** 13
>>> base = torch.float32
>>> target = torch.bfloat16
>>> A = torch.randn(N, N).to(dtype=base)
>>> from lazylinop.butterfly import ksd, Chain
>>> L = ksd(A, Chain.square_dyadic((N, N)))
>>> from lazylinop.quantization import qbutterfly
>>> Lq, rerr = qbutterfly(L, target)
>>> L.ks_values[0].dtype
torch.bfloat16
>>> Lq.ks_values[0].dtype
torch.bfloat16

where Lq is the quantized Butterfly decomposition of matrix A.

References

[1] Rémi Gribonval, Theo Mary, Elisa Riccietti. Optimal quantization of rank-one matrices in floating-point arithmetic—with applications to butterfly factorizations. 2023. hal-04125381 https://inria.hal.science/hal-04125381v1/document

[2] Maël Chaumette, Rémi Gribonval, Elisa Riccietti. CROQuant: Complex Rank-One Quantization Algorithm, with Application to Butterfly Factorizations. 2026. hal-05520926 https://hal.science/hal-05520926

Rank one quantization

lazylinop.quantization.qrank_one()

lazylinop.quantization.qrank_one(x, y, target, delta=0, t=None)

Find vectors \(\hat{x}\) and \(\hat{y}\) in the target format (real or complex) such that \(\hat{x}\hat{y}^H\) is closest to \(xy^H\) using algorithms of references [1, 2].

Args:

x, y: 1d array

Column vectors of shape (m, 1) and (n, 1) to be quantized belonging to the same namespace, with the same dtype, raise an Exception otherwize.

target: NumPy/CuPy or torch dtype

Target dtype of the quantization. If target is complex, x and y must be complex too.

delta: int, optional

Parameter that controls the number of breaklines in the complex quantization algorithm. The default value is delta = 0 to only have the centroids on the accumulation lines [2]. delta has no effect when x and y have real entries or when target is not a complex dtype.

t: int, optional

If t is smaller than the number of bits in mantissa of target, use t to quantize and then cast the results \(\hat{x}\) and \(\hat{y}\) to match the dtype given by target. Default value is None corresponding to t = finfo(target).nmant.

Returns:

A tuple (xh, yh, rerr) where xh is \(\hat{x}\) the quantized \(x\), yh is \(\hat{y}\) the quantized \(y\) and rerr the quantization relative error.

Note

• The dtype of xh and yh is target.

• By construction, xh is not necessarily close to x, yh is not necessarily close to y but xh @ yh.H is close to x @ y.H.

• If target is a torch.dtype and x, y are not torch tensors convert x and y before proceeding.

• If x.dtype promotes to target return (xc, yc, rerr) where xc and yc are x and y casted to target dtype.

Examples:

>>> import numpy as np
>>> from lazylinop.quantization import qrank_one
>>> from lazylinop.quantization import chop
>>> m = np.random.randint(2, high=10)
>>> n = np.random.randint(2, high=10)
>>> base = 'float64'
>>> x = np.random.randn(m, 1).astype(base)
>>> y = np.random.randn(n, 1).astype(base)
>>> # Quantize using float16.
>>> target = 'float16'
>>> xh, yh, rerr = qrank_one(x, y, target)
>>> xh.shape == x.shape
True
>>> yh.shape == y.shape
True
>>> xh.dtype == target
True
>>> yh.dtype == target
True
>>> # RTN strategy for comparison.
>>> from lazylinop.quantization import chop, finfo
>>> t = finfo(target).nmant
>>> n = np.linalg.norm(x @ y.T, ord='fro')
>>> xr = chop(x, t)
>>> yr = chop(y, t)
>>> rtn_rerr = np.linalg.norm(x @ y.T - xr @ yr.T, ord='fro') / n
>>> # Optimal error versus RTN error.
>>> bool(rerr <= rtn_rerr)
True

See also

• lazylinop.quantization.qmonarch(),

• lazylinop.quantization.qbutterfly().

References:

[1] Rémi Gribonval, Theo Mary, Elisa Riccietti. Optimal quantization of rank-one matrices in floating-point arithmetic—with applications to butterfly factorizations. 2023. hal-04125381 https://inria.hal.science/hal-04125381v1/document

[2] Maël Chaumette, Rémi Gribonval, Elisa Riccietti. CROQuant: Complex Rank-One Quantization Algorithm, with Application to Butterfly Factorizations. 2026. hal-05520926 https://hal.science/hal-05520926

Monarch quantization

lazylinop.quantization.qmonarch()

lazylinop.quantization.qmonarch(L, target, delta=0, t=None)

Optimal quantization of a Monarch operator L to a target format (real or complex) using algorithms from reference [1, 2] with t = finfo(target).nmant bits in mantissa of target dtype.

Args:

L:

Monarch operator to be quantized, described either as:

• a list of two 4d arrays ks_values of shapes compatible with lazylinop.butterfly.Chain.monarch().

• or a lazy linear operator L obtained either with lazylinop.butterfly.ksm() or with lazylinop.butterfly.ksd() using a Monarch chain.

• an instance of lazylinop.butterfly.KsmLazyLinOp

• an instance of lazylinop.butterfly.KsmPermLazyLinOp

target: NumPy/CuPy or torch dtype

Target dtype of the quantization. If target is complex, L.ks_values must be complex too.

delta: int, optional

Parameter that controls the number of breaklines in the complex quantization algorithm. The default value is delta = 0 to only have the centroids on the accumulation lines [2]. delta has no effect when L has real entries or when target is not a complex dtype.

t: int, optional

If t is smaller than the number of bits in mantissa of target, use t to quantize and then cast the results \(\hat{x}\) and \(\hat{y}\) to match the dtype given by target. Default value is None corresponding to t = finfo(target).nmant.

Returns:

A tuple (Lq, rerr) where Lq is of the same nature as L but with ks_values of dtype target and rerr is the quantization relative error.

Note

• The dtype of ks_values of Lq is target.

• If target is a torch.dtype and ks_values of L are not torch tensors convert ks_values before proceeding.

• If ks_values.dtype of L promotes to target return (Lc, rerr) where Lc is L casted to target dtype.

Examples:

>>> import numpy as np
>>> from lazylinop import LazyLinOp
>>> from lazylinop.butterfly import Chain, ksm
>>> from lazylinop.quantization import qmonarch
>>> M, N = 4, 10
>>> # base and target dtype.
>>> base = 'float64'
>>> target = 'float16'
>>> chain = Chain.monarch((M, N))
>>> ksp = chain.ks_patterns
>>> # List of two random 4d arrays compatible with the Monarch chain.
>>> ksv = [np.random.randn(*ksp[i]).astype(base) for i in range(2)]
>>> Lq, rerr = qmonarch(ksv, target)
>>> isinstance(Lq, list)
True
>>> len(Lq) == 2
True
>>> Lq[0].dtype == target
True
>>> # ksm such that len(L.ks_values) == 2.
>>> L = ksm(ksv, backend='xp')
>>> Lq, rerr = qmonarch(L, target)
>>> isinstance(Lq, LazyLinOp)
True
>>> len(Lq.ks_values) == 2
True
>>> Lq.ks_values[0].dtype == target
True
>>> # Compute L @ X.
>>> X = np.random.randn(L.shape[1], 128).astype(base)
>>> Y = L @ X
>>> # RTN strategy for comparison.
>>> from lazylinop.quantization import chop, finfo
>>> t = finfo(target).nmant
>>> _ksv = [chop(k, t) for k in ksv]
>>> Lr = ksm(_ksv, backend='xp')
>>> Yr = Lr @ X
>>> n = np.linalg.norm(Y, ord='fro')
>>> rtn_rerr = np.linalg.norm(Yr - Y, ord='fro') / n
>>> # Optimal error versus RTN error.
>>> bool(rerr <= rtn_rerr)
True

See also

• lazylinop.butterfly.ksm(),

• lazylinop.butterfly.ksd(),

• lazylinop.butterfly.Chain(),

• lazylinop.quantization.finfo(),

• lazylinop.quantization.qrank_one(),

• lazylinop.quantization.qbutterfly().

References:

[1] Rémi Gribonval, Theo Mary, Elisa Riccietti. Optimal quantization of rank-one matrices in floating-point arithmetic—with applications to butterfly factorizations. 2023. hal-04125381 https://inria.hal.science/hal-04125381v1/document

[2] Maël Chaumette, Rémi Gribonval, Elisa Riccietti. CROQuant: Complex Rank-One Quantization Algorithm, with Application to Butterfly Factorizations. 2026. hal-05520926 https://hal.science/hal-05520926

Butterfly quantization

lazylinop.quantization.qbutterfly()

lazylinop.quantization.qbutterfly(L, target, order='l2r', delta=0, t=None)

Optimal quantization of a Butterfly operator L to a target format (real or complex) using algorithms of references [1, 2] with t = finfo(target).nmant bits in mantissa of target dtype.

Args:

L:

Butterfly operator to be quantized, described either as:

• a list of \(p\) 4d arrays ks_values of shapes compatible with lazylinop.butterfly.Chain.square_dyadic().

• or a lazy linear operator L obtained either with lazylinop.butterfly.ksm() or with lazylinop.butterfly.ksd() using a square-dyadic chain.

• an instance of lazylinop.butterfly.KsmLazyLinOp

• an instance of lazylinop.butterfly.KsmPermLazyLinOp

target: NumPy/CuPy or torch dtype

Target dtype of the quantization. If target is complex, L.ks_values must be complex too.

order: str, optional

Possible choices are 'l2r' (default) corresponding to Algorithm 7.3 or 'pairwise' corresponding to Algorithm 7.2 of reference [1, 2].

delta: int, optional

Parameter that controls the number of breaklines in the complex quantization algorithm. The default value is delta = 0 to only have the centroids on the accumulation lines [2]. delta has no effect when L has real entries or when target is not a complex dtype.

t: int, optional

If t is smaller than the number of bits in mantissa of target, use t to quantize and then cast the results \(\hat{x}\) and \(\hat{y}\) to match the dtype given by target. Default value is None corresponding to t = finfo(target).nmant.

Returns:

A tuple (Lq, rerr) where Lq is of the same nature as L but with ks_values of dtype target and rerr is the quantization relative error.

Note

• The dtype of ks_values of Lq is target.

• If target is a torch.dtype and ks_values of L are not torch tensors convert ks_values before proceeding.

• If ks_values.dtype of L promotes to target return (Lc, rerr) where Lc is L casted to target dtype.

Examples:

>>> import numpy as np
>>> from lazylinop import LazyLinOp
>>> from lazylinop.butterfly import Chain, ksm
>>> from lazylinop.quantization import qbutterfly
>>> p = 4
>>> N = 2 ** p
>>> # base and target dtype.
>>> base = 'float64'
>>> target = 'float16'
>>> # Square-dyadic chain.
>>> chain = Chain.square_dyadic((N, N))
>>> ksp = chain.ks_patterns
>>> # List of p random 4d arrays compatible with square-dyadic chain.
>>> ksv = [np.random.randn(*ksp[i]).astype(base) for i in range(p)]
>>> Lq, rerr = qbutterfly(ksv, target, 'l2r')
>>> len(Lq) == p
True
>>> Lq[0].dtype == target
True
>>> # ksm such that len(L.ks_values) == p.
>>> L = ksm(ksv, backend='xp')
>>> Lq, rerr = qbutterfly(L, target, 'l2r')
>>> isinstance(Lq, LazyLinOp)
True
>>> len(Lq.ks_values) == p
True
>>> Lq.ks_values[0].dtype == target
True
>>> # Compute L @ X.
>>> X = np.random.randn(L.shape[1], 128).astype(base)
>>> Y = L @ X
>>> # RTN strategy for comparison.
>>> from lazylinop.quantization import chop, finfo
>>> t = finfo(target).nmant
>>> _ksv = [chop(k, t) for k in ksv]
>>> Lr = ksm(_ksv, backend='xp')
>>> Yr = Lr @ X
>>> n = np.linalg.norm(Y, ord='fro')
>>> rtn_rerr = np.linalg.norm(Yr - Y, ord='fro') / n
>>> # Optimal error versus RTN error.
>>> bool(rerr <= rtn_rerr)
True

See also

• lazylinop.butterfly.ksm(),

• lazylinop.butterfly.ksd(),

• lazylinop.butterfly.Chain(),

• lazylinop.quantization.finfo(),

• lazylinop.quantization.qrank_one(),

• lazylinop.quantization.qmonarch().

References:

[1] Rémi Gribonval, Theo Mary, Elisa Riccietti. Optimal quantization of rank-one matrices in floating-point arithmetic—with applications to butterfly factorizations. 2023. hal-04125381 https://inria.hal.science/hal-04125381v1/document

[2] Maël Chaumette, Rémi Gribonval, Elisa Riccietti. CROQuant: Complex Rank-One Quantization Algorithm, with Application to Butterfly Factorizations. 2026. hal-05520926 https://hal.science/hal-05520926

Utils

1. lazylinop.quantization.chop()

2. lazylinop.quantization.upcast_downcast()

3. lazylinop.quantization.finfo()

4. lazylinop.quantization.promote_types()

lazylinop.quantization.chop(x, target=24, stochastic=False, fixed=False)

Round matrix elements using round(x * pow(2, t - e)) * pow(2, e - t) where t is the number of bits in the mantissa of target dtype and where e = floor(log2(x) + 1).

Args:

x: array

chop(x, t) is the matrix obtained by rounding the elements of x to t significant binary places.

target: NumPy/CuPy/torch dtype

It can be either:

• Number of bits in mantissa. target must be lesser than the number of bits in mantissa of x.dtype`. Default value is 24, corresponding to IEEE single precision.

• A NumPy/CuPy/torch dtype. The number of bits in mantissa of target must be lesser than the number of bits in mantissa of x.dtype`.

stochastic: bool, optional

Quantizes the elements of the complex array x with t significant bits with a stochastic rounding strategy. The default value if False.

fixed: bool, optional

Performs fixed-point quantization of the complex array x with t significant bits. The default value if False.

Returns:

An array obtained by rounding the elements of x to t significant binary places.

lazylinop.quantization.upcast_downcast(x, target)

Round-as the element of x.

Args:

x: array

upcast_downcat(x, target) is the matrix obtained by down-casting followed by an up-casting of the elements of x to t significant binary places. t is the number of bits in mantissa of target.

target: NumPy/CuPy or torch dtype

Down-cast x to target.

Returns:

An array of dtype x.dtype with rounded elements according to target dtype.

lazylinop.quantization.finfo(dtype)

Get info about the numerical properties of dtype. dtype is a floating point type. If dtype is a torch.dtype add nmant and nexp attributes.

Args:

dtype:

Get info about this dtype.

Returns:

An object that stores the numerical properties of dtype.

Examples:

>>> import torch
>>> from lazylinop.quantization import finfo
>>> f = finfo(torch.bfloat16)
>>> f.nmant
7

See also

• NumPy finfo,

• PyTorch finfo.

lazylinop.quantization.promote_types(t1, t2)

Promote to the type in which both t1 and t2 may be safely cast.

Args:

t1, t2: NumPy/CuPy or torch dtype

t1 and t2 are two dtypes.

Returns:

If the two types share the same namespace and t1 and t2 are not torch.float8_* return the promoted type, otherwise return None.

Examples:

>>> import numpy as np
>>> t1 = np.dtype('float16')
>>> t2 = np.dtype('float64')
>>> promote_types(t1, t2) == t2
True
>>> import torch
>>> t1 = torch.float16
>>> t2 = torch.bfloat16
>>> promote_types(t1, t2) == torch.float32
True
>>> t1 = torch.float32
>>> t2 = torch.bfloat16
>>> promote_types(t1, t2) == t1
True
>>> t1 = torch.float8_e4m3fn
>>> t2 = torch.float32
>>> promote_types(t1, t2) is None
True

See also

• NumPy promote,

• PyTorch promote.