fix o class ut (#2079)

Author: lvyufeng

mindnlp/core/nn/functional.py

@@ -350,7 +350,7 @@ def _nllloss_nd(input, target, weight=None, ingore_index=-100, reduction='mean')
         raise ValueError(f"input bacth_size should be equal to target batch_size, but got {input.shape[0]} and "
                          f"{target.shape[0]}")
     if input_dim == 1 or input_dim == 2:
-        return nllloss_impl(input, target, weight, reduction, ingore_index)[0]
+        return nllloss_impl(input.float(), target, weight.float(), reduction, ingore_index)[0]
     if input_dim == 4:
         return nllloss_2d_op(input, target, weight, reduction, ingore_index)[0]
     # input_dim==3 or input_dim>4
@@ -1490,7 +1490,7 @@ def pixel_shuffle(input, upscale_factor):
     return ops.pixel_shuffle(input, upscale_factor)
 
 def pixel_unshuffle(input, downscale_factor):
-    return ops.pixel_shuffle(input, downscale_factor)
+    return ops.pixel_unshuffle(input, downscale_factor)
 
 def grid_sample(input, grid, mode='bilinear', padding_mode='zeros', align_corners=False):
     if use_pyboost():

mindnlp/core/ops/comparison.py

@@ -7,7 +7,8 @@
 
 from ._inner import call_ms_func
 
-sort_out = namedtuple('stor_out', ['sorted', 'indices'])
+sort_out = namedtuple('sort_out', ['sorted', 'indices'])
+topk_out = namedtuple('topk_out', ['values', 'indices'])
 # allclose
 has_allclose = hasattr(mindspore.mint, 'allclose')
 def allclose(input, other, rtol=1e-05, atol=1e-08, equal_nan=False):
@@ -173,15 +174,18 @@ def not_equal(input, other):
 def sort(input, *, dim=-1, descending=False, stable=False):
     if use_pyboost() and has_sort:
         out = mindspore.mint.sort(input, dim=dim, descending=descending, stable=stable)
-    out = ops.sort(input, dim, descending)
+    else:
+        out = ops.sort(input, dim, descending)
     return sort_out(sorted=out[0], indices=out[1])
 
 # topk
 has_topk = hasattr(mindspore.mint, 'topk')
 def topk(input, k, dim=-1, largest=True, sorted=True):
     if use_pyboost() and has_topk:
-        return mindspore.mint.topk(input, k, dim, largest, sorted)
-    return ops.topk(input, k, dim, largest, sorted)
+        out = mindspore.mint.topk(input, k, dim, largest, sorted)
+    else:
+        out = ops.topk(input, k, dim, largest, sorted)
+    return topk_out(values=out[0], indices=out[1])
 
 # msort
 def msort(input):

mindnlp/core/serialization.py

@@ -903,6 +903,7 @@ def _maybe_decode_ascii(bytes_str: Union[bytes, str]) -> str:
 
 
 dtype_map = {
+    "DoubleStorage": np.float64,
     "HalfStorage": np.float16,
     "FloatStorage": np.float32,
     'BFloat16Storage': bfloat16,
@@ -912,6 +913,7 @@ def _maybe_decode_ascii(bytes_str: Union[bytes, str]) -> str:
 }
 
 storage_map = {
+    mindspore.float64: "DoubleStorage",
     mindspore.float16: "HalfStorage",
     mindspore.float32: "FloatStorage",
     mindspore.bfloat16: 'BFloat16Storage',

mindnlp/core/types.py

@@ -30,6 +30,7 @@ def __init__(self, type=None, index=None):
                 _target = type.type
                 _id = type.index
             else:
+                print(type)
                 raise TypeError("core.device(): `type` must be type of 'str' or 'core.device'.")
         else:
             raise ValueError("core.device(): `type` can not be None")