From 4ebd647590b079942dfdcd5b2513c74615870ee5 Mon Sep 17 00:00:00 2001
From: huangmengxi <jaclyn.huang@huawei.com>
Date: Fri, 18 Jun 2021 10:19:46 +0800
Subject: [PATCH] fix

---
 mindspore/numpy/array_ops.py | 3 +++
 mindspore/numpy/math_ops.py  | 4 ++--
 2 files changed, 5 insertions(+), 2 deletions(-)

diff --git a/mindspore/numpy/array_ops.py b/mindspore/numpy/array_ops.py
index 56e7d70275..507e0a7077 100644
--- a/mindspore/numpy/array_ops.py
+++ b/mindspore/numpy/array_ops.py
@@ -1825,6 +1825,8 @@ def _check_indices(dims, indices, mode, allow_negative_index=True):
         _raise_unimplemented_error('"raise" mode is not implemented')
     if mode == 'wrap':
         return _mod(indices, F.fill(mstype.float32, shape, dims)).astype(dtype)
+    if mode != 'clip':
+        _raise_value_error('invalid mode. Expected "raise", "wrap", or "clip"')
     zeros = F.fill(dtype, shape, 0)
     clipped = F.select(out_of_lowerbounds, zeros, indices)
     clipped = F.select(out_of_upperbounds, upperbounds, clipped)
@@ -2180,6 +2182,7 @@ def choose(a, choices, mode='clip'):
     else:
         choices = _to_tensor(choices)
         shape_choice = _infer_out_shape(F.shape(a), F.shape(choices)[1:])
+        choices = F.reshape(choices, choices.shape[:1] + _add_unit_axes(choices.shape[1:], len(shape_choice)))
         choices = broadcast_to(choices, (F.shape(choices)[0],) + shape_choice)
 
     if F.rank(a) == 0 or F.rank(choices) == 0:
diff --git a/mindspore/numpy/math_ops.py b/mindspore/numpy/math_ops.py
index b6a82d5044..8b7507572c 100644
--- a/mindspore/numpy/math_ops.py
+++ b/mindspore/numpy/math_ops.py
@@ -4516,7 +4516,7 @@ def digitize(x, bins, right=False):
 
     Args:
         x (Union[int, float, bool, list, tuple, Tensor]): Input array to be binned.
-        bins (Union[int, float, bool, list, tuple, Tensor]): Array of bins. It has to
+        bins (Union[list, tuple, Tensor]): Array of bins. It has to
             be 1-dimensional and monotonic.
         right (boolean, optional): Indicating whether the intervals include the right
             or the left bin edge. Default behavior is ``(right==False)`` indicating
@@ -4539,7 +4539,7 @@ def digitize(x, bins, right=False):
         [1 3 3 4 5]
     """
     x, bins = _to_tensor(x, bins)
-    if F.rank(bins) > 1:
+    if F.rank(bins) != 1:
         _raise_value_error('bins should be 1-dimensional')
     if x.size == 0:
         return x