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

username: afshin67
cooked: <p>I was wondering if I can use numpy APIs in a function which is going to be scripted by <code>torch.jit.script</code>. I have this simple function which does not work:</p>
<p>import torch<br>
import torch.nn as nn</p>
<p>class MyModule(nn.Module):<br>
def <strong>init</strong>(self):<br>
super(MyModule, self).<strong>init</strong>()</p>
<pre><code>@torch.jit.ignore
def call_np():
    return torch.jit.export(np.random.choice(2, p=[.95,.05]))

def forward(self):
    pass 

@torch.jit.export
def func(self):
    done = self.call_np()
    print (done)
</code></pre>
<blockquote>
<p>scripted_module = torch.jit.script(MyModule())<br>
scripted_module.func()</p>
</blockquote>
<p>which results in:</p>
<pre><code class="lang-auto">---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
&lt;ipython-input-133-ab1ce37d6edc&gt; in &lt;module&gt;()
     18         print (done)
     19 
---&gt; 20 scripted_module = torch.jit.script(MyModule())
     21 scripted_module.func()

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

C:\ProgramData\Anaconda3\lib\site-packages\torch\jit\_recursive.py in recursive_script(mod, exclude_methods)
    171     filtered_methods = filter(ignore_overloaded, methods)
    172     stubs = list(map(make_stub, filtered_methods))
--&gt; 173     return copy_to_script_module(mod, overload_stubs + stubs)
    174 
    175 

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

C:\ProgramData\Anaconda3\lib\site-packages\torch\jit\__init__.py 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]
-&gt; 1423     self._c._create_methods(self, defs, rcbs, defaults)
   1424 
   1425 # For each user-defined class that subclasses ScriptModule this meta-class,

RuntimeError: Unable to cast Python instance of type &lt;class 'int'&gt; to C++ type 'unsigned __int64'
</code></pre>
<p>I appreciate any help or comment.</p>

username: driazati
cooked: <p>TorchScript only supports PyTorch and the <code>math</code> module, so <code>numpy</code> functions won’t work natively and can’t be exported. You can use <code>torch.jit.ignore</code> as you have done to leave a call to the Python interpreter. Modifying your example slightly and running with the latest version of PyTorch:</p>
<pre><code class="lang-python">@torch.jit.ignore
def call_np() -&gt; int:
    return np.random.choice(2, p=[.95,.05])

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

    @torch.jit.export
    def func(self):
        done = call_np()
        print (done)

scripted_module = torch.jit.script(MyModule())
print(scripted_module.func.graph)
scripted_module.func()
</code></pre>
<p>prints</p>
<pre><code class="lang-auto">graph(%self : __torch__.MyModule):
  %3 : None = prim::Constant() # ../test.py:184:4
  %done.1 : int = ^call_np()() # ../test.py:185:15
   = prim::Print(%done.1) # ../test.py:186:8
  return (%3)

0
</code></pre>
<p>In which we can see <code>^call_np()</code> which is an upcall to Python. You must also type annotate <code>call_np</code> since all arguments / returns are assumed to be <code>Tensor</code>s unless otherwise specified. <code>torch.jit.export</code> is also intended to be used only as a decorator on a function (<a href="https://pytorch.org/docs/master/jit.html#torch.jit.export" rel="nofollow noopener">see here for details</a>).</p>
<p>I couldn’t repro your exact error, could you <a href="https://github.com/pytorch/pytorch/issues/new?template=bug-report.md" rel="nofollow noopener">file an issue on GitHub</a> 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 <a href="https://github.com/pytorch/pytorch/issues/30625" rel="nofollow noopener">this issue</a> while trying to repro it which seems like another bug.</p>

username: afshin67
cooked: <p>Thanks for the answer. I’ll fine the the bug on github page.</p>