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

Added mx.stack c++ frontend impl (#123)

Browse Source 
* stack C++ operation + python bindings
This commit is contained in:
Jason
2023-12-14 16:21:19 -05:00
committed by GitHub
parent e5851e52b1
commit e28b57e371
9 changed files with 191 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
mlx/utils.h
View File

@@ -16,6 +16,15 @@ std::vector<int> broadcast_shapes(
const std::vector<int>& s1,
const std::vector<int>& s2);
bool is_same_shape(const std::vector<array>& arrays);
/**
* Returns the axis normalized to be in the range [0, ndim).
* Based on numpy's normalize_axis_index. See
* https://numpy.org/devdocs/reference/generated/numpy.lib.array_utils.normalize_axis_index.html
*/
int normalize_axis(int axis, int ndim);
std::ostream& operator<<(std::ostream& os, const Device& d);
std::ostream& operator<<(std::ostream& os, const Stream& s);
std::ostream& operator<<(std::ostream& os, const Dtype& d);