lalamo 0.2.5__py3-none-any.whl → 0.2.6__py3-none-any.whl

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.6"
 
 __all__ = [
     "REPO_TO_MODEL",

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/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/utils.py

@@ -1,36 +1,8 @@
 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)
+__all__ = ["jax_uint4_to_packed_uint8"]
 
 
 def jax_uint4_to_packed_uint8(array: Array) -> Array:

{lalamo-0.2.5.dist-info → lalamo-0.2.6.dist-info}/METADATA

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

{lalamo-0.2.5.dist-info → lalamo-0.2.6.dist-info}/RECORD

@@ -1,9 +1,9 @@
-lalamo/__init__.py,sha256=HNdYYIO9ypGyZ53AUbiDIp-i3W61ZupGATapp3rEFYQ,217
+lalamo/__init__.py,sha256=VhtKkFHsGy_199K6ToPyIuepm2yVsi7knqRNCG5PiXk,217
 lalamo/common.py,sha256=uYLw68V4AF3zlENG3KAIKRpOFXVHv8xX_n0cc3qJnj4,1877
 lalamo/language_model.py,sha256=GiA_BDQuYCgVBFHljb_ltW_M7g3I1Siwm111M3Jc8MM,9286
 lalamo/main.py,sha256=K2RLyTcxvBCP0teSsminssj_oUkuQAQ5y9ixa1uOqas,9546
 lalamo/quantization.py,sha256=8o6ryIZLzzDYQuvBTboPfaVVdfijAKGpTxOcg3GKVD8,2752
-lalamo/utils.py,sha256=QzkT0_82nd9pS5p0e7yOOdL_ZeKQr_Ftj4kFrWF35R8,1754
+lalamo/utils.py,sha256=ihV9ojDMlAf2_Ja5kNZMIYLMQxpQXBlNOd9TIdMq0yM,815
 lalamo/model_import/__init__.py,sha256=Z8pS9rbKKx1QgUy7KZtHxiNWlZhII3mdovT9d37vAxg,168
 lalamo/model_import/common.py,sha256=sHXEGQUtVb6TRT5FOGtJG9pz1Ohy5v_LtunubVxZKqQ,3303
 lalamo/model_import/configs/__init__.py,sha256=JYXeco_kfzKZuWqEmG24qxeYWs-FuE1W1kNgoFNrBEw,461
@@ -22,7 +22,7 @@ lalamo/model_import/loaders/common.py,sha256=2FigeDMUwlMPUebX8DAK2Yh9aLgVtsfTj0S
 lalamo/model_import/loaders/executorch.py,sha256=nSvpylK8QL3nBk78P3FabLoyA87E3kv5CCpMfvuZe6Q,8886
 lalamo/model_import/loaders/huggingface.py,sha256=Ze_qB0fSxY8lH4ovH0t8jd5jiteasUWkS9HdgMZXCrs,10523
 lalamo/model_import/model_specs/__init__.py,sha256=_sJthAH1xXl5B9JPhRqMVP2t5KkhzqmKFHSRlOiFg8s,915
-lalamo/model_import/model_specs/common.py,sha256=Ob3yTMDczKUHMWBH0PaClbSvHJhKfZ-zbv2Z04YqMVg,3806
+lalamo/model_import/model_specs/common.py,sha256=H56uwiyDsm6GD-1uxCm-zGBLuxE_EN84ivHi80niL0g,3822
 lalamo/model_import/model_specs/deepseek.py,sha256=9l3pVyC-ZoIaFG4xWhPDCbKkD2TsND286o0KzO0uxKo,788
 lalamo/model_import/model_specs/gemma.py,sha256=y4aDeaGGl4JPIanAgPMOlyfD_cx3Q7rpTKgDgx5AsX0,2299
 lalamo/model_import/model_specs/huggingface.py,sha256=ktDJ_qZxSGmHREydrYQaWi71bXJZiHqzHDoZeORENno,784
@@ -44,10 +44,11 @@ lalamo/modules/linear.py,sha256=loUGFu3wx-iGqDqGMphQorhqBm7b9lAqT4B0jAmoamk,2408
 lalamo/modules/mlp.py,sha256=bV8qJTjsQFGv-CA7d32UQFn6BX5zmCKWC5pgm29-W3U,2631
 lalamo/modules/normalization.py,sha256=BWCHv6ycFJ_qMGfxkusGfay9dWzUlbpuwmjbLy2rI68,2380
 lalamo/modules/rope.py,sha256=Vdt2J_W0MPDK52nHsroLVCfWMHyHW3AfrKZCZAE4VYs,9369
+lalamo/modules/torch_interop.py,sha256=-mujd1zI4ec2w92Hd50RtDa0K3jl6ZSnPxc5r3Fp9nU,916
 lalamo/modules/utils.py,sha256=5QTdi34kEI5jix7TfTdB0mOYZbzZUul_T1y8eWCA6lQ,262
-lalamo-0.2.5.dist-info/licenses/LICENSE,sha256=diHRfjSEJHD1nnEeMIfMRCjR3UERf8bT3eseD6b1ayA,1072
-lalamo-0.2.5.dist-info/METADATA,sha256=Uh-z7iYbNur26j9mBoH5OInfmBqXQJ2pRdSKRiru1xg,2645
-lalamo-0.2.5.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-lalamo-0.2.5.dist-info/entry_points.txt,sha256=qli7qTfnBk5WP10rOGXXEckHMtt-atJMDWd8jN89Uks,43
-lalamo-0.2.5.dist-info/top_level.txt,sha256=VHvWL5JN5XRG36NsN_MieJ7EwRihEOrEjyDaTdFJ-aI,7
-lalamo-0.2.5.dist-info/RECORD,,
+lalamo-0.2.6.dist-info/licenses/LICENSE,sha256=diHRfjSEJHD1nnEeMIfMRCjR3UERf8bT3eseD6b1ayA,1072
+lalamo-0.2.6.dist-info/METADATA,sha256=3IJMF8cNEPM2dmIZfDcNlpOT-SQkvJBkZnWPxjA3CWY,2645
+lalamo-0.2.6.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+lalamo-0.2.6.dist-info/entry_points.txt,sha256=qli7qTfnBk5WP10rOGXXEckHMtt-atJMDWd8jN89Uks,43
+lalamo-0.2.6.dist-info/top_level.txt,sha256=VHvWL5JN5XRG36NsN_MieJ7EwRihEOrEjyDaTdFJ-aI,7
+lalamo-0.2.6.dist-info/RECORD,,