haliax 1.4.dev336__tar.gz → 1.4.dev337__tar.gz

{haliax-1.4.dev336 → haliax-1.4.dev337}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: haliax
-Version: 1.4.dev336
+Version: 1.4.dev337
 Summary: Named Tensors for Legible Deep Learning in JAX
 Project-URL: Homepage, https://github.com/stanford-crfm/haliax
 Project-URL: Bug Tracker, https://github.com/stanford-crfm/haliax/issues/

haliax-1.4.dev337/src/haliax/__about__.py

@@ -0,0 +1 @@
+__version__ = "1.4.dev337"

{haliax-1.4.dev336 → haliax-1.4.dev337}/src/haliax/quantization.py

@@ -14,9 +14,10 @@ import jax.random as jrandom
 from aqt.jax.v2.aqt_dot_general import DotGeneral
 from jax import numpy as jnp
 from jax.tree_util import DictKey, FlattenedIndexKey, GetAttrKey, SequenceKey
-from jax.typing import DTypeLike
+from jaxtyping import DTypeLike, PyTree
 
 import haliax.nn as hnn
+from haliax.state_dict import StateDict
 from haliax.types import PrecisionLike
 
 from ._src.fp8 import dot_general_with_precision, in_qdq, out_qdq
@@ -206,6 +207,10 @@ class Int8DotGeneralOp(OverwriteWithGradient):
         cfg = aqt_config.set_context(self.cfg, jrandom.PRNGKey(42), train_step=None)
         return cfg(lhs, rhs, dimension_numbers, precision, preferred_element_type)
 
+    def to_state_dict(tree: PyTree, prefix: Optional[str] = None) -> StateDict:
+        warnings.warn("Ignore all int8 states (if any) for now.")
+        return {}
+
 
 @dataclass(frozen=True)
 class QuantizationConfig:

haliax-1.4.dev336/src/haliax/__about__.py

@@ -1 +0,0 @@
-__version__ = "1.4.dev336"