Quantize embedding (#994)

* quantize embedding * rename as_linear + comment * consistency in docs * fix test
2025-09-01 04:24:36 +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/_templates/nn-module-template.rst
+++ b/docs/src/_templates/nn-module-template.rst
@@ -0,0 +1,20 @@
+{{ fullname | escape | underline}}
+
+.. currentmodule:: {{ module }}
+
+.. autoclass:: {{ objname }}
+
+   {% block methods %}
+
+   {% if methods %}
+   .. rubric:: {{ _('Methods') }}
+
+   .. autosummary::
+   {% for item in methods %}
+      {%- if item not in inherited_members and item != "__init__" %}
+         ~{{ name }}.{{ item }}
+      {%- endif %}
+   {%- endfor %}
+   {% endif %}
+   {% endblock %}
+
--- 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