Quantize embedding (#994)

* quantize embedding * rename as_linear + comment * consistency in docs * fix test
2025-09-26 15:58:14 +08:00 · 2024-04-15 16:42:10 -07:00
parent 2e7c02d5cd
commit cd9e184529
9 changed files with 269 additions and 54 deletions
--- a/docs/src/python/nn.rst
+++ b/docs/src/python/nn.rst
@@ -173,6 +173,7 @@ In detail:
   :toctree: _autosummary

   value_and_grad
+   quantize

 .. toctree::

--- a/docs/src/python/nn/layers.rst
+++ b/docs/src/python/nn/layers.rst
@@ -31,6 +31,7 @@ Layers
   Mish
   MultiHeadAttention
   PReLU
+   QuantizedEmbedding
   QuantizedLinear
   RMSNorm
   ReLU
@@ -43,4 +44,4 @@ Layers
   Softshrink
   Step
   Transformer
-   Upsample
+   Upsample
--- a/docs/src/python/tree_utils.rst
+++ b/docs/src/python/tree_utils.rst
@@ -19,3 +19,4 @@ return python trees will be using the default python ``dict``, ``list`` and
   tree_flatten
   tree_unflatten
   tree_map
+   tree_map_with_path