Define class (torch.nn.Module) correctly

username: Soroush_Najafian
cooked: <p>Hi. I am try to define import torch <code>class LuminDist(torch.nn.Module):</code> for data fitting with pymc3.<br>
here is my code</p>
<pre><code class="lang-auto">import numpy as np
import matplotlib.pyplot as plt
plt.style.use('seaborn-darkgrid')
import pymc3 as pm
from scipy.integrate import quad
import theano.tensor as tt
import theano
import pdb
from pymc3.backends import Text
import torch
def integrand(z, H0, Om):
	h = H0/100.0
	Or = 4.16e-5/h**2
	# For LCDM, putting following line.
	fz = 1
	# fz = (1 + z)**(3 * (1 + w0 + wa)) * np.exp(-3 * wa * z/(1+z))
	Hz = (Or*(1+z)**4 + Om*(1+z)**3 + (1 - Om - Or)*fz)**0.5
	return 1.0 / Hz

    def integral(z, H0, Om,M):
        como_dist= quad(integrand(z,H0,Om), 0, z, args=(H0, Om))[0]

        D_L = (1 + z) * como_dist
        return 5*np.log10(D_L) + M
import torch

class LuminDist(torch.nn.Module):
    
    def __init__(self):
        super(LuminDist, self).__init__()
        
        self.H0 = torch.nn.Parameter(torch.tensor(0.0, dtype=torch.float64))
        self.Om = torch.nn.Parameter(torch.tensor(0.0, dtype=torch.float64))
        self.M = torch.nn.Parameter(torch.tensor(0.0, dtype=torch.float64))
        self.logs = torch.nn.Parameter(torch.tensor(0.0, dtype=torch.float64))
      
   

    def forward(self, x, y):
        mean = integral(z,H0,Om,M)
        loglike = -0.5 * torch.sum((y - mean) ** 2 * torch.exp(-2 * self.logs) + 2 * self.logs)
        return loglike


import theano
import theano.tensor as tt

class TorchOp(tt.Op):
    
    def __init__(self, module, params, args=None):
        self.module = module
        self.params = list(params)
        if args is None:
            self.args = tuple()
        else:
            self.args = tuple(args)
    
    def make_node(self, *args):
        if len(args) != len(self.params):
            raise ValueError("dimension mismatch")
        args = [tt.as_tensor_variable(a) for a in args]
        return theano.gof.Apply(self, args, [tt.dscalar().type()] + [a.type() for a in args])
    
    def infer_shape(self, node, shapes):
        return tuple([()] + list(shapes))
    
    def perform(self, node, inputs, outputs):
        for p, value in zip(self.params, inputs):
            p.data = torch.tensor(value)
            if p.grad is not None:
                p.grad.detach_()
                p.grad.zero_()

        result = self.module(*self.args)
        result.backward()
        
        outputs[0][0] = result.detach().numpy()
        for i, p in enumerate(self.params):
            outputs[i+1][0] = p.grad.numpy()

    def grad(self, inputs, gradients):
        for i, g in enumerate(gradients[1:]):
            if not isinstance(g.type, theano.gradient.DisconnectedType):
                raise ValueError(
                    "can't propagate gradients wrt parameter {0}".format(i + 1)
                )
        return [gradients[0] * d for d in self(*inputs)[1:]]

</code></pre>
<p>I do not know how define <code>mean=integral(z,H0,Om,M)</code> correctly in this class.<br>
Any advice and suggestions will be greatly appreciated.</p>

username: JuanFMontesinos
cooked: <p>I don’t really know what does quad do but you can define it like that if z H0 and Om are torch tensors.<br>
Just replace np log by torch.</p>

username: Soroush_Najafian
cooked: <p>Thanks for your response. quad is integration methods from scipy . I changed <code>np </code> to <code>torch</code>,but when I run this part of program:</p>
<pre><code class="lang-auto"># Instantiate the PyTorch model
model = torch.jit.script(LuminDist())

# It's useful to select the parameters directly instead of using model.parameters()
# so that we make sure that the order is as expected
params = [model.H0, model.Om,model.M, model.logs]

# The "forward" method of the torch op takes the data as well
args = [torch.tensor(x, dtype=torch.double), torch.tensor(y_obs, dtype=torch.double)]

</code></pre>
<p>I get this error:</p>
<pre><code class="lang-auto">UnsupportedNodeError                      Traceback (most recent call last)

&lt;ipython-input-42-0a5cf026ea8a&gt; in &lt;module&gt;()
      1 # Instantiate the PyTorch model
----&gt; 2 model = torch.jit.script(LuminDist())
      3 
      4 # It's useful to select the parameters directly instead of using model.parameters()
      5 # so that we make sure that the order is as expected

12 frames

/usr/local/lib/python3.7/dist-packages/torch/jit/frontend.py in __call__(self, ctx, node)
    276         method = getattr(self, 'build_' + node.__class__.__name__, None)
    277         if method is None:
--&gt; 278             raise UnsupportedNodeError(ctx, node)
    279         return method(ctx, node)
    280 

UnsupportedNodeError: try blocks aren't supported:
  File "/usr/local/lib/python3.7/dist-packages/scipy/integrate/quadpack.py", line 378
                   5: "The integral over this cycle is probably divergent or slowly convergent."}

    try:
    ~~~ &lt;--- HERE
        msg = msgs[ier]
    except KeyError:
'integral' is being compiled since it was called from 'LuminDist.forward'
  File "&lt;ipython-input-40-efadf1247617&gt;", line 16
    def forward(self, x, y):
        mean = integral(z,H0,Om,M)
        ~~~~~~~~~~~~~~~~~~~~~~~~~ &lt;--- HERE
        loglike = -0.5 * torch.sum((y - mean) ** 2 * torch.exp(-2 * self.logs) + 2 * self.logs)
        return loglike

</code></pre>

username: JuanFMontesinos
cooked: <p>Yeh I see.<br>
All the ops have to be pytorch’s in order the autograd to work.</p>
<p>In case it is useful there is a bayesian optimization library for pytorch <a href="https://botorch.org/tutorials/optimize_stochastic" class="inline-onebox" rel="noopener nofollow ugc">BoTorch · Bayesian Optimization in PyTorch</a></p>
<p>I’m afraid I don’t use theano. But in short the story is that pytorch’s autograd generates and backpropagates gradients automatically. To do so you have to use pytorch ops (as they already have backprop and gradients coded.) If you want to use 3rd party libraries which aren’t based in pytorch you should implement their backward function the autograd to work.</p>
<p>You can implement a custom backward function following this tutorial.<br>
<a href="https://pytorch.org/tutorials/beginner/examples_autograd/two_layer_net_custom_function.html" class="onebox" target="_blank" rel="noopener nofollow ugc">https://pytorch.org/tutorials/beginner/examples_autograd/two_layer_net_custom_function.html</a><br>
Altough I don’t really know if jit supports that.</p>

username: Soroush_Najafian
cooked: <p>Thanks for your feedback.</p>