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

Add the roll op (#1455)

Browse Source
This commit is contained in:
Angelos Katharopoulos
2024-10-07 17:21:42 -07:00
committed by GitHub
parent f374b6ca4d
commit 9b12093739
6 changed files with 231 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

24
mlx/ops.h
View File

@@ -1456,6 +1456,30 @@ array right_shift(const array& a, const array& b, StreamOrDevice s = {});
array operator>>(const array& a, const array& b);
array view(const array& a, const Dtype& dtype, StreamOrDevice s = {});
/** Roll elements along an axis and introduce them on the other side */
array roll(const array& a, int shift, StreamOrDevice s = {});
array roll(
const array& a,
const std::vector<int>& shift,
StreamOrDevice s = {});
array roll(const array& a, int shift, int axis, StreamOrDevice s = {});
array roll(
const array& a,
int shift,
const std::vector<int>& axes,
StreamOrDevice s = {});
array roll(
const array& a,
const std::vector<int>& shift,
int axis,
StreamOrDevice s = {});
array roll(
const array& a,
const std::vector<int>& shift,
const std::vector<int>& axes,
StreamOrDevice s = {});
/** @} */
} // namespace mlx::core