quantize-py 0.3.0__tar.gz → 0.3.1__tar.gz

{quantize_py-0.3.0 → quantize_py-0.3.1}/Cargo.lock

@@ -1411,14 +1411,14 @@ dependencies = [
 
 [[package]]
 name = "quantize"
-version = "0.3.0"
+version = "0.3.1"
 dependencies = [
  "half",
 ]
 
 [[package]]
 name = "quantize-py"
-version = "0.3.0"
+version = "0.3.1"
 dependencies = [
  "half",
  "numpy",

{quantize_py-0.3.0 → quantize_py-0.3.1}/Cargo.toml

@@ -17,16 +17,13 @@ include = [
     "/README.md",
 ]
 
-[package.metadata.docs.rs]
-features = ["std"]
-
 [workspace]
 members = ["python"]
 resolver = "3"
 
 [workspace.package]
 # Shared with the Python package. Between releases it ends in -dev.
-version = "0.3.0"
+version = "0.3.1"
 
 [workspace.dependencies]
 half = "2"
@@ -38,8 +35,8 @@ all = { level = "deny", priority = -1 }
 
 [features]
 default = ["std"]
-# Reserved for future no_std support. For now it only implements
-# std::error::Error for Error.
+# Does nothing, since the crate always needs the standard library. Kept so
+# that a Cargo.toml that names it still builds.
 std = []
 
 [dependencies]

{quantize_py-0.3.0 → quantize_py-0.3.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: quantize-py
-Version: 0.3.0
+Version: 0.3.1
 Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
@@ -27,7 +27,7 @@ python bindings for [quantize](https://github.com/aksheyd/quantize), a simple, f
 pip install quantize-py
 ```
 
-requires python 3.12 or newer. numpy is installed with it.
+requires python 3.12 or newer. numpy is installed with it. where no wheel fits, like free-threaded python or alpine linux, pip builds it from source, which needs rust 1.88 or newer.
 
 ```python
 from quantize import Scale, quantize
@@ -48,12 +48,16 @@ the other schemes return the same `Quantized` type:
 - `asymmetric.quantize(weights, bits=8, block=32)` adds a zero-point per block, for values that aren't centered on zero
 - `adaptive.quantize(weights, tolerance=0.1 * weights.std())` gives each block the fewest bits, from 2 to 8, that round every weight within `tolerance`, in the weights' own units. a tenth of their standard deviation gives about 5 bits a block. for a list, use `np.std(weights)`
 - `learned.refine(q, weights)` refits each block's scale, and its zero-point if it has one, to lower the mean squared error. it changes `q` in place, so call `q.copy()` first to keep the original
-- `learned.alternate(q, weights)` refits too, then rounds each value to the nearest code on its block's new line, and repeats until no code moves. it also changes `q` in place. both can raise the worst error past an adaptive tensor's tolerance
+- `learned.alternate(q, weights)` refits too, then rounds each value to the nearest code on its block's new line, and repeats until no code moves. it also changes `q` in place. both can raise the worst error past an adaptive tensor's tolerance, and lowering the error of the weights doesn't always lower the error of a model's outputs, so check those too
 - `Scheme.Q4_32.quantize(weights)` picks a scheme at run time
 
 a block with outliers can need more than 8 bits, which raises `ToleranceTooTightError`. retrying with its `smallest_tolerance` works, but loosens every block, not just that one:
 
 ```python
+import numpy as np
+from quantize import ToleranceTooTightError, adaptive
+
+weights = np.array(weights)  # a list has no .std(). skip this line for a pytorch tensor
 try:
     q = adaptive.quantize(weights, tolerance=0.1 * weights.std())
 except ToleranceTooTightError as error:
@@ -71,7 +75,7 @@ np.savez("layer.npz", **{name: part for name, part in parts.items() if part is n
 q = Quantized.from_parts(**np.load("layer.npz"))
 ```
 
-to load quantized values from a `torch.save` checkpoint, call `torch.serialization.add_safe_globals([Quantized])` before `torch.load`.
+to keep quantized values in a `torch.save` checkpoint, store `torch.frombuffer(bytearray(q.to_bytes()), dtype=torch.uint8)`, and load each back with `Quantized(t)`. the checkpoint is then as small as the bytes, and `torch.load` reads it without `add_safe_globals`. pickling `q` itself makes the checkpoint about 1.5 times larger, since `torch.save` stores bytes as text, and needs `torch.serialization.add_safe_globals([Quantized])` before `torch.load`.
 
 each value decodes as `code * scale`, or `(code - zero_point) * scale` with zero-points, using the scale and zero-point of its block. codes are signed and `bits` wide, and `q.codes` packs them low bits first. scales can be negative, since a symmetric block puts its value farthest from zero on the most negative code. `help(Quantized)` has the details.
 

{quantize_py-0.3.0 → quantize_py-0.3.1}/README.md

@@ -59,7 +59,7 @@ against [candle](https://github.com/huggingface/candle)'s `Q4_0`, `Q5_0`, and `Q
 ### quality
 
 ```
-cargo run --release --example compare
+cargo run --release -p benchmarks --example compare
 ```
 
 quantize two random matrices, reconstruct them, then matmul. mse is the mean squared error against the exact f32 result: smaller is better, and it varies a little between runs. bits/value counts the scales too: 4-bit codes plus one f16 scale per 32 values is 4 + 16/32 = 4.5.
@@ -74,15 +74,21 @@ quantize two random matrices, reconstruct them, then matmul. mse is the mean squ
 
 <!-- comparison:end -->
 
-on WikiText-2 with SmolLM-135M, 4-bit perplexity is 26.44 against candle `Q4_0`'s 26.46, and fp32 is 18.93. lower is better, and `cargo run --release --example wikitext --features benchmarks/workload` reproduces them.
+on WikiText-2 with SmolLM-135M, 4-bit perplexity is 26.44 against candle `Q4_0`'s 26.46, and fp32 is 18.93. perplexity is roughly how many tokens the model is choosing between when it guesses the next one, so lower is better. reproducing these numbers downloads the model and dataset from hugging face, then scores the whole test set, which takes hours:
+
+```
+cargo run --release -p benchmarks --example wikitext --features workload
+```
+
+add `-- --max-tokens 512` to score only the first 512 tokens. that takes about a minute, though its numbers won't match.
 
 ### speed
 
 ```
-cargo run --release --example throughput
+cargo run --release -p benchmarks --example throughput
 ```
 
-quantize and dequantize with f16 scales. both libraries allocate their output on every call. aarch64 is an apple M5 Max and x86_64 an intel xeon. the hand-written simd only targets aarch64, so x86_64 runs plain loops and is slower.
+quantize and dequantize with f16 scales. both libraries allocate their output on every call. aarch64 is an apple M5 Max and x86_64 an intel xeon. the hand-written simd only targets aarch64, so x86_64 runs plain loops and is slower. the run also times this crate's `dot`, `matmul` on 16 inputs, and row-by-row decoding of an adaptive matrix, which aren't compared with candle.
 
 <!-- speed:start -->
 

{quantize_py-0.3.0 → quantize_py-0.3.1}/python/README.md

@@ -6,7 +6,7 @@ python bindings for [quantize](https://github.com/aksheyd/quantize), a simple, f
 pip install quantize-py
 ```
 
-requires python 3.12 or newer. numpy is installed with it.
+requires python 3.12 or newer. numpy is installed with it. where no wheel fits, like free-threaded python or alpine linux, pip builds it from source, which needs rust 1.88 or newer.
 
 ```python
 from quantize import Scale, quantize
@@ -27,12 +27,16 @@ the other schemes return the same `Quantized` type:
 - `asymmetric.quantize(weights, bits=8, block=32)` adds a zero-point per block, for values that aren't centered on zero
 - `adaptive.quantize(weights, tolerance=0.1 * weights.std())` gives each block the fewest bits, from 2 to 8, that round every weight within `tolerance`, in the weights' own units. a tenth of their standard deviation gives about 5 bits a block. for a list, use `np.std(weights)`
 - `learned.refine(q, weights)` refits each block's scale, and its zero-point if it has one, to lower the mean squared error. it changes `q` in place, so call `q.copy()` first to keep the original
-- `learned.alternate(q, weights)` refits too, then rounds each value to the nearest code on its block's new line, and repeats until no code moves. it also changes `q` in place. both can raise the worst error past an adaptive tensor's tolerance
+- `learned.alternate(q, weights)` refits too, then rounds each value to the nearest code on its block's new line, and repeats until no code moves. it also changes `q` in place. both can raise the worst error past an adaptive tensor's tolerance, and lowering the error of the weights doesn't always lower the error of a model's outputs, so check those too
 - `Scheme.Q4_32.quantize(weights)` picks a scheme at run time
 
 a block with outliers can need more than 8 bits, which raises `ToleranceTooTightError`. retrying with its `smallest_tolerance` works, but loosens every block, not just that one:
 
 ```python
+import numpy as np
+from quantize import ToleranceTooTightError, adaptive
+
+weights = np.array(weights)  # a list has no .std(). skip this line for a pytorch tensor
 try:
     q = adaptive.quantize(weights, tolerance=0.1 * weights.std())
 except ToleranceTooTightError as error:
@@ -50,7 +54,7 @@ np.savez("layer.npz", **{name: part for name, part in parts.items() if part is n
 q = Quantized.from_parts(**np.load("layer.npz"))
 ```
 
-to load quantized values from a `torch.save` checkpoint, call `torch.serialization.add_safe_globals([Quantized])` before `torch.load`.
+to keep quantized values in a `torch.save` checkpoint, store `torch.frombuffer(bytearray(q.to_bytes()), dtype=torch.uint8)`, and load each back with `Quantized(t)`. the checkpoint is then as small as the bytes, and `torch.load` reads it without `add_safe_globals`. pickling `q` itself makes the checkpoint about 1.5 times larger, since `torch.save` stores bytes as text, and needs `torch.serialization.add_safe_globals([Quantized])` before `torch.load`.
 
 each value decodes as `code * scale`, or `(code - zero_point) * scale` with zero-points, using the scale and zero-point of its block. codes are signed and `bits` wide, and `q.codes` packs them low bits first. scales can be negative, since a symmetric block puts its value farthest from zero on the most negative code. `help(Quantized)` has the details.
 

{quantize_py-0.3.0 → quantize_py-0.3.1}/python/src/lib.rs

@@ -147,7 +147,11 @@ fn adaptive_quantize(
     quantize_values(py, values, scale, |_| Scheme::Adaptive { block, tolerance })
 }
 
-#[pymodule]
+// `refine` and `alternate` borrow a tensor mutably while they change it, and
+// only the GIL keeps other threads from borrowing it at the same time, which
+// would panic. So free-threaded Python turns the GIL back on when it imports
+// this module.
+#[pymodule(gil_used = true)]
 #[pyo3(name = "_native")]
 fn native(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add("__version__", env!("CARGO_PKG_VERSION"))?;

{quantize_py-0.3.0 → quantize_py-0.3.1}/python/src/quantized/inner.rs

@@ -1,3 +1,5 @@
+use std::sync::Arc;
+
 use half::{bf16, f16};
 use pyo3::prelude::*;
 use quantize::{Quantized, Scale, Scheme, learned};
@@ -6,11 +8,16 @@ use super::parts::Parts;
 use crate::error::from_quantize;
 use crate::scale::PyScale;
 
+/// The tensor, with one of three scale types, behind an `Arc` so that a
+/// clone shares its codes and scales instead of copying them. `refine` and
+/// `alternate` change it through `Arc::make_mut`, which copies a shared
+/// tensor first, so a `copy()`, or a `dot` or `matmul` still running, keeps
+/// the values it had.
 #[derive(Clone, PartialEq)]
 pub(crate) enum QuantizedInner {
-    F32(Quantized<f32>),
-    F16(Quantized<f16>),
-    Bf16(Quantized<bf16>),
+    F32(Arc<Quantized<f32>>),
+    F16(Arc<Quantized<f16>>),
+    Bf16(Arc<Quantized<bf16>>),
 }
 
 macro_rules! with_inner {
@@ -64,9 +71,15 @@ impl QuantizedInner {
         scale: PyScale,
     ) -> PyResult<Self> {
         match scale {
-            PyScale::F32 => quantize_shaped(scheme, values, shape).map(Self::F32),
-            PyScale::F16 => quantize_shaped(scheme, values, shape).map(Self::F16),
-            PyScale::Bf16 => quantize_shaped(scheme, values, shape).map(Self::Bf16),
+            PyScale::F32 => quantize_shaped(scheme, values, shape)
+                .map(Arc::new)
+                .map(Self::F32),
+            PyScale::F16 => quantize_shaped(scheme, values, shape)
+                .map(Arc::new)
+                .map(Self::F16),
+            PyScale::Bf16 => quantize_shaped(scheme, values, shape)
+                .map(Arc::new)
+                .map(Self::Bf16),
         }
         .map_err(from_quantize)
     }
@@ -92,11 +105,15 @@ impl QuantizedInner {
     }
 
     pub(crate) fn refine(&mut self, values: &[f32]) -> quantize::Result<()> {
-        with_inner!(self, |quantized| learned::refine(quantized, values))
+        with_inner!(self, |quantized| {
+            learned::refine(Arc::make_mut(quantized), values)
+        })
     }
 
     pub(crate) fn alternate(&mut self, values: &[f32]) -> quantize::Result<bool> {
-        with_inner!(self, |quantized| learned::alternate(quantized, values))
+        with_inner!(self, |quantized| {
+            learned::alternate(Arc::make_mut(quantized), values)
+        })
     }
 
     pub(crate) fn to_bytes(&self) -> Vec<u8> {
@@ -107,17 +124,17 @@ impl QuantizedInner {
     /// their header names.
     pub(crate) fn from_bytes(bytes: &[u8]) -> quantize::Result<Self> {
         match saved_scale(bytes) {
-            PyScale::F32 => Quantized::from_bytes(bytes).map(Self::F32),
-            PyScale::F16 => Quantized::from_bytes(bytes).map(Self::F16),
-            PyScale::Bf16 => Quantized::from_bytes(bytes).map(Self::Bf16),
+            PyScale::F32 => Quantized::from_bytes(bytes).map(Arc::new).map(Self::F32),
+            PyScale::F16 => Quantized::from_bytes(bytes).map(Arc::new).map(Self::F16),
+            PyScale::Bf16 => Quantized::from_bytes(bytes).map(Arc::new).map(Self::Bf16),
         }
     }
 
     pub(crate) fn from_parts(parts: Parts, scale: PyScale) -> PyResult<Self> {
         match scale {
-            PyScale::F32 => parts.into_quantized().map(Self::F32),
-            PyScale::F16 => parts.into_quantized().map(Self::F16),
-            PyScale::Bf16 => parts.into_quantized().map(Self::Bf16),
+            PyScale::F32 => parts.into_quantized().map(Arc::new).map(Self::F32),
+            PyScale::F16 => parts.into_quantized().map(Arc::new).map(Self::F16),
+            PyScale::Bf16 => parts.into_quantized().map(Arc::new).map(Self::Bf16),
         }
     }
 }

{quantize_py-0.3.0 → quantize_py-0.3.1}/python/src/quantized/methods.rs

@@ -39,9 +39,10 @@ fn bits<S: Scale>(quantized: &Quantized<S>) -> Option<u32> {
     }
 }
 
-// Methods that release the GIL take `slf` instead of `&self`, and clone the
-// tensor out of a short borrow first. A borrow held while the GIL is released
-// would keep `learned.refine` on another thread from borrowing it mutably.
+// Methods that release the GIL take `slf` instead of `&self`, and clone
+// `inner` out of a short borrow first, which shares the tensor instead of
+// copying it. A borrow held while the GIL is released would keep
+// `learned.refine` on another thread from borrowing it mutably.
 #[pymethods]
 impl PyQuantized {
     /// Decode the values into an array of the tensor's `shape`. `out`, if
@@ -72,9 +73,8 @@ impl PyQuantized {
     /// The dot product of the decoded values with `values`, an array of the
     /// tensor's `shape`. For a matrix times a vector, use `matmul`.
     ///
-    /// Symmetric tensors with 4-bit or 8-bit codes are decoded straight from
-    /// the packed codes, without storing the decoded values. Other tensors are
-    /// first unpacked into a buffer as large as the decoded values.
+    /// Each code is read straight from the packed bytes as it's multiplied,
+    /// so the decoded values are never stored, whatever the scheme.
     fn dot(slf: &Bound<'_, Self>, values: Bound<'_, PyAny>) -> PyResult<f32> {
         let (array, values_shape) = as_f32_array(&values)?;
         let values = array.as_slice()?;
@@ -93,6 +93,10 @@ impl PyQuantized {
     /// `inputs` is one vector of shape `(columns,)` or a batch of shape
     /// `(batch, columns)`. The result is `inputs @ W.T`, of shape `(rows,)`
     /// or `(batch, rows)`.
+    ///
+    /// Each call decodes the matrix one row at a time, straight from the
+    /// packed codes, and multiplies each row by every input before moving
+    /// on, so the whole matrix is never decoded at once.
     fn matmul<'py>(
         slf: &Bound<'py, Self>,
         inputs: Bound<'_, PyAny>,

{quantize_py-0.3.0 → quantize_py-0.3.1}/python/src/scale.rs

@@ -2,6 +2,7 @@
 
 use pyo3::exceptions::PyValueError;
 use pyo3::prelude::*;
+use pyo3::types::PyType;
 
 use crate::error::QuantizeError;
 
@@ -9,7 +10,10 @@ use crate::error::QuantizeError;
 /// exactly, in 4 bytes each. `Scale.F16` and `Scale.BF16` round them to 2
 /// bytes: f16 keeps more digits, and bf16 more range. Rounded zero-points cap
 /// the accuracy of asymmetric codes above about 10 bits, so use `Scale.F32`
-/// there. `scale=` also takes the name that `name` returns.
+/// there. `scale=` also takes the name that `name` returns, and
+/// `Scale(name)` gives the scale back, like `Scale("f16")`. Pickles load
+/// through it, so `torch.load` accepts them once
+/// `torch.serialization.add_safe_globals([Scale])` allows the class.
 #[pyclass(
     eq,
     frozen,
@@ -65,6 +69,11 @@ impl FromPyObject<'_, '_> for PyScale {
 
 #[pymethods]
 impl PyScale {
+    #[new]
+    fn new(name: PyScale) -> Self {
+        name
+    }
+
     /// The name that `scale=` also accepts: `'f32'`, `'f16'`, or `'bf16'`.
     #[getter]
     fn name(&self) -> &'static str {
@@ -79,13 +88,15 @@ impl PyScale {
         self.name()
     }
 
+    /// Pickles saved by quantize-py 0.3.0 and earlier call this with the name.
     #[staticmethod]
     fn _from_pickle(name: &str) -> PyResult<Self> {
         Self::from_name(name).ok_or_else(|| PyValueError::new_err("malformed pickle state"))
     }
 
-    fn __reduce__<'py>(slf: &Bound<'py, Self>) -> PyResult<(Bound<'py, PyAny>, (&'static str,))> {
-        let callable = slf.as_any().getattr("_from_pickle")?;
-        Ok((callable, (slf.get().name(),)))
+    // Pickles call the class with the name, so that `torch.load` loads them
+    // once `add_safe_globals([Scale])` allows it, as with `Quantized`.
+    fn __reduce__<'py>(slf: &Bound<'py, Self>) -> (Bound<'py, PyType>, (&'static str,)) {
+        (slf.as_any().get_type(), (slf.get().name(),))
     }
 }

{quantize_py-0.3.0 → quantize_py-0.3.1}/python/src/scheme.rs

@@ -4,6 +4,7 @@ use pyo3::exceptions::PyValueError;
 use pyo3::prelude::*;
 use pyo3::types::PyType;
 
+use crate::error::from_quantize;
 use crate::input::{bits_argument, block_argument};
 use crate::quantize_values;
 use crate::quantized::PyQuantized;
@@ -13,6 +14,17 @@ use crate::scale::PyScale;
 /// `Scheme.symmetric`, `Scheme.asymmetric`, and `Scheme.adaptive` build one,
 /// and `quantize` runs it. `Scheme.Q8_32` and `Scheme.Q4_32` are symmetric
 /// 8-bit and 4-bit codes, with blocks of 32.
+///
+/// `Scheme(text)` reads a scheme written like a call to one of those
+/// methods, without the `Scheme.`, such as
+/// `Scheme("adaptive(block=32, tolerance=0.002)")`, or a constant's name,
+/// such as `Scheme("Q4_32")`. Text that isn't a scheme, or that holds a
+/// value `quantize` would reject, like `bits=99`, raises `QuantizeError`.
+///
+/// Pickles and copies load through `Scheme(text)`, so a scheme that
+/// `quantize` would reject, like `Scheme.symmetric(bits=99)`, raises the
+/// same error when it's unpickled or copied. `torch.load` accepts pickles
+/// once `torch.serialization.add_safe_globals([Scheme])` allows the class.
 #[pyclass(frozen, name = "Scheme", module = "quantize", eq, skip_from_py_object)]
 #[derive(Clone, Copy, PartialEq)]
 pub struct PyScheme {
@@ -35,6 +47,12 @@ impl PyScheme {
 
 #[pymethods]
 impl PyScheme {
+    #[new]
+    fn new(text: &str) -> PyResult<Self> {
+        let inner = text.parse().map_err(from_quantize)?;
+        Ok(Self { inner })
+    }
+
     #[classattr]
     #[pyo3(name = "Q8_32")]
     fn q8_32() -> Self {
@@ -165,6 +183,8 @@ impl PyScheme {
         self.pickle_parts()
     }
 
+    /// Pickles saved by quantize-py 0.3.0 and earlier call this with the
+    /// parts that `__getstate__` returns.
     #[staticmethod]
     fn _from_pickle(
         kind: &str,
@@ -186,8 +206,10 @@ impl PyScheme {
         }
     }
 
-    fn __reduce__<'py>(slf: &Bound<'py, Self>) -> PyResult<(Bound<'py, PyAny>, SchemePickle)> {
-        let callable = slf.getattr("_from_pickle")?;
-        Ok((callable, slf.get().pickle_parts()))
+    // Pickles call the class with the scheme's text, so that `torch.load`
+    // loads them once `add_safe_globals([Scheme])` allows it, as with
+    // `Quantized`.
+    fn __reduce__<'py>(slf: &Bound<'py, Self>) -> (Bound<'py, PyType>, (String,)) {
+        (slf.get_type(), (slf.get().inner.to_string(),))
     }
 }

{quantize_py-0.3.0 → quantize_py-0.3.1}/python/tests/test_learned.py

@@ -127,6 +127,16 @@ def test_except_quantize_error_catches_length():
     assert isinstance(raised.value, ValueError)
 
 
+def test_a_copy_keeps_the_original_through_refine_and_alternate():
+    weights = np.linspace(-0.5, 0.5, 64, dtype=np.float32).reshape(2, 32)
+    for refit in [learned.refine, learned.alternate]:
+        quantized = quantize(weights, bits=4)
+        original = quantized.copy()
+        refit(quantized, weights)
+        assert quantized != original
+        assert original == quantize(weights, bits=4)
+
+
 def test_refine_and_alternate_work_while_another_thread_uses_the_tensor():
     weights = np.random.default_rng(0).standard_normal((256, 256)).astype(np.float32)
     quantized = quantize(weights, bits=4)

{quantize_py-0.3.0 → quantize_py-0.3.1}/python/tests/test_quantize.py

@@ -410,6 +410,19 @@ def test_scheme_factory_does_not_validate():
     assert scheme.bits == 1
     with pytest.raises(InvalidBitsError):
         scheme.quantize([0.1])
+    data = pickle.dumps(scheme)
+    with pytest.raises(InvalidBitsError):
+        pickle.loads(data)
+
+
+def test_scheme_reads_text_written_like_its_class_methods():
+    assert Scheme("symmetric(bits=4, block=32)") == Scheme("Q4_32") == Scheme.Q4_32
+    assert Scheme("asymmetric(bits=3, block=7)") == Scheme.asymmetric(bits=3, block=7)
+    assert Scheme("adaptive(block=32, tolerance=0.002)") == Scheme.adaptive(tolerance=0.002)
+    with pytest.raises(QuantizeError, match=r"isn't a scheme; write one like symmetric\(bits=4"):
+        Scheme("symmetric:4:32")
+    with pytest.raises(InvalidBitsError):
+        Scheme("symmetric(bits=99, block=32)")
 
 
 def test_quantize_rejects_other_dimensions():
@@ -437,9 +450,12 @@ def test_scale_enum_selects_storage():
 def test_scale_can_be_given_by_name():
     for scale in [Scale.F32, Scale.F16, Scale.BF16]:
         assert quantize([0.1], scale=scale.name).scale == scale
+        assert Scale(scale.name) == scale
     assert Scale.BF16.name == "bf16"
     with pytest.raises(QuantizeError, match="scale must be a Scale or its name"):
         quantize([0.1], scale="float32")
+    with pytest.raises(QuantizeError, match="scale must be a Scale or its name"):
+        Scale("float32")
 
 
 def test_bad_values_raise_value_errors_that_name_the_argument():
@@ -521,25 +537,38 @@ def test_pickles_call_the_class_with_the_bytes_that_to_bytes_saves():
     assert Quantized(data) == Quantized.from_bytes(data) == quantized
 
 
-class OnlyQuantizedUnpickler(pickle.Unpickler):
+class OnlyOurClassesUnpickler(pickle.Unpickler):
     """Stands in for `torch.load`, which by default refuses any global it
-    doesn't trust, after `torch.serialization.add_safe_globals([Quantized])`.
+    doesn't trust, after
+    `torch.serialization.add_safe_globals([Quantized, Scale, Scheme])`.
     It also trusts `_codecs.encode`, as `torch.load` does, since pickles below
     protocol 3 store bytes with it."""
 
     def find_class(self, module, name):
-        if (module, name) == ("quantize", "Quantized"):
-            return Quantized
+        classes = {"Quantized": Quantized, "Scale": Scale, "Scheme": Scheme}
+        if module == "quantize" and name in classes:
+            return classes[name]
         if (module, name) == ("_codecs", "encode"):
             return codecs.encode
         raise pickle.UnpicklingError(f"{module}.{name} isn't allowed")
 
 
-def test_pickles_load_when_only_the_class_is_allowed_as_in_torch_load():
-    checkpoint = {"layer": quantize(weight_matrix(8, 32), bits=4, scale=Scale.F16)}
+def test_pickles_load_when_only_the_classes_are_allowed_as_in_torch_load():
+    checkpoint = {
+        "layer": quantize(weight_matrix(8, 32), bits=4, scale=Scale.F16),
+        "scheme": Scheme.adaptive(block=32, tolerance=0.002),
+        "scale": Scale.BF16,
+    }
     for protocol in range(pickle.HIGHEST_PROTOCOL + 1):
         data = pickle.dumps(checkpoint, protocol)
-        assert OnlyQuantizedUnpickler(io.BytesIO(data)).load() == checkpoint
+        assert OnlyOurClassesUnpickler(io.BytesIO(data)).load() == checkpoint
+
+
+def test_scales_and_schemes_pickle_as_a_call_to_their_class():
+    for value in [Scale.F16, Scheme.Q4_32, Scheme.adaptive(block=32, tolerance=0.002)]:
+        rebuild, arguments = value.__reduce__()
+        assert rebuild is type(value)
+        assert rebuild(*arguments) == value
 
 
 def test_from_bytes_rejects_bytes_that_do_not_hold_a_tensor():
@@ -599,6 +628,26 @@ def test_a_pickle_through_from_bytes_still_loads():
     assert pickle.loads(PICKLED_THROUGH_FROM_BYTES) == quantize(weights, bits=8, block=4)
 
 
+# [Scale.F16, Scheme.adaptive(block=32, tolerance=0.002), Scheme.Q4_32], pickled
+# by quantize-py 0.3.0, which called their _from_pickle methods.
+SCALE_AND_SCHEMES_PICKLED_BY_0_3_0 = (
+    b"\x80\x04\x95\xad\x00\x00\x00\x00\x00\x00\x00]\x94(\x8c\x08builtins\x94\x8c"
+    b"\x07getattr\x94\x93\x94\x8c\x08quantize\x94\x8c\x05Scale\x94\x93\x94\x8c\x0c"
+    b"_from_pickle\x94\x86\x94R\x94\x8c\x03f16\x94\x85\x94R\x94h\x03\x8c\x08quanti"
+    b"ze\x94\x8c\x06Scheme\x94\x93\x94\x8c\x0c_from_pickle\x94\x86\x94R\x94(\x8c"
+    b"\x08adaptive\x94NK G?`bM\xe0\x00\x00\x00t\x94R\x94h\x12(\x8c\tsymmetric\x94K"
+    b"\x04K Nt\x94R\x94e."
+)
+
+
+def test_scales_and_schemes_pickled_by_0_3_0_still_load():
+    assert pickle.loads(SCALE_AND_SCHEMES_PICKLED_BY_0_3_0) == [
+        Scale.F16,
+        Scheme.adaptive(block=32, tolerance=0.002),
+        Scheme.Q4_32,
+    ]
+
+
 def test_quantized_compares_by_value():
     weights = weight_matrix(4, 32)
     quantized = quantize(weights, bits=4)