.. |check| replace:: :octicon:`check-circle;1em;sd-text-success`
.. |x| replace:: :octicon:`x-circle;1em;sd-text-info`
.. |gear| replace:: :octicon:`gear;1em;sd-text-danger`

Quantization module
-------------------

.. automodule:: lazylinop.quantization

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:

.. image:: _static/qbutterfly_accuracy_gain.svg
   :height: 275
   :align: left

.. image:: _static/benchmark_quantization_N2048.svg
   :height: 275
   :align: right

.. admonition:: legend

   **(left)** Accuracy gain achieved by the *optimal* quantization over the naive strategy that *rounds to the nearest* (see :ref:`[1] <references>` 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``. :download:`code to reproduce <../../benchmark/quantization/accuracy_gain.py>` **(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. :download:`code to reproduce <../../benchmark/quantization/benchmark_quantization.py>`

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

.. table::
   :align: center
   :widths: auto

   +-------------------+----------------------------------------------------------------------+
   |                   |                              target                                  |
   +                   +-----------------------------+----------------------------------------+
   |                   | torch only                  |     NumPy/CuPy and torch               |
   +-------------------+---------------+-------------+----------+---------+---------+---------+
   | base              | float8_e4m3fn | float8_e5m2 | bfloat16 | float16 | float32 | float64 |
   +===================+===============+=============+==========+=========+=========+=========+
   | **float8_e4m3fn** | |check|       | |gear|      | |x|      | |x|     | |x|     | |x|     |
   +-------------------+---------------+-------------+----------+---------+---------+---------+
   | **float8_e5m2**   | |gear|        | |check|     | |x|      | |x|     | |x|     | |x|     |
   +-------------------+---------------+-------------+----------+---------+---------+---------+
   | **bfloat16**      | |check|       | |check|     | |check|  | |check| | |x|     | |x|     |
   +-------------------+---------------+-------------+----------+---------+---------+---------+
   | **float16**       | |check|       | |check|     | |x|      | |check| | |x|     | |x|     |
   +-------------------+---------------+-------------+----------+---------+---------+---------+
   | **float32**       | |check|       | |check|     | |check|  | |check| | |check| | |x|     |
   +-------------------+---------------+-------------+----------+---------+---------+---------+
   | **float64**       | |check|       | |check|     | |check|  | |check| | |check| | |check| |
   +-------------------+---------------+-------------+----------+---------+---------+---------+

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

.. table::
   :align: center
   :widths: auto

   +-------------------+----------------------------------------+
   |                   |                              target    |
   +                   +-------------+--------------------------+
   |                   | torch only  |     NumPy/CuPy and torch |
   +-------------------+-------------+-------------+------------+
   | base              | complex32   | complex64   | complex128 |
   +===================+=============+=============+============+
   | **complex32**     | |check|     | |x|         | |x|        |
   +-------------------+-------------+-------------+------------+
   | **complex64**     | |check|     | |check|     | |x|        |
   +-------------------+-------------+-------------+------------+
   | **complex128**    | |check|     | |check|     | |check|    |
   +-------------------+-------------+-------------+------------+

.. admonition:: legend

   - |x| 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.
   - |check| If *base* ``dtype`` is a NumPy/CuPy ``dtype`` and *target* is a ``torch.dtype``, the quantization functions (:py:func:`qrank_one`, :py:func:`qbutterfly`, :py:func:`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.
   - |gear| 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``:

.. code-block:: python

    >>> 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:

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
~~~~~~~~~~~~~~~~~~~~~

:meth:`lazylinop.quantization.qrank_one`

.. autofunction:: lazylinop.quantization.qrank_one

Monarch quantization
~~~~~~~~~~~~~~~~~~~~

:meth:`lazylinop.quantization.qmonarch`

.. autofunction:: lazylinop.quantization.qmonarch

Butterfly quantization
~~~~~~~~~~~~~~~~~~~~~~

:meth:`lazylinop.quantization.qbutterfly`

.. autofunction:: lazylinop.quantization.qbutterfly

Utils
~~~~~

1. :meth:`lazylinop.quantization.chop`
2. :meth:`lazylinop.quantization.upcast_downcast`
3. :meth:`lazylinop.quantization.finfo`
4. :meth:`lazylinop.quantization.promote_types`

.. autofunction:: lazylinop.quantization.chop
.. autofunction:: lazylinop.quantization.upcast_downcast
.. autofunction:: lazylinop.quantization.finfo
.. autofunction:: lazylinop.quantization.promote_types