[GitHub] [incubator-mxnet] barry-jin commented on a change in pull request #20105: [FFI] part4: npx.embedding, npx.topk, npx.layer_norm, npx.leaky_relu

[GitBox <gi...@apache.org>]

barry-jin commented on a change in pull request #20105:
URL: https://github.com/apache/incubator-mxnet/pull/20105#discussion_r610274484



##########
File path: python/mxnet/ndarray/numpy_extension/_op.py
##########
@@ -354,9 +355,12 @@ def batch_norm(x, gamma, beta, running_mean, running_var, eps=1e-3, momentum=0.9
     out : NDArray or list of NDArrays
         The output of this function.
     """
-    return _api_internal.batch_norm(x, gamma, beta, running_mean, running_var, eps, momentum,
-                                    fix_gamma, use_global_stats, output_mean_var, axis,
-                                    cudnn_off, min_calib_range, max_calib_range)
+    out = _api_internal.batch_norm(x, gamma, beta, running_mean, running_var, eps, momentum,
+                                   fix_gamma, use_global_stats, output_mean_var, axis,
+                                   cudnn_off, min_calib_range, max_calib_range)
+    if isinstance(out, NDArrayBase):
+        return out
+    return list(out)

Review comment:
       Because the out can be NDArraBase type or ADT type. This branch will return out directly if its NDArray, otherwise convert ADT to python list and return list of NDArrays. 




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org