How to call SciPy functions?

f3ba: <p>pytorch doesn’t have the <code>erfcx(x)</code> function implemented (<a href="https://github.com/pytorch/pytorch/issues/31945" rel="nofollow noopener">https://github.com/pytorch/pytorch/issues/31945</a>). However this function is available in scipy (<a href="https://docs.scipy.org/doc/scipy/reference/generated/scipy.special.erfcx.html" rel="nofollow noopener">https://docs.scipy.org/doc/scipy/reference/generated/scipy.special.erfcx.html</a>). So I am trying to call the scipy version from pytorch.</p>
<p>I was thinking of following this tutorial: <a href="https://pytorch.org/tutorials/advanced/numpy_extensions_tutorial.html" rel="nofollow noopener">https://pytorch.org/tutorials/advanced/numpy_extensions_tutorial.html</a>, to implement a pytorch version of <code>erfcx</code> that calls the scipy version and is automatically differentiable.</p>
<p>Is this the recommended approach? I would like to do something that results in the greatest performance possible here.</p>

LeviViana: <p>Maybe creating the extension using scipy directly will crack the case for you. And maybe is the best thing to do in your context.</p>
<p>However, if you want to explore how to implement a PyTorch extension, I suggest you to take a look at the first commit of <a href="https://github.com/LeviViana/torch_sampling/tree/090029f8f1d21671716c38b0edd30f5351dfcda7" rel="nofollow noopener">this repo</a>. It’s a simple extension I implemented last year. Things changed quite a lot on the backend since then, but the idea on how it works is pretty much the same.</p>
<p>I guess, once you get it working really nice on an extension, you can try to submit a PR with your implementation. There are a lot of subtleties when it comes to implement your well working extension natively on PyTorch, so I guess is better to just know how to make a standalone extension work first.</p>

f3ba: <p>Here is what I did:</p>
<pre><code class="lang-python">class ErfcxFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input):
        result = scipy.special.erfcx(input)
        ctx.save_for_backward(input, result)
        return result
    
    @staticmethod
    def backward(ctx, grad_output):
        input, result = ctx.saved_tensors
        df = -2 / np.sqrt(np.pi) + 2 * input * result
        return df * grad_output

erfcx = ErfcxFunction.apply

def logerfc(x):
    return torch.where(x &gt; 0, erfcx(x).log() - x**2, x.erfc().log())

def logerfcx(x):
    return torch.where(x &lt; 0, x.erfc().log() + x**2, erfcx(x.log()))
</code></pre>
<p>However this is very slow:</p>
<pre><code class="lang-auto">A = torch.randn(50,50)
%timeit logerfcx(A) # 428 µs ± 22.9 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
%timeit torch.erfc(A) # 13.9 µs ± 561 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
</code></pre>
<p>Anything I can do to make this faster?</p>

Yaroslav_Bulatov: <p>The slowness is probably due to scipy implementation of erfcx. In colab env it’s 13x slower than torch.erfc.</p>
<p>You could do a naive reimplementation of erfcx, although it will have numerical precision issues</p>
<pre><code class="lang-auto">torch.manual_seed(123456)
def erfcx(A): return torch.exp(A*A)*torch.erfc(A)
A = torch.randn(50,50)
A0=A.numpy()

torch.max(abs(erfcx(A)-torch.tensor(scipy.special.erfcx(A0)))) # tensor(0.0625)
</code></pre>

f3ba: <p>I need <code>erfcx</code> because otherwise <code>logerfc</code> won’t work on the tails of the Gaussian.</p>