Use numpy in script class (torch.jit.script)

I was wondering if I can use numpy APIs in a function which is going to be scripted by torch.jit.script. I have this simple function which does not work:

import torch
import torch.nn as nn

class MyModule(nn.Module):
def init(self):
super(MyModule, self).init()

def call_np():
    return torch.jit.export(np.random.choice(2, p=[.95,.05]))

def forward(self):

def func(self):
    done = self.call_np()
    print (done)

scripted_module = torch.jit.script(MyModule())

which results in:

RuntimeError                              Traceback (most recent call last)
<ipython-input-133-ab1ce37d6edc> in <module>()
     18         print (done)
---> 20 scripted_module = torch.jit.script(MyModule())
     21 scripted_module.func()

C:\ProgramData\Anaconda3\lib\site-packages\torch\jit\ in script(obj, optimize, _frames_up, _rcb)
   1202     if isinstance(obj, torch.nn.Module):
-> 1203         return torch.jit.torch.jit._recursive.recursive_script(obj)
   1205     qualified_name = _qualified_name(obj)

C:\ProgramData\Anaconda3\lib\site-packages\torch\jit\ in recursive_script(mod, exclude_methods)
    171     filtered_methods = filter(ignore_overloaded, methods)
    172     stubs = list(map(make_stub, filtered_methods))
--> 173     return copy_to_script_module(mod, overload_stubs + stubs)

C:\ProgramData\Anaconda3\lib\site-packages\torch\jit\ in copy_to_script_module(original, stubs)
     93             setattr(script_module, name, item)
---> 95     torch.jit._create_methods_from_stubs(script_module, stubs)
     97     # Now that methods have been compiled, take methods that have been compiled

C:\ProgramData\Anaconda3\lib\site-packages\torch\jit\ in _create_methods_from_stubs(self, stubs)
   1421     rcbs = [m.resolution_callback for m in stubs]
   1422     defaults = [get_default_args(m.original_method) for m in stubs]
-> 1423     self._c._create_methods(self, defs, rcbs, defaults)
   1425 # For each user-defined class that subclasses ScriptModule this meta-class,

RuntimeError: Unable to cast Python instance of type <class 'int'> to C++ type 'unsigned __int64'

I appreciate any help or comment.

TorchScript only supports PyTorch and the math module, so numpy functions won’t work natively and can’t be exported. You can use torch.jit.ignore as you have done to leave a call to the Python interpreter. Modifying your example slightly and running with the latest version of PyTorch:

def call_np() -> int:
    return np.random.choice(2, p=[.95,.05])

class MyModule(nn.Module):
    def forward(self):

    def func(self):
        done = call_np()
        print (done)

scripted_module = torch.jit.script(MyModule())


graph(%self : __torch__.MyModule):
  %3 : None = prim::Constant() # ../
  %done.1 : int = ^call_np()() # ../
   = prim::Print(%done.1) # ../
  return (%3)


In which we can see ^call_np() which is an upcall to Python. You must also type annotate call_np since all arguments / returns are assumed to be Tensors unless otherwise specified. torch.jit.export is also intended to be used only as a decorator on a function (see here for details).

I couldn’t repro your exact error, could you file an issue on GitHub with a full repro to produce your issue? Even if what’s happening should throw an error the message could be made more clear. I also ran into this issue while trying to repro it which seems like another bug.

1 Like

Thanks for the answer. I’ll fine the the bug on github page.