lalamo 0.2.6__tar.gz → 0.2.7__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (61) hide show
  1. {lalamo-0.2.6 → lalamo-0.2.7}/PKG-INFO +1 -1
  2. {lalamo-0.2.6 → lalamo-0.2.7}/lalamo/__init__.py +1 -1
  3. {lalamo-0.2.6 → lalamo-0.2.7}/lalamo/model_import/model_specs/common.py +3 -2
  4. {lalamo-0.2.6 → lalamo-0.2.7}/lalamo.egg-info/PKG-INFO +1 -1
  5. {lalamo-0.2.6 → lalamo-0.2.7}/LICENSE +0 -0
  6. {lalamo-0.2.6 → lalamo-0.2.7}/README.md +0 -0
  7. {lalamo-0.2.6 → lalamo-0.2.7}/lalamo/common.py +0 -0
  8. {lalamo-0.2.6 → lalamo-0.2.7}/lalamo/language_model.py +0 -0
  9. {lalamo-0.2.6 → lalamo-0.2.7}/lalamo/main.py +0 -0
  10. {lalamo-0.2.6 → lalamo-0.2.7}/lalamo/model_import/__init__.py +0 -0
  11. {lalamo-0.2.6 → lalamo-0.2.7}/lalamo/model_import/common.py +0 -0
  12. {lalamo-0.2.6 → lalamo-0.2.7}/lalamo/model_import/configs/__init__.py +0 -0
  13. {lalamo-0.2.6 → lalamo-0.2.7}/lalamo/model_import/configs/common.py +0 -0
  14. {lalamo-0.2.6 → lalamo-0.2.7}/lalamo/model_import/configs/executorch.py +0 -0
  15. {lalamo-0.2.6 → lalamo-0.2.7}/lalamo/model_import/configs/huggingface/__init__.py +0 -0
  16. {lalamo-0.2.6 → lalamo-0.2.7}/lalamo/model_import/configs/huggingface/common.py +0 -0
  17. {lalamo-0.2.6 → lalamo-0.2.7}/lalamo/model_import/configs/huggingface/gemma2.py +0 -0
  18. {lalamo-0.2.6 → lalamo-0.2.7}/lalamo/model_import/configs/huggingface/gemma3.py +0 -0
  19. {lalamo-0.2.6 → lalamo-0.2.7}/lalamo/model_import/configs/huggingface/llama.py +0 -0
  20. {lalamo-0.2.6 → lalamo-0.2.7}/lalamo/model_import/configs/huggingface/mistral.py +0 -0
  21. {lalamo-0.2.6 → lalamo-0.2.7}/lalamo/model_import/configs/huggingface/qwen2.py +0 -0
  22. {lalamo-0.2.6 → lalamo-0.2.7}/lalamo/model_import/configs/huggingface/qwen3.py +0 -0
  23. {lalamo-0.2.6 → lalamo-0.2.7}/lalamo/model_import/loaders/__init__.py +0 -0
  24. {lalamo-0.2.6 → lalamo-0.2.7}/lalamo/model_import/loaders/common.py +0 -0
  25. {lalamo-0.2.6 → lalamo-0.2.7}/lalamo/model_import/loaders/executorch.py +0 -0
  26. {lalamo-0.2.6 → lalamo-0.2.7}/lalamo/model_import/loaders/huggingface.py +0 -0
  27. {lalamo-0.2.6 → lalamo-0.2.7}/lalamo/model_import/model_specs/__init__.py +0 -0
  28. {lalamo-0.2.6 → lalamo-0.2.7}/lalamo/model_import/model_specs/deepseek.py +0 -0
  29. {lalamo-0.2.6 → lalamo-0.2.7}/lalamo/model_import/model_specs/gemma.py +0 -0
  30. {lalamo-0.2.6 → lalamo-0.2.7}/lalamo/model_import/model_specs/huggingface.py +0 -0
  31. {lalamo-0.2.6 → lalamo-0.2.7}/lalamo/model_import/model_specs/llama.py +0 -0
  32. {lalamo-0.2.6 → lalamo-0.2.7}/lalamo/model_import/model_specs/mistral.py +0 -0
  33. {lalamo-0.2.6 → lalamo-0.2.7}/lalamo/model_import/model_specs/pleias.py +0 -0
  34. {lalamo-0.2.6 → lalamo-0.2.7}/lalamo/model_import/model_specs/polaris.py +0 -0
  35. {lalamo-0.2.6 → lalamo-0.2.7}/lalamo/model_import/model_specs/qwen.py +0 -0
  36. {lalamo-0.2.6 → lalamo-0.2.7}/lalamo/model_import/model_specs/reka.py +0 -0
  37. {lalamo-0.2.6 → lalamo-0.2.7}/lalamo/modules/__init__.py +0 -0
  38. {lalamo-0.2.6 → lalamo-0.2.7}/lalamo/modules/activations.py +0 -0
  39. {lalamo-0.2.6 → lalamo-0.2.7}/lalamo/modules/attention.py +0 -0
  40. {lalamo-0.2.6 → lalamo-0.2.7}/lalamo/modules/common.py +0 -0
  41. {lalamo-0.2.6 → lalamo-0.2.7}/lalamo/modules/decoder.py +0 -0
  42. {lalamo-0.2.6 → lalamo-0.2.7}/lalamo/modules/decoder_layer.py +0 -0
  43. {lalamo-0.2.6 → lalamo-0.2.7}/lalamo/modules/embedding.py +0 -0
  44. {lalamo-0.2.6 → lalamo-0.2.7}/lalamo/modules/kv_cache.py +0 -0
  45. {lalamo-0.2.6 → lalamo-0.2.7}/lalamo/modules/linear.py +0 -0
  46. {lalamo-0.2.6 → lalamo-0.2.7}/lalamo/modules/mlp.py +0 -0
  47. {lalamo-0.2.6 → lalamo-0.2.7}/lalamo/modules/normalization.py +0 -0
  48. {lalamo-0.2.6 → lalamo-0.2.7}/lalamo/modules/rope.py +0 -0
  49. {lalamo-0.2.6 → lalamo-0.2.7}/lalamo/modules/torch_interop.py +0 -0
  50. {lalamo-0.2.6 → lalamo-0.2.7}/lalamo/modules/utils.py +0 -0
  51. {lalamo-0.2.6 → lalamo-0.2.7}/lalamo/quantization.py +0 -0
  52. {lalamo-0.2.6 → lalamo-0.2.7}/lalamo/utils.py +0 -0
  53. {lalamo-0.2.6 → lalamo-0.2.7}/lalamo.egg-info/SOURCES.txt +0 -0
  54. {lalamo-0.2.6 → lalamo-0.2.7}/lalamo.egg-info/dependency_links.txt +0 -0
  55. {lalamo-0.2.6 → lalamo-0.2.7}/lalamo.egg-info/entry_points.txt +0 -0
  56. {lalamo-0.2.6 → lalamo-0.2.7}/lalamo.egg-info/requires.txt +0 -0
  57. {lalamo-0.2.6 → lalamo-0.2.7}/lalamo.egg-info/top_level.txt +0 -0
  58. {lalamo-0.2.6 → lalamo-0.2.7}/pyproject.toml +0 -0
  59. {lalamo-0.2.6 → lalamo-0.2.7}/setup.cfg +0 -0
  60. {lalamo-0.2.6 → lalamo-0.2.7}/tests/test_generation.py +0 -0
  61. {lalamo-0.2.6 → lalamo-0.2.7}/tests/test_huggingface_models.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: lalamo
3
- Version: 0.2.6
3
+ Version: 0.2.7
4
4
  Summary: JAX library for optimization and export of models for use with the UZU inference engine.
5
5
  Requires-Python: <4,>=3.12
6
6
  Description-Content-Type: text/markdown
@@ -1,7 +1,7 @@
1
1
  from lalamo.model_import import REPO_TO_MODEL, ModelSpec, import_model
2
2
  from lalamo.modules import Decoder
3
3
 
4
- __version__ = "0.2.6"
4
+ __version__ = "0.2.7"
5
5
 
6
6
  __all__ = [
7
7
  "REPO_TO_MODEL",
@@ -7,7 +7,6 @@ from jaxtyping import Array, DTypeLike
7
7
  from safetensors.flax import load_file as load_safetensors
8
8
 
9
9
  from lalamo.model_import.configs import ForeignConfig
10
- from lalamo.modules.torch_interop import torch_to_jax
11
10
  from lalamo.quantization import QuantizationMode
12
11
 
13
12
  __all__ = [
@@ -38,6 +37,8 @@ class WeightsType(Enum):
38
37
 
39
38
  import torch
40
39
 
40
+ from lalamo.modules.torch_interop import torch_to_jax
41
+
41
42
  torch_weights = torch.load(filename, map_location="cpu", weights_only=True)
42
43
  return {k: cast_if_float(torch_to_jax(v), float_dtype) for k, v in torch_weights.items()}
43
44
 
@@ -95,7 +96,7 @@ def awq_model_spec(
95
96
  )
96
97
 
97
98
 
98
- def build_quantized_models(model_specs: list[ModelSpec]):
99
+ def build_quantized_models(model_specs: list[ModelSpec]) -> list[ModelSpec]:
99
100
  quantization_compatible_repos: list[str] = [
100
101
  "Qwen/Qwen2.5-3B-Instruct",
101
102
  "Qwen/Qwen2.5-7B-Instruct",
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: lalamo
3
- Version: 0.2.6
3
+ Version: 0.2.7
4
4
  Summary: JAX library for optimization and export of models for use with the UZU inference engine.
5
5
  Requires-Python: <4,>=3.12
6
6
  Description-Content-Type: text/markdown
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes