zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-10-19 08:38:09 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

CPU binary reduction + Nits (#1242)

Browse Source 
* very minor nits

* reduce binary

* fix test
This commit is contained in:
Awni Hannun
2024-06-28 13:50:42 -07:00
committed by GitHub
parent d6383a1c6a
commit 20bb301195
5 changed files with 26 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
python/mlx/nn/layers/base.py
View File

@@ -157,8 +157,8 @@ class Module(dict):
Args:
file_or_weights (str or list(tuple(str, mx.array))): The path to
the weights ``.npz`` file (``.npz`` or ``.safetensors``) or a list of pairs of parameter names
and arrays.
the weights ``.npz`` file (``.npz`` or ``.safetensors``) or a list
of pairs of parameter names and arrays.
strict (bool, optional): If ``True`` then checks that the provided
weights exactly match the parameters of the model. Otherwise,
only the weights actually contained in the model are loaded and
@@ -222,7 +222,7 @@ class Module(dict):
if v_new.shape != v.shape:
raise ValueError(
f"Expected shape {v.shape} but received "
f" shape {v_new.shape} for parameter {k}"
f"shape {v_new.shape} for parameter {k}"
)
self.update(tree_unflatten(weights))

2
python/src/fast.cpp
View File

@@ -83,7 +83,7 @@ void init_fast(nb::module_& parent_module) {
"offset"_a,
"stream"_a = nb::none(),
nb::sig(
"def rope(a: array, dims: int, *, traditinoal: bool, base: float, scale: float, offset: int, stream: Union[None, Stream, Device] = None) -> array"),
"def rope(a: array, dims: int, *, traditional: bool, base: float, scale: float, offset: int, stream: Union[None, Stream, Device] = None) -> array"),
R"pbdoc(
Apply rotary positional encoding to the input.