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

minor fixes (#631)

Browse Source 
* minor fixes

* var with ddof >= nelements
This commit is contained in:
Awni Hannun
2024-02-05 13:27:49 -08:00
committed by GitHub
parent d75ae52ecd
commit d40a04f8dc
5 changed files with 40 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
mlx/backend/accelerate/matmul.cpp
View File

@@ -46,6 +46,9 @@ inline void matmul_cblas_general(
size_t N = b.shape(-1);
size_t K = a.shape(-1);
if (M == 0 || N == 0) {
return;
}
if (K == 0) {
std::memset(static_cast<void*>(out.data<float>()), 0, out.nbytes());
return;
@@ -94,6 +97,9 @@ inline void matmul_bnns_general(
size_t N = b.shape(-1);
size_t K = a.shape(-1);
if (M == 0 || N == 0) {
return;
}
if (K == 0) {
std::memset(static_cast<void*>(out.data<float>()), 0, out.nbytes());
return;

4
mlx/backend/common/default_primitives.cpp
View File

@@ -132,7 +132,9 @@ inline void matmul_common_general(
size_t M = a.shape(-2);
size_t N = b.shape(-1);
size_t K = a.shape(-1);
if (M == 0 || N == 0) {
return;
}
if (K == 0) {
std::memset(static_cast<void*>(out.data<float>()), 0, out.nbytes());
return;

6
mlx/ops.cpp
View File

@@ -1323,8 +1323,8 @@ array mean(
for (int axis : axes) {
if (axis < -ndim || axis >= ndim) {
std::ostringstream msg;
msg << "[mean] axis " << axis + " is out of bounds for array with "
<< ndim + " dimensions.";
msg << "[mean] axis " << axis << " is out of bounds for array with "
<< ndim << " dimensions.";
throw std::invalid_argument(msg.str());
}
}
@@ -1364,7 +1364,7 @@ array var(
if (ddof != 0) {
auto nelements = compute_number_of_elements(a, axes);
float factor = nelements / (nelements - ddof);
auto factor = nelements / static_cast<float>(std::max(nelements - ddof, 0));
v = multiply(v, array(factor, dtype), s);
}