haliax 1.4.dev450__tar.gz → 1.4.dev452__tar.gz

{haliax-1.4.dev450 → haliax-1.4.dev452}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: haliax
-Version: 1.4.dev450
+Version: 1.4.dev452
 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.dev450 → haliax-1.4.dev452}/src/haliax/__about__.py

@@ -3,4 +3,4 @@
 # SPDX-License-Identifier: Apache-2.0
 
 
-__version__ = "1.4.dev450"
+__version__ = "1.4.dev452"

{haliax-1.4.dev450 → haliax-1.4.dev452}/src/haliax/nn/linear.py

@@ -159,12 +159,14 @@ class Linear(ModuleWithStateDictSerialization, ReparamEnabled):
 
     def to_state_dict(self, prefix: Optional[str] = None) -> StateDict:
         # weight can be None for certain filtering things like LoRA
-        scaled = dataclasses.replace(self, weight=self.weight * self.reparam.active_scale if self.weight is not None else None)
+        scaled = dataclasses.replace(
+            self, weight=self.weight * self.reparam.active_scale if self.weight is not None else None
+        )
         return default_eqx_module_to_state_dict(scaled, prefix)
 
     def from_state_dict(self: Mod, state_dict: StateDict, prefix: Optional[str] = None) -> Mod:
         unscaled = default_eqx_module_from_state_dict(self, state_dict, prefix)
-        if unscaled.weight is not None: 
+        if unscaled.weight is not None:
             unscaled = dataclasses.replace(unscaled, weight=unscaled.weight / self.reparam.active_scale)
         return unscaled