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

Fixed shift operations issue (#2080)

Browse Source 
* Fixed shift operations issue

* Added tests and fixes

* Fixed loop syntax error

* Added tests for bool

* Fixed typo
This commit is contained in:
Param Thakkar
2025-04-19 02:58:33 +05:30
committed by GitHub
parent 55935ccae7
commit 5f04c0f818
2 changed files with 46 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
mlx/ops.cpp
View File

@@ -4915,8 +4915,10 @@ array operator^(const array& a, const array& b) {
}
array left_shift(const array& a, const array& b, StreamOrDevice s /* = {} */) {
// Bit shift on bool always up-casts to uint8
auto t = promote_types(result_type(a, b), uint8);
auto t = result_type(a, b);
if (t == bool_) {
t = uint8;
}
return bitwise_impl(
astype(a, t, s),
astype(b, t, s),
@@ -4929,8 +4931,10 @@ array operator<<(const array& a, const array& b) {
}
array right_shift(const array& a, const array& b, StreamOrDevice s /* = {} */) {
// Bit shift on bool always up-casts to uint8
auto t = promote_types(result_type(a, b), uint8);
auto t = result_type(a, b);
if (t == bool_) {
t = uint8;
}
return bitwise_impl(
astype(a, t, s),
astype(b, t, s),