address reviews

Author: JyotinderSingh

guides/ipynb/writing_quantization_compatible_layers.ipynb

@@ -191,7 +191,7 @@
     "\n",
     "- `self._kernel` as an INT8 vector of shape `(input_dim,)` (the same shape as\n",
     "  the original full-precision kernel).\n",
-    "- `self.scale` as the scalar quantization scale in the layer's compute dtype,\n",
+    "- `self.scale` as the scalar quantization scale in the layer's variable dtype,\n",
     "  which is FP32 in this case."
    ]
   },
@@ -801,4 +801,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 0
-}
+}

guides/md/writing_quantization_compatible_layers.md

@@ -137,7 +137,7 @@ INT8 variables. It should allocate:
 
 - `self._kernel` as an INT8 vector of shape `(input_dim,)` (the same shape as
   the original full-precision kernel).
-- `self.scale` as the scalar quantization scale in the layer's compute dtype,
+- `self.scale` as the scalar quantization scale in the layer's variable dtype,
   which is FP32 in this case.
 
 

guides/writing_quantization_compatible_layers.py

@@ -135,7 +135,7 @@ def quantize(self, mode, **kwargs):
 
 - `self._kernel` as an INT8 vector of shape `(input_dim,)` (the same shape as
   the original full-precision kernel).
-- `self.scale` as the scalar quantization scale in the layer's compute dtype,
+- `self.scale` as the scalar quantization scale in the layer's variable dtype,
   which is FP32 in this case.
 """