lalamo 0.2.5__tar.gz → 0.2.7__tar.gz

{lalamo-0.2.5 → lalamo-0.2.7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: lalamo
-Version: 0.2.5
+Version: 0.2.7
 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

{lalamo-0.2.5 → lalamo-0.2.7}/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.5"
+__version__ = "0.2.7"
 
 __all__ = [
     "REPO_TO_MODEL",

{lalamo-0.2.5 → lalamo-0.2.7}/lalamo/model_import/model_specs/common.py

@@ -8,7 +8,6 @@ from safetensors.flax import load_file as load_safetensors
 
 from lalamo.model_import.configs import ForeignConfig
 from lalamo.quantization import QuantizationMode
-from lalamo.utils import torch_to_jax
 
 __all__ = [
     "HUGGINFACE_GENERATION_CONFIG_FILE",
@@ -38,6 +37,8 @@ class WeightsType(Enum):
 
         import torch
 
+        from lalamo.modules.torch_interop import torch_to_jax
+
         torch_weights = torch.load(filename, map_location="cpu", weights_only=True)
         return {k: cast_if_float(torch_to_jax(v), float_dtype) for k, v in torch_weights.items()}
 
@@ -95,7 +96,7 @@ def awq_model_spec(
     )
 
 
-def build_quantized_models(model_specs: list[ModelSpec]):
+def build_quantized_models(model_specs: list[ModelSpec]) -> list[ModelSpec]:
     quantization_compatible_repos: list[str] = [
         "Qwen/Qwen2.5-3B-Instruct",
         "Qwen/Qwen2.5-7B-Instruct",

lalamo-0.2.7/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.7/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.5 → lalamo-0.2.7}/lalamo.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: lalamo
-Version: 0.2.5
+Version: 0.2.7
 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

{lalamo-0.2.5 → lalamo-0.2.7}/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.5/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)