haliax 1.4.dev332__py3-none-any.whl → 1.4.dev334__py3-none-any.whl

haliax/__about__.py

@@ -1 +1 @@
-__version__ = "1.4.dev332"
+__version__ = "1.4.dev334"

haliax/_src/state_dict.py

@@ -395,7 +395,7 @@ def flatten_linear_layers(tree: T) -> T:
         new_In: Axis = flatten_axes(layer.In, "__IN__")
 
         if weight.array is not None:
-            out_first = layer.out_first
+            out_first = layer._out_first
             weight = weight.flatten_axes(layer.Out, new_Out).flatten_axes(layer.In, new_In)
 
             if out_first:

haliax/axis.py

@@ -35,6 +35,8 @@ def make_axes(**kwargs: int) -> Tuple[Axis, ...]:
     Example:
     ```
     X, Y = axes(X=10, Y=20)
+    ```
+
     """
     return tuple(Axis(name, size) for name, size in kwargs.items())
 

haliax/nn/embedding.py

@@ -24,26 +24,60 @@ class Embedding(eqx.Module):
 
     @staticmethod
     def init(Vocab: Axis, Embed: AxisSpec, *, init_scale: float = 1, key, initializer_range: Optional[float] = None):
+        """
+        Initialize an Embedding module.
+
+        An embedding module is a simple lookup table that maps integer indices to vectors or tensors.
+        Weights are initialized with a truncated normal distribution with a standard deviation of
+          `init_scale / output_size`.
+
+        Args:
+            Vocab: Size of the vocabulary
+            Embed: Shape of the embedding vectors. May be a single axis or a full AxisSpec
+            init_scale: Scale of the initialization
+            key: PRNG key
+            initializer_range: Deprecated. Use init_scale instead.
+        """
         if initializer_range is not None:
             warnings.warn("initializer_range is deprecated. Use init_std instead.", DeprecationWarning)
             init_scale = initializer_range
 
         all_axes = (Vocab,) + ensure_tuple(Embed)
         output_size = hax.axis_size(Embed)
-        weight = hax.random.truncated_normal(key, all_axes, -3, 3) * (init_scale / math.sqrt(output_size))
+        weight = hax.random.truncated_normal(key, all_axes, -3, 3) * (init_scale / output_size)
         return Embedding(weight=weight, Vocab=Vocab, Embed=Embed)
 
-    def __call__(self, input_ids, *, key: Optional[PRNGKeyArray] = None):
+    def __call__(self, input_ids: NamedArray, *, key: Optional[PRNGKeyArray] = None):
+        """Alias for `embed`. key is ignored."""
         return self.embed(input_ids)
 
     @named_call
-    def embed(self, input_ids):
+    def embed(self, input_ids: NamedArray):
+        """
+        Args:
+            input_ids: token IDs with shape > {Vocab}
+        """
         input_embeds = self.weight.take(self.Vocab, input_ids)
         return input_embeds
 
-    def unembed(self, input_embeds):
+    def unembed(self, input_embeds: NamedArray):
+        """
+        Unembed the input embeddings back to the vocabulary space.
+
+        Equivalent to `input_embeds.dot(self.weight, axis=self.Embed)`.
+        """
         return input_embeds.dot(self.weight, axis=self.Embed)
 
     def resize_embeddings(self, new_size: int, key: Optional[PRNGKeyArray] = None):
+        """
+        Resize the embedding layer to a new size.
+        Args:
+            new_size: New size of the vocabulary
+            key: PRNG key for initialization of any new weights
+
+        Returns:
+            Embedding: Resized embedding layer
+
+        """
         new_weights = resize_axis(self.weight, self.Vocab, new_size, key=key)
         return dataclasses.replace(self, Vocab=self.Vocab.resize(new_size), weight=new_weights)  # type: ignore

haliax/nn/linear.py

@@ -72,7 +72,10 @@ class Linear(eqx.Module):
         return q
 
     @property
-    def out_first(self):
+    def _out_first(self):
+        """
+        Returns: bool: Whether the output axes are first in the weight matrix
+        """
         # We do it this way because of scan layers
         if isinstance(self.Out, hax.Axis):
             return self.weight.axes[-1] != self.Out

haliax/nn/normalization.py

@@ -58,6 +58,9 @@ def logsumexp(a: A, axis: Optional[AxisSelection] = None) -> A:
     return wrap_reduction_call(jnn.logsumexp, a, axis=axis, single_axis_only=False, supports_where=False)
 
 
+# TODO: support where in softmax, etc
+
+
 def softmax(a: A, axis: Optional[AxisSelection] = None) -> A:
     return wrap_axiswise_call(jnn.softmax, a, axis=axis, single_axis_only=False)
 

{haliax-1.4.dev332.dist-info → haliax-1.4.dev334.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: haliax
-Version: 1.4.dev332
+Version: 1.4.dev334
 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.dev332.dist-info → haliax-1.4.dev334.dist-info}/RECORD

@@ -1,6 +1,6 @@
-haliax/__about__.py,sha256=1Bc7sSYtp49Iv6Rc1ZOV2dLFJ-vdQxZRlu4f9wKlkYg,27
+haliax/__about__.py,sha256=QmiPi6dv5hegVnrb07SVyu_ESGszJnp91zXeTEgMyZA,27
 haliax/__init__.py,sha256=CQHrLfNXSO8hd27zqqzDeJP6Vzls_BKAE65VA48wlAc,29075
-haliax/axis.py,sha256=U8Ugd2d-u55RoQQqRiITKm1mzsZkDh3rThfAgkTWOwo,20757
+haliax/axis.py,sha256=tEfVlIqFzJRYfMVtHNw_iM68Otehzlb9cGRLFcrXr_s,20766
 haliax/core.py,sha256=d-9nppDP_-IcxF1yHXDzzVNbP7WIadA1ojb3xHP_-SQ,71000
 haliax/debug.py,sha256=0qEGgIsEw3Jkp40oxgBAjD0Ps-cxi3QNP42uMRGlM84,3612
 haliax/hof.py,sha256=jDcii2IwAfhoNYTCYsdeS8vpSKmtzioPKZ6LBAVuWn8,18205
@@ -22,21 +22,21 @@ haliax/_src/einsum.py,sha256=IXBynFdGYFYP3kbNjQfQLovqwjhNeDAoDnUiEqATV0I,15410
 haliax/_src/fp8.py,sha256=wtyOK7TrXgg4UlwydDklCMXvkbGMoixB3mKOk3IIJvM,4951
 haliax/_src/parsing.py,sha256=9JDE3unCjS8AN_B5kiAy1yV0AemYW0jL8A_T_bdcLHs,11189
 haliax/_src/rearrange.py,sha256=SWZAIMSrqQzqwZDxrMk25q6IOaK4PyDUEm-VOSe0N18,19498
-haliax/_src/state_dict.py,sha256=YcOtA_7B99_0VYBI1-WLHb4ZcxIfxQ1f1UyollTXlPI,17276
+haliax/_src/state_dict.py,sha256=hKE0EVPN4zXvyBM40FyUSrmvd88SUFqCJ7ZLb4ZqNy4,17277
 haliax/_src/util.py,sha256=pXizGeJfRcx7QA1YVAxvru2EnMGIXAhZdcPGkwatP_M,1566
 haliax/nn/__init__.py,sha256=WeiO9JGBK-Am0NZ-s2UW7OmR4cFZbQAakuTDcMMnLg4,2875
 haliax/nn/activations.py,sha256=9h1uJ2pUWmb_OsRhqUGUJ6s95WZgxABDx7YWFD2tn5A,1723
 haliax/nn/attention.py,sha256=5P6IvQ5G-hpStKFnZ4mEV379Af6pkCyIA0qWXMM3Ed4,11942
 haliax/nn/conv.py,sha256=ozdvYBnb20CDz_xMKiXdpOW0ZMBZKlwadx21XafFw3E,15131
 haliax/nn/dropout.py,sha256=2BETs0MeHMdzUuGof3vdVRZSwC7vQwKP9BT-Ix0lDRY,3682
-haliax/nn/embedding.py,sha256=RfNNO2GZim7o0pxqzhpwv9JcPLAArEXKvKwi8mt0tQo,1687
-haliax/nn/linear.py,sha256=DQH4YdoQ4nOPpXihSZfPuoidIMQheTCp3aGDKksc7Ww,2722
+haliax/nn/embedding.py,sha256=e-U3U7-nI0puNPLr-e-LH_-GfplgIuTHHGUzqMAsr-U,2880
+haliax/nn/linear.py,sha256=21sorQpSjg4zPMG3-SHg7mqyAg86E5QKI3bxdSuw-Ag,2825
 haliax/nn/loss.py,sha256=OsXpiidiKOD6gslXC54hywn9qm3shH9qyhx8baaps0Y,4565
 haliax/nn/mlp.py,sha256=KW_7phxbt0C8CLEVahP5Zoq4SQruiGyg7ltRvLrTdPg,3695
-haliax/nn/normalization.py,sha256=4e1rpVvZKW9w5eUsn8j0tBF6lAIs5YljsvOE_HaoOD8,2652
+haliax/nn/normalization.py,sha256=jjPVrs2wBlaK47WwaAm18Y53EnREvy6HLpAgtdqCajI,2692
 haliax/nn/pool.py,sha256=DHeswGramsqNvYufl36NbwgWgI-Iht9bnJy170Pwfu0,8530
 haliax/nn/scan.py,sha256=cm5GHVjdv8nWKKlo9hHzlyoj0Ubq9770EgqsdbUc7Z0,29163
-haliax-1.4.dev332.dist-info/METADATA,sha256=Qcs512OIq18GGpSwTgAhbstqbxdVfhfm5RjuueAYPoA,7663
-haliax-1.4.dev332.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-haliax-1.4.dev332.dist-info/licenses/LICENSE,sha256=bJiay7Nn5SHQ2n_4ZIT3AE0W1RGq4O7pxOApgBsaT64,11349
-haliax-1.4.dev332.dist-info/RECORD,,
+haliax-1.4.dev334.dist-info/METADATA,sha256=u8d5xnDqZtYJq6Bx1caefXQ45_e8pXpPydwjCAy0ykg,7663
+haliax-1.4.dev334.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+haliax-1.4.dev334.dist-info/licenses/LICENSE,sha256=bJiay7Nn5SHQ2n_4ZIT3AE0W1RGq4O7pxOApgBsaT64,11349
+haliax-1.4.dev334.dist-info/RECORD,,