lalamo 0.2.4__tar.gz → 0.2.6__tar.gz

{lalamo-0.2.4 → lalamo-0.2.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: lalamo
-Version: 0.2.4
+Version: 0.2.6
 Summary: JAX library for optimization and export of models for use with the UZU inference engine.
 Requires-Python: <4,>=3.12
 Description-Content-Type: text/markdown
@@ -18,6 +18,7 @@ Requires-Dist: optax>=0.2.4
 Requires-Dist: rich>=14.0.0
 Requires-Dist: thefuzz>=0.22.1
 Requires-Dist: typer>=0.15.1
+Requires-Dist: safetensors>=0.6.2
 Dynamic: license-file
 
 <p align="center">

{lalamo-0.2.4 → lalamo-0.2.6}/lalamo/__init__.py

@@ -1,7 +1,7 @@
 from lalamo.model_import import REPO_TO_MODEL, ModelSpec, import_model
 from lalamo.modules import Decoder
 
-__version__ = "0.2.4"
+__version__ = "0.2.6"
 
 __all__ = [
     "REPO_TO_MODEL",

{lalamo-0.2.4 → lalamo-0.2.6}/lalamo/model_import/model_specs/common.py

@@ -7,8 +7,8 @@ from jaxtyping import Array, DTypeLike
 from safetensors.flax import load_file as load_safetensors
 
 from lalamo.model_import.configs import ForeignConfig
+from lalamo.modules.torch_interop import torch_to_jax
 from lalamo.quantization import QuantizationMode
-from lalamo.utils import torch_to_jax
 
 __all__ = [
     "HUGGINFACE_GENERATION_CONFIG_FILE",

lalamo-0.2.6/lalamo/modules/torch_interop.py

@@ -0,0 +1,29 @@
+import jax.numpy as jnp
+import torch
+from jaxtyping import Array
+
+__all__ = ["jax_to_torch", "torch_to_jax"]
+
+
+@torch.no_grad()
+def _torch_to_jax_bfloat16(tensor: torch.Tensor) -> Array:
+    if tensor.dtype != torch.bfloat16:
+        raise ValueError("Trying to convert non-bfloat16 tensor to bfloat16")
+    intermediate_tensor = tensor.view(torch.uint16)
+    return jnp.array(intermediate_tensor).view("bfloat16")
+
+
+def torch_to_jax(array: torch.Tensor) -> Array:
+    array = array.detach().cpu()
+    if array.dtype == torch.bfloat16:
+        return _torch_to_jax_bfloat16(array)
+    return jnp.array(array.numpy())
+
+
+def jax_to_torch(array: Array) -> torch.Tensor:
+    from torch.utils import dlpack as _dlpack
+
+    if array.dtype == jnp.bfloat16:
+        intermediate_array = array.view(jnp.uint16)
+        return _dlpack.from_dlpack(intermediate_array).view(torch.bfloat16)
+    return _dlpack.from_dlpack(array)

lalamo-0.2.6/lalamo/utils.py

@@ -0,0 +1,27 @@
+import einops
+import jax.numpy as jnp
+from jaxtyping import Array
+
+__all__ = ["jax_uint4_to_packed_uint8"]
+
+
+def jax_uint4_to_packed_uint8(array: Array) -> Array:
+    if array.dtype != jnp.uint4:
+        raise ValueError(f"Input array must have dtype jnp.uint4, but got {array.dtype}")
+
+    if not array.shape:
+        raise ValueError("Input array cannot be a scalar and must have at least one dimension.")
+
+    *_, last_dim = array.shape
+    if last_dim % 2 != 0:
+        raise ValueError(f"The last dimension of the input array must be even, but got shape {array.shape}")
+
+    low_nibbles, high_nibbles = einops.rearrange(
+        array.astype(jnp.uint8),
+        "... (dim_half two) -> two ... dim_half",
+        two=2,
+    )
+
+    packed = (high_nibbles << 4) | low_nibbles
+
+    return packed.astype(jnp.uint8)

{lalamo-0.2.4 → lalamo-0.2.6}/lalamo.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: lalamo
-Version: 0.2.4
+Version: 0.2.6
 Summary: JAX library for optimization and export of models for use with the UZU inference engine.
 Requires-Python: <4,>=3.12
 Description-Content-Type: text/markdown
@@ -18,6 +18,7 @@ Requires-Dist: optax>=0.2.4
 Requires-Dist: rich>=14.0.0
 Requires-Dist: thefuzz>=0.22.1
 Requires-Dist: typer>=0.15.1
+Requires-Dist: safetensors>=0.6.2
 Dynamic: license-file
 
 <p align="center">

{lalamo-0.2.4 → lalamo-0.2.6}/lalamo.egg-info/SOURCES.txt

@@ -53,6 +53,7 @@ lalamo/modules/linear.py
 lalamo/modules/mlp.py
 lalamo/modules/normalization.py
 lalamo/modules/rope.py
+lalamo/modules/torch_interop.py
 lalamo/modules/utils.py
 tests/test_generation.py
 tests/test_huggingface_models.py

{lalamo-0.2.4 → lalamo-0.2.6}/lalamo.egg-info/requires.txt

@@ -9,6 +9,7 @@ optax>=0.2.4
 rich>=14.0.0
 thefuzz>=0.22.1
 typer>=0.15.1
+safetensors>=0.6.2
 
 [:sys_platform == "darwin"]
 jax>=0.4.38

{lalamo-0.2.4 → lalamo-0.2.6}/pyproject.toml

@@ -16,6 +16,7 @@ dependencies = [
     "rich>=14.0.0",
     "thefuzz>=0.22.1",
     "typer>=0.15.1",
+    "safetensors>=0.6.2",
     # "jax-metal>=0.1.1 ; sys_platform == 'darwin'",
 ]
 readme = { file = "README.md", content-type = "text/markdown" }

lalamo-0.2.4/lalamo/utils.py

@@ -1,55 +0,0 @@
-import einops
-import jax.numpy as jnp
-import torch.utils.dlpack
-from jaxtyping import Array
-
-__all__ = [
-    "jax_to_torch",
-    "jax_uint4_to_packed_uint8",
-    "torch_to_jax",
-]
-
-
-@torch.no_grad()
-def _torch_to_jax_bfloat16(tensor: torch.Tensor) -> Array:
-    # Credit: https://github.com/jax-ml/ml_dtypes/issues/81#issuecomment-2399636232
-    if tensor.dtype != torch.bfloat16:
-        raise ValueError("Trying to convert non-bfloat16 tensor to bfloat16")
-    intermediate_tensor = tensor.view(torch.uint16)
-    return jnp.array(intermediate_tensor).view("bfloat16")
-
-
-def torch_to_jax(array: torch.Tensor) -> Array:
-    array = array.detach().cpu()
-    if array.dtype == torch.bfloat16:
-        return _torch_to_jax_bfloat16(array)
-    return jnp.array(array.numpy())
-
-
-def jax_to_torch(array: Array) -> torch.Tensor:
-    if array.dtype == jnp.bfloat16:
-        intermediate_array = array.view(jnp.uint16)
-        return torch.utils.dlpack.from_dlpack(intermediate_array).view(torch.bfloat16)
-    return torch.utils.dlpack.from_dlpack(array)
-
-
-def jax_uint4_to_packed_uint8(array: Array) -> Array:
-    if array.dtype != jnp.uint4:
-        raise ValueError(f"Input array must have dtype jnp.uint4, but got {array.dtype}")
-
-    if not array.shape:
-        raise ValueError("Input array cannot be a scalar and must have at least one dimension.")
-
-    *_, last_dim = array.shape
-    if last_dim % 2 != 0:
-        raise ValueError(f"The last dimension of the input array must be even, but got shape {array.shape}")
-
-    low_nibbles, high_nibbles = einops.rearrange(
-        array.astype(jnp.uint8),
-        "... (dim_half two) -> two ... dim_half",
-        two=2,
-    )
-
-    packed = (high_nibbles << 4) | low_nibbles
-
-    return packed.astype(jnp.uint8)