[GitHub] [incubator-mxnet] mk-61 commented on a change in pull request #19488: Add AMP patching of npi ops in _api_internal module

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

mk-61 commented on a change in pull request #19488:
URL: https://github.com/apache/incubator-mxnet/pull/19488#discussion_r520194107



##########
File path: python/mxnet/amp/amp.py
##########
@@ -80,20 +80,27 @@ def _get_nd_fun_to_wrap(name, module, submodule_dict):
     else:
         func_name = name
         cur_module = module
-    return func_name, cur_module
+    return func_name, [cur_module]
 
 def _get_np_fun_to_wrap(name, ns_prefix):
     for pre, mod, subs in ((_NP_OP_PREFIX, 'numpy', _NP_OP_SUBMODULE_LIST),
                            (_NP_EXT_OP_PREFIX, 'numpy_extension', _NP_EXT_OP_SUBMODULE_LIST),
                            (_NP_INTERNAL_OP_PREFIX, 'numpy._internal', [])):
         if name.startswith(pre):
-            name = name[len(pre):]
+            nm = name[len(pre):]
             for sub in subs:
-                if name.startswith(sub):
-                    return name[len(sub):], sys.modules[f'{ns_prefix}.{mod}.{sub[1:-1]}']
-            return name, sys.modules[f'{ns_prefix}.{mod}']
-    assert False
-    return None  # for pylint
+                if nm.startswith(sub):
+                    func, modules = nm[len(sub):], [sys.modules[f'{ns_prefix}.{mod}.{sub[1:-1]}']]
+                    break
+            else:

Review comment:
       What are your doubts? `Else` clause on line 95 aligns with `for` on line 91 - it gets executed if the loop completes without a break.




----------------------------------------------------------------
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