From 6b2d5448f230c73ba70996fce188498ce2c83e3f Mon Sep 17 00:00:00 2001
From: 1ndig0 <1090891928@qq.com>
Date: Sat, 26 Apr 2025 00:14:28 +0800
Subject: [PATCH] Fix the error message in `mx.right_shift` and `mx.left_shift`
 (#2121)

* update right_shift and lef_shift

* simplify

---------

Co-authored-by: Awni Hannun <awni@apple.com>
---
 mlx/ops.cpp | 15 ++++++---------
 1 file changed, 6 insertions(+), 9 deletions(-)

diff --git a/mlx/ops.cpp b/mlx/ops.cpp
index c2aa4786f..f8308c2d5 100644
--- a/mlx/ops.cpp
+++ b/mlx/ops.cpp
@@ -4882,8 +4882,9 @@ array bitwise_impl(
     const array& b,
     BitwiseBinary::Op op,
     const std::string& op_name,
-    const StreamOrDevice& s) {
-  auto out_type = promote_types(a.dtype(), b.dtype());
+    const StreamOrDevice& s,
+    std::optional<Dtype> out_type_ = std::nullopt) {
+  auto out_type = out_type_ ? *out_type_ : promote_types(a.dtype(), b.dtype());
   if (!(issubdtype(out_type, integer) || out_type == bool_)) {
     std::ostringstream msg;
     msg << "[" << op_name
@@ -4928,12 +4929,7 @@ array left_shift(const array& a, const array& b, StreamOrDevice s /* = {} */) {
   if (t == bool_) {
     t = uint8;
   }
-  return bitwise_impl(
-      astype(a, t, s),
-      astype(b, t, s),
-      BitwiseBinary::Op::LeftShift,
-      "left_shift",
-      s);
+  return bitwise_impl(a, b, BitwiseBinary::Op::LeftShift, "left_shift", s, t);
 }
 array operator<<(const array& a, const array& b) {
   return left_shift(a, b);
@@ -4949,7 +4945,8 @@ array right_shift(const array& a, const array& b, StreamOrDevice s /* = {} */) {
       astype(b, t, s),
       BitwiseBinary::Op::RightShift,
       "right_shift",
-      s);
+      s,
+      t);
 }
 array operator>>(const array& a, const array& b) {
   return right_shift(a, b);