lalamo 0.3.0__tar.gz → 0.3.1__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 (66) hide show
  1. {lalamo-0.3.0 → lalamo-0.3.1}/PKG-INFO +1 -1
  2. {lalamo-0.3.0 → lalamo-0.3.1}/lalamo/__init__.py +1 -1
  3. {lalamo-0.3.0 → lalamo-0.3.1}/lalamo/model_import/model_specs/common.py +2 -1
  4. {lalamo-0.3.0 → lalamo-0.3.1}/lalamo.egg-info/PKG-INFO +1 -1
  5. {lalamo-0.3.0 → lalamo-0.3.1}/LICENSE +0 -0
  6. {lalamo-0.3.0 → lalamo-0.3.1}/README.md +0 -0
  7. {lalamo-0.3.0 → lalamo-0.3.1}/lalamo/common.py +0 -0
  8. {lalamo-0.3.0 → lalamo-0.3.1}/lalamo/language_model.py +0 -0
  9. {lalamo-0.3.0 → lalamo-0.3.1}/lalamo/main.py +0 -0
  10. {lalamo-0.3.0 → lalamo-0.3.1}/lalamo/message_processor.py +0 -0
  11. {lalamo-0.3.0 → lalamo-0.3.1}/lalamo/model_import/__init__.py +0 -0
  12. {lalamo-0.3.0 → lalamo-0.3.1}/lalamo/model_import/common.py +0 -0
  13. {lalamo-0.3.0 → lalamo-0.3.1}/lalamo/model_import/decoder_configs/__init__.py +0 -0
  14. {lalamo-0.3.0 → lalamo-0.3.1}/lalamo/model_import/decoder_configs/common.py +0 -0
  15. {lalamo-0.3.0 → lalamo-0.3.1}/lalamo/model_import/decoder_configs/executorch.py +0 -0
  16. {lalamo-0.3.0 → lalamo-0.3.1}/lalamo/model_import/decoder_configs/huggingface/__init__.py +0 -0
  17. {lalamo-0.3.0 → lalamo-0.3.1}/lalamo/model_import/decoder_configs/huggingface/common.py +0 -0
  18. {lalamo-0.3.0 → lalamo-0.3.1}/lalamo/model_import/decoder_configs/huggingface/gemma2.py +0 -0
  19. {lalamo-0.3.0 → lalamo-0.3.1}/lalamo/model_import/decoder_configs/huggingface/gemma3.py +0 -0
  20. {lalamo-0.3.0 → lalamo-0.3.1}/lalamo/model_import/decoder_configs/huggingface/llama.py +0 -0
  21. {lalamo-0.3.0 → lalamo-0.3.1}/lalamo/model_import/decoder_configs/huggingface/mistral.py +0 -0
  22. {lalamo-0.3.0 → lalamo-0.3.1}/lalamo/model_import/decoder_configs/huggingface/qwen2.py +0 -0
  23. {lalamo-0.3.0 → lalamo-0.3.1}/lalamo/model_import/decoder_configs/huggingface/qwen3.py +0 -0
  24. {lalamo-0.3.0 → lalamo-0.3.1}/lalamo/model_import/huggingface_generation_config.py +0 -0
  25. {lalamo-0.3.0 → lalamo-0.3.1}/lalamo/model_import/huggingface_tokenizer_config.py +0 -0
  26. {lalamo-0.3.0 → lalamo-0.3.1}/lalamo/model_import/loaders/__init__.py +0 -0
  27. {lalamo-0.3.0 → lalamo-0.3.1}/lalamo/model_import/loaders/common.py +0 -0
  28. {lalamo-0.3.0 → lalamo-0.3.1}/lalamo/model_import/loaders/executorch.py +0 -0
  29. {lalamo-0.3.0 → lalamo-0.3.1}/lalamo/model_import/loaders/huggingface.py +0 -0
  30. {lalamo-0.3.0 → lalamo-0.3.1}/lalamo/model_import/model_specs/__init__.py +0 -0
  31. {lalamo-0.3.0 → lalamo-0.3.1}/lalamo/model_import/model_specs/deepseek.py +0 -0
  32. {lalamo-0.3.0 → lalamo-0.3.1}/lalamo/model_import/model_specs/gemma.py +0 -0
  33. {lalamo-0.3.0 → lalamo-0.3.1}/lalamo/model_import/model_specs/huggingface.py +0 -0
  34. {lalamo-0.3.0 → lalamo-0.3.1}/lalamo/model_import/model_specs/llama.py +0 -0
  35. {lalamo-0.3.0 → lalamo-0.3.1}/lalamo/model_import/model_specs/mistral.py +0 -0
  36. {lalamo-0.3.0 → lalamo-0.3.1}/lalamo/model_import/model_specs/pleias.py +0 -0
  37. {lalamo-0.3.0 → lalamo-0.3.1}/lalamo/model_import/model_specs/polaris.py +0 -0
  38. {lalamo-0.3.0 → lalamo-0.3.1}/lalamo/model_import/model_specs/qwen.py +0 -0
  39. {lalamo-0.3.0 → lalamo-0.3.1}/lalamo/model_import/model_specs/reka.py +0 -0
  40. {lalamo-0.3.0 → lalamo-0.3.1}/lalamo/modules/__init__.py +0 -0
  41. {lalamo-0.3.0 → lalamo-0.3.1}/lalamo/modules/activations.py +0 -0
  42. {lalamo-0.3.0 → lalamo-0.3.1}/lalamo/modules/attention.py +0 -0
  43. {lalamo-0.3.0 → lalamo-0.3.1}/lalamo/modules/common.py +0 -0
  44. {lalamo-0.3.0 → lalamo-0.3.1}/lalamo/modules/decoder.py +0 -0
  45. {lalamo-0.3.0 → lalamo-0.3.1}/lalamo/modules/decoder_layer.py +0 -0
  46. {lalamo-0.3.0 → lalamo-0.3.1}/lalamo/modules/embedding.py +0 -0
  47. {lalamo-0.3.0 → lalamo-0.3.1}/lalamo/modules/kv_cache.py +0 -0
  48. {lalamo-0.3.0 → lalamo-0.3.1}/lalamo/modules/linear.py +0 -0
  49. {lalamo-0.3.0 → lalamo-0.3.1}/lalamo/modules/mlp.py +0 -0
  50. {lalamo-0.3.0 → lalamo-0.3.1}/lalamo/modules/normalization.py +0 -0
  51. {lalamo-0.3.0 → lalamo-0.3.1}/lalamo/modules/rope.py +0 -0
  52. {lalamo-0.3.0 → lalamo-0.3.1}/lalamo/modules/torch_interop.py +0 -0
  53. {lalamo-0.3.0 → lalamo-0.3.1}/lalamo/modules/utils.py +0 -0
  54. {lalamo-0.3.0 → lalamo-0.3.1}/lalamo/quantization.py +0 -0
  55. {lalamo-0.3.0 → lalamo-0.3.1}/lalamo/sampling.py +0 -0
  56. {lalamo-0.3.0 → lalamo-0.3.1}/lalamo/utils.py +0 -0
  57. {lalamo-0.3.0 → lalamo-0.3.1}/lalamo.egg-info/SOURCES.txt +0 -0
  58. {lalamo-0.3.0 → lalamo-0.3.1}/lalamo.egg-info/dependency_links.txt +0 -0
  59. {lalamo-0.3.0 → lalamo-0.3.1}/lalamo.egg-info/entry_points.txt +0 -0
  60. {lalamo-0.3.0 → lalamo-0.3.1}/lalamo.egg-info/requires.txt +0 -0
  61. {lalamo-0.3.0 → lalamo-0.3.1}/lalamo.egg-info/top_level.txt +0 -0
  62. {lalamo-0.3.0 → lalamo-0.3.1}/pyproject.toml +0 -0
  63. {lalamo-0.3.0 → lalamo-0.3.1}/setup.cfg +0 -0
  64. {lalamo-0.3.0 → lalamo-0.3.1}/tests/test_generation.py +0 -0
  65. {lalamo-0.3.0 → lalamo-0.3.1}/tests/test_huggingface_models.py +0 -0
  66. {lalamo-0.3.0 → lalamo-0.3.1}/tests/test_parameter_tree.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: lalamo
3
- Version: 0.3.0
3
+ Version: 0.3.1
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.3.0"
4
+ __version__ = "0.3.1"
5
5
 
6
6
  __all__ = [
7
7
  "REPO_TO_MODEL",
@@ -10,7 +10,6 @@ from jaxtyping import Array, DTypeLike
10
10
  from safetensors.flax import load_file as load_safetensors
11
11
 
12
12
  from lalamo.model_import.decoder_configs import ForeignConfig
13
- from lalamo.modules.torch_interop import torch_to_jax
14
13
  from lalamo.quantization import QuantizationMode
15
14
  from lalamo.utils import MapDictValues
16
15
 
@@ -41,6 +40,8 @@ class WeightsType(Enum):
41
40
 
42
41
  import torch
43
42
 
43
+ from lalamo.modules.torch_interop import torch_to_jax
44
+
44
45
  torch_weights = torch.load(filename, map_location="cpu", weights_only=True)
45
46
  return MapDictValues(lambda v: cast_if_float(torch_to_jax(v), float_dtype), torch_weights)
46
47
 
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: lalamo
3
- Version: 0.3.0
3
+ Version: 0.3.1
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
File without changes