TypeError: MaxAndArgmax needs a constant axis for arbitrary / custom distribution function
<p>I am trying to define a multivariate, continuous, custom/arbitrary distribution function that is quite complex and sample from it using the NUTS. I am receiving the following error when I do:    </p>

<pre><code>ERROR (theano.gof.opt): SeqOptimizer apply &lt;theano.scan_module.scan_opt.PushOutNonSeqScan object at 0x000000000AC0E390&gt;

ERROR:theano.gof.opt:SeqOptimizer apply &lt;theano.scan_module.scan_opt.PushOutNonSeqScan object at 0x000000000AC0E390&gt;

ERROR (theano.gof.opt): Traceback:
ERROR:theano.gof.opt:Traceback:

ERROR (theano.gof.opt): Traceback (most recent call last):
  File "C:\Anaconda\envs\py3conda\lib\site-packages\theano\gof\opt.py", line 196, in apply
    sub_prof = optimizer.optimize(fgraph)
  File "C:\Anaconda\envs\py3conda\lib\site-packages\theano\gof\opt.py", line 82, in optimize
    ret = self.apply(fgraph, *args, **kwargs)
  File "C:\Anaconda\envs\py3conda\lib\site-packages\theano\scan_module\scan_opt.py", line 216, in apply
    self.process_node(fgraph, node)
  File "C:\Anaconda\envs\py3conda\lib\site-packages\theano\scan_module\scan_opt.py", line 282, in process_node
    **dict(return_list=True))[0].owner
  File "C:\Anaconda\envs\py3conda\lib\site-packages\theano\gof\op.py", line 507, 
in __call__
    node = self.make_node(*inputs, **kwargs)
  File "C:\Anaconda\envs\py3conda\lib\site-packages\theano\tensor\basic.py", line 1252, in make_node
    raise TypeError("MaxAndArgmax needs a constant axis")

TypeError: MaxAndArgmax needs a constant axis
</code></pre>

<p>I have made sure that the output of the custom logp function is a theano variable scalar and, outside of my custom distribution declaration, I am able to successfully compile a theano function that returns the logp and then also the gradient with respect to the multivariate input. My custom distribution function does not use a scan op which seems to play a part in generating the above error. I also receive an infer shape warning as well with regards to a scan op (I can follow up here with that error message).</p>

<p>From the research I've done, it seems to me that the problem is possibly coming from the calculation of the hessian by the NUTS, but I'm not sure. I'm inexperienced with pymc3 and theano so any insight here would be greatly appreciated.</p>

<p>Thanks!</p>

<p>Alex</p>

<p>UPDATE:</p>

<p>I'm now fairly confident that the hessian is causing this error and so this becomes more of a theano issue than a pymc3 one. I was able to test this by evaluating the <code>theano.gradient.hessian()</code> function outside of pymc3 entirely, and I received exactly the same error above. Also, akin to this <a href="https://groups.google.com/forum/#!searchin/theano-users/Maxandargmax$20hessian/theano-users/Qp9dF0ZFo9k/uRXYJiX7vsgJ" rel="nofollow">issue</a>, I was using a <code>tensor.min()</code> in evaluating my likelihood function which caused the above error. However, I replaced min() with sum() to double check and theano failed to compile my function still but didn't throw an error (just blocked)... but only on the gpu. Still keeping the "fake" sum() operation in there, the hessian successfully evaluated on my cpu. Anyway, I think there are issues with my likelihood function as well with regards to continuity and convexity so won't be travelling down this path.</p>