haliax 1.4.dev441__tar.gz → 1.4.dev443__tar.gz

{haliax-1.4.dev441 → haliax-1.4.dev443}/PKG-INFO

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

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

{haliax-1.4.dev441 → haliax-1.4.dev443}/src/haliax/nn/embedding.py

@@ -24,7 +24,12 @@ class Embedding(eqx.Module, ReparamEnabled):
     # axes
     Vocab: Axis = eqx.field(static=True)
     Embed: AxisSpec = eqx.field(static=True)
-    reparam: AbstractEmbeddingReparam = eqx.field(static=True)
+
+    _reparam_cls: type[AbstractEmbeddingReparam] = eqx.field(static=True, default=EmbeddingStandardParam)
+
+    @property
+    def reparam(self) -> AbstractEmbeddingReparam:
+        return self._reparam_cls(self.Embed, self.Vocab)
 
     @staticmethod
     def init(
@@ -61,7 +66,7 @@ class Embedding(eqx.Module, ReparamEnabled):
         weight = hax.random.truncated_normal(key, all_axes, -3, 3) * (
             init_scale * reparam_cls.init_scale(Vocab, Embed)
         )
-        return Embedding(weight=weight, Vocab=Vocab, Embed=Embed, reparam=reparam_cls(Embed, Vocab))
+        return Embedding(weight=weight, Vocab=Vocab, Embed=Embed, _reparam_cls=reparam_cls)
 
     def __call__(self, input_ids: NamedArray, *, key: PRNGKeyArray | None = None):
         """Alias for `embed`. key is ignored."""

{haliax-1.4.dev441 → haliax-1.4.dev443}/src/haliax/nn/linear.py

@@ -45,9 +45,14 @@ class Linear(ModuleWithStateDictSerialization, ReparamEnabled):
 
     In: AxisSpec = eqx.field(static=True)
     Out: AxisSpec = eqx.field(static=True)
-    reparam: AbstractLinearReparam = eqx.field(static=True)
     dot_general: DotGeneralOp = eqx.field(default_factory=DotGeneralOp.default)
 
+    _reparam_cls: type[AbstractLinearReparam] = eqx.field(static=True, default=LinearStandardParam)
+
+    @property
+    def reparam(self) -> AbstractLinearReparam:
+        return self._reparam_cls(self.In, self.Out)
+
     @staticmethod
     def init(
         In: AxisSpec,
@@ -77,7 +82,7 @@ class Linear(ModuleWithStateDictSerialization, ReparamEnabled):
         if dot_general is None:
             dot_general = DotGeneralOp.default()
 
-        return Linear(weight, bias, In, Out, dot_general=dot_general, reparam=reparam_cls(In, Out))
+        return Linear(weight, bias, In, Out, dot_general=dot_general, _reparam_cls=reparam_cls)
 
     @named_call
     def __call__(self, inputs, *, key: PRNGKeyArray | None = None):

{haliax-1.4.dev441 → haliax-1.4.dev443}/tests/test_mup_embedding.py

@@ -31,7 +31,7 @@ def test_mup_embedding_unembedding_scale():
     Embed = (hax.Axis("E", 3),)
 
     weight = hax.ones(hax.concat_axis_specs(Vocab, Embed))
-    layer = Embedding(weight=weight, Vocab=Vocab, Embed=Embed, reparam=EmbeddingMup(Embed, Vocab))
+    layer = Embedding(weight=weight, Vocab=Vocab, Embed=Embed, _reparam_cls=EmbeddingMup)
 
     scale = layer.reparam.unembed_active_scale
     assert scale == pytest.approx(1.0 / hax.axis_size(Embed))

{haliax-1.4.dev441 → haliax-1.4.dev443}/tests/test_mup_linear.py

@@ -11,7 +11,12 @@ import pytest
 
 import haliax as hax
 from haliax.nn import Linear
-from haliax.nn.mup import InputLinearMup, LinearStandardParam, HiddenLinearMup, OutputLinearMup
+from haliax.nn.mup import (
+    InputLinearMup,
+    LinearStandardParam,
+    HiddenLinearMup,
+    OutputLinearMup,
+)
 
 
 @pytest.mark.parametrize("out_first", [True, False])
@@ -37,8 +42,8 @@ def test_mup_linear_call_matches_linear():
     weight = hax.ones(hax.concat_axis_specs(Out, In)) * 0.5
     bias = hax.full(Out, 0.25)
 
-    linear = Linear(weight, bias, In, Out, reparam=LinearStandardParam(In, Out))
-    mup = Linear(weight, bias, In, Out, reparam=InputLinearMup(In, Out))
+    linear = Linear(weight, bias, In, Out, _reparam_cls=LinearStandardParam)
+    mup = Linear(weight, bias, In, Out, _reparam_cls=InputLinearMup)
 
     inputs = hax.full(hax.concat_axis_specs(Batch, In), 2.0)
 
@@ -109,8 +114,8 @@ def test_input_linear_behaves_like_base_linear():
     weight = hax.ones((Out, In)) * 0.1
     bias = hax.zeros(Out)
 
-    linear = Linear(weight, bias, In, Out, reparam=LinearStandardParam(In, Out))
-    input_linear = Linear(weight, bias, In, Out, reparam=InputLinearMup(In, Out))
+    linear = Linear(weight, bias, In, Out, _reparam_cls=LinearStandardParam)
+    input_linear = Linear(weight, bias, In, Out, _reparam_cls=InputLinearMup)
 
     inputs = hax.random.normal(jrandom.PRNGKey(5), (Batch, In))