Thanks for the suggestion for pm.Deterministic. Unfortunately i cannot do a .sum() since it would be incompatible with my model objective, so that rules out pm.Potential.
But i am still struggling to make aesara op/Deterministic work for vectors i.e. an aesara op that returns a vector that can be used inside Deterministic. It all works for a scalar!
It’s difficult to post the full-code, but here is a pseudo-code with a simple scalar to vector change that actually breaks. The code is adapted from this v3 code.
class MyOp:
make_node():
# 1st value is return type and it WORKS if its at.dscalar.type()
out_t = [at.dvector.type()] + [...]
Apply(.., out_t)
perform():
outstorage[0][0] = blackbox_output # this is a vector from blackbox
outstorage[1][0] = blackbox_grads # these are grads from blackbox
# this is expected to use the return value of the custom op
pm.Deterministic("obs", my_op(*model_vars)[0])
If I use at.dscalar as return type and blackbox_output is a scalar, with both pm.Potential and pm.Deterministic all works correctly.
However, if out_t is a at.dvector (with blackbox_output as a vector), i get a compile_pymc error:
AssertionError: Something inferred a shape with 0 dimensions for a variable with 1 dimensions for the variable.
MyOp.0 [id A] <TensorType(float64, (None,))>
My input parameters to aesara op are all vectors, so i am confident that vectors work correctly. Its only the return value as vectors that I am struggling. Not sure if I have to pass anything extra (e.g. a shape parameter somewhere). Is there such an example of theano op that returns a vector and is wrapped from pymc with potential/deterministic?