Wrap Ops during etuplization

[rlouf]

Evaluating etuplized objects fails for some RandomVariable ops whose
__call__ function does not defer to make_node.

In this commit we wrap Ops during etuplization with a class that always
defers __call__ to make_node. We also add a dispatch rule for
etuplize so it also wraps RandomVariable with the same class.

Closes #1002.

Thank you for opening a PR!

Here are a few important guidelines and requirements to check before your PR can be merged:

• There is an informative high-level description of the changes.
• The description and/or commit message(s) references the relevant GitHub issue(s).
• pre-commit is installed and set up.
• The commit messages follow these guidelines.
• The commits correspond to relevant logical changes, and there are no commits that fix changes introduced by other commits in the same branch/BR.
• There are tests covering the changes introduced in the PR.

[rlouf]

Some tests that are unrelated to etuplizeation expectedly fail. This is easily fixed by wrapping an op with MakeRandomVariableNodeOp only if it is a subclass of RandomVariable, but still it makes we wonder whether we shouldn't come up with something in etuples so that this fix only requires the following code:

@dataclass
class MakeRVNodeOp:
    op: Op

    def __call__(self, *args):
        return self.op.make_node(*args)

@etuplize.register(RandomVariable)
def etuplize_random(*args, **kwargs):
    return etuple(MakeRVNodeOp, etuplize.funcs[(object,)](*args, **kwargs))

[codecov]

Codecov Report

Merging #1036 (c84149c) into main (69c1044) will increase coverage by 0.01%.
The diff coverage is 93.61%.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1036      +/-   ##
==========================================
+ Coverage   79.13%   79.14%   +0.01%     
==========================================
  Files         173      173              
  Lines       48492    48565      +73     
  Branches    10966    10327     -639     
==========================================
+ Hits        38374    38439      +65     
- Misses       7625     7632       +7     
- Partials     2493     2494       +1     

Impacted Files	Coverage Δ
aesara/graph/rewriting/unify.py	97.35% <93.61%> (-1.99%)	⬇️
aesara/link/numba/dispatch/scan.py	94.11% <0.00%> (-1.44%)	⬇️
aesara/sparse/basic.py	82.76% <0.00%> (-0.12%)	⬇️
aesara/scan/op.py	85.41% <0.00%> (-0.07%)	⬇️
aesara/tensor/random/basic.py	99.00% <0.00%> (ø)
aesara/link/numba/dispatch/basic.py	92.46% <0.00%> (+0.01%)	⬆️
aesara/link/utils.py	61.84% <0.00%> (+0.71%)	⬆️
aesara/link/numba/dispatch/scalar.py	87.33% <0.00%> (+1.33%)	⬆️

[rlouf]

Codecov complains about a file that is not affected by the changes. @brandonwillard this is ready for review.

[brandonwillard]

I would like to consider some other approaches before fully committing to this one, especially given the way this one complicates the use of Ops in etuple form (i.e. by adding an indirection).

In general, it's not good that Ops can't be used directly. For example, etuple(op, ...) will no longer unify with the results of a shallow etuplize(op(...)) due to the wrapper now introduced by the latter.

Also, since Ops themselves can be converted to etuples (e.g. in order to unify against configurable Op properties), these wrapper objects need to be traversed and unified explicitly, which incurs a cost of its own. I might be possible to simplify or altogether avoid some of this by altering the wrapper's __eq__ so that wrapped objects are equal to the Ops they wrap. This might introduce other complications, though.

Perhaps a better approach would involve the construction of entirely new Op instance clones during etuplize; ones that are equivalent to the cloned Ops aside from their Op.__call__s. One problem with this approach: if we do this by constructing new types, we might need to cache the generated clone types in order to avoid producing an abundance of redundant types; however, we're doing something almost exactly the same as that in this PR (e.g. a distinct wrapper instance is produced for each occurence of the same Op). (It's never been clear to me how type creation is handled memory-wise in Python, so this distinction could be practically irrelevant for all I know.)

[rlouf]

We should be able to unify etuple(op, ...) with the result of etuplize(op(...)), so I don't think this solution (wrapping the Ops) is acceptable.

I gave it more thought, and I think that the solution that introduces the least amount of change to Aesara while not being too hacky is to have a function handle the object's evaluation in etuples and dispatch this function depending on the type of op.

In the present case we would only need to add a function inaesara.graph.rewriting.unify that would be along the lines of:

def eval_Op(op, evaled_args, evaled_kwargs):
    return op.make_node(*evaled_args, **evaled_kwargs)
    
_eval.add((Op,X, X), eval_Op)

In fact, couldn't we used apply for that? It is already dispatched for Op.

[brandonwillard]

I gave it more thought, and I think that the solution that introduces the least amount of change to Aesara while not being too hacky is to have a function handle the object's evaluation in etuples and dispatch this function depending on the type of op.

One of the first things I considered was a change to ExpressionTuple._eval_step that would allow one to specify something other than the use of __call__ (e.g. op(*args)) for evaluation; however, that's a very critical point in the ExpressionTuple evaluation process, and we don't need/want dispatching there—at least if it's not necessary.

I left that approach with the idea that we could simply subclass ExpressionTuple and make that change. We could expose an interface for that exact kind of customization, though, which would make such subclasses very simple. For instance, adding a ExpressionTuple._eval_apply that takes op and op_args values would help.

[brandonwillard]

pythological/etuples#19 illustrates the ExpressionTuple changes. We would need to follow that up here by creating a custom ExpressionTuple subclass that uses Op.make_node in its ExpressionTuple._eval_apply.

[rlouf]

@brandonwillard I made the necessary changes following pythological/etuples#21. Provided the tests pass this is ready on my end.

[brandonwillard]

Looks good. The only thing I would consider changing is the limited use of a custom ExpressionTuple. We should always use Op.make_node.

[rlouf]

I considered that, it's just that the apply node created by RandomVariables have the rng as well as the output variable as outputs, and we're only interested in returning the latter. I could add an OpExpressionTuple for other Ops.

[brandonwillard]

I considered that, it's just that the apply node created by RandomVariables have the rng as well as the output variable as outputs, and we're only interested in returning the latter. I could add an OpExpressionTuple for other Ops.

We're always interested in all outputs of a node. We might not be using them in our current rewrites, but they need to be available at some point. Regardless, we really shouldn't be "special casing" RandomVariables in this way.

[rlouf]

Regardless, we really shouldn't be "special casing" RandomVariables in this way.

I'm not sure about this one. Remember that this PR was created to address an issue with etuplization of RandomVariables specifically.

Nevermind, I agree.

[rlouf]

Evaluating Ops now returns an Apply node, as opposed to a TensorVariable (or an error in the case of RandomVariables) on the main branch. Is that what we want? Asking before making changes elsewhere in the codebase.

[brandonwillard]

Evaluating Ops now returns an Apply node, as opposed to a TensorVariable (or an error in the case of RandomVariables) on the main branch. Is that what we want? Asking before making changes elsewhere in the codebase.

For ease of use with kanren, we might want to return the outputs list.

[rlouf]

It is trickier than I thought because there is an assumption throughout the code that inputs to Apply nodes are passed as arguments individually. This causes issues when evaluating nested OpExpressionTuples as Apply nodes interpret [Var] as MakeVector(Var):

import aesara
import aesara.tensor as at
from etuples import etuplize, etuple

a = at.scalar('a')
b = at.scalar('b')
c = at.scalar('c')

x_et = etuple(at.mul, etuple(at.add, b, b), etuple(at.add, a, a))
x_val = x_et.evaled_obj
print(x_val)
# [Elemwise{mul,no_inplace}.0] expected

print(x_val[0].owner.inputs)
# [MakeVector{dtype='float64'}.0, MakeVector{dtype='float64'}.0]

print(aesara.dprint(x_val[0].owner.inputs))
# MakeVector{dtype='float64'} [id A]
# |Elemwise{add,no_inplace} [id B]
#   |b [id C]
#   |b [id C]
# MakeVector{dtype='float64'} [id D]

I am not even sure it makes sense to return all the outputs, e.g. for random variables:

z_rv = at.random.normal(0, 1)
z_et = etuplize(z_rv)
z_et._evaled_obj = z_et.null
print(z_et.evaled_obj)
# [normal_rv{0, (0, 0), floatX, False}.0, normal_rv{0, (0, 0), floatX, False}.out]

there is bound to be some confusion here. I think that in terms of returned value we have to do something similar to what the corresponding Op does with __call__. Because it behaves differently from the other Ops we would have no choice but to "special case" it. We'd only be propagating design decisions that were made upstream.

But I feel that's a hack and this is touching on a deeper issue with Aesara.

[rlouf]

Following up on a discussion with @brandonwillard

This indeed touches on a much deeper issue, that is the fact that Aesara's IR lacks the ability to represent selections of specific multi-outputs outputs. All the logic is hidden in Op.__call__ as the graph is being built.

To illustrate the problem, let's look again at the etuplization of RandomVariables with the current state of the PR:

import aesara.tensor as at
from etuples import etuple, etuplize

u_rv = at.random.normal(0, 1)
u_et = etuplize(u_rv)
print(u_et)
# : oExpressionTuple((ExpressionTuple((<class 'aesara.tensor.random.ba...ormalRV'>, 'normal', 0, (0, 0), 'floatX', False)), RandomGeneratorSharedVariable(<Generator(PCG64) at 0x7FA6605E5EE0>), TensorConstant{[]}, TensorConstant{11}, TensorConstant{0}, TensorConstant{1}))

# reset the saved evaled obj
u_et._evaled_obj = u_et.null
print(u_et.evaled_obj)
# [normal_rv{0, (0, 0), floatX, False}.0, normal_rv{0, (0, 0), floatX, False}.out]

We get two outputs, and we should only be getting the second one! Of course we can hack our way around this, "special casing" the random variables as was mentioning above in the thread. It's unsatisfactory, but there's worse. Let's look at what happens when we etuplize one of the outputs of a multi-output Op and evaluate the resulting object:

import aesara.tensor as at
from etuples import etuple, etuplize

a = at.matrix('a')
u, v, w = at.nlinalg.svd(a)
u_et = etuplize(u)
print(u_et)
# oExpressionTuple((ExpressionTuple((<class 'aesara.tensor.nlinalg.SVD'>, 1, 1)), a))

# reset the saved evaled obj
u_et._evaled_obj = u_et.null
print(u_et.evaled_obj)
# [SVD{full_matrices=1, compute_uv=1}.0, SVD{full_matrices=1, compute_uv=1}.1, SVD{full_matrices=1, compute_uv=1}.2]

We have lost some precious information in the process of etuplization, which is that u is the first output of the apply node created by at.nlinalg.svd, and are thus unable to evaluate the expression tuple to this variable. What we need to solve this issue is a more principled approach, and we can already have an idea about what this would look like in the etuplized version of the graph.

What we essentially need is adding an indexing operator when we etuplize the Ops, so that etuplizing the RandomVariable above would return

ExpressionTuple(
    nth,
    1,
    oExpressionTuple((ExpressionTuple((<class 'aesara.tensor.random.ba...ormalRV'>, 'normal', 0, (0, 0), 'floatX', False)), RandomGeneratorSharedVariable(<Generator(PCG64) at 0x7FA6605E5EE0>), TensorConstant{[]}, TensorConstant{11}, TensorConstant{0}, TensorConstant{1}))
)

where oExpressionTuple(...) evaluates to the list of outputs of the Apply nodes. the nth function returns the element at index 1 in this list so the output is the RandomVariable returned by the Apply node.

This also works in the at.nlinalg.svd example:

ExpressionTuple(
    nth,
    0,
    oExpressionTuple((ExpressionTuple((<class 'aesara.tensor.nlinalg.SVD'>, 1, 1)), a))
)

which evaluates to the first Variable in the outputs list of the Apply node. We may be able to get by temporarily using the information in the Op's default_output attribute, and the index attribute of each Variable, but this doesn't solve the root issue that is the inability to represent the selection of outputs in Aesara's IR.

[brandonwillard]

This indeed touches on a much deeper issue, that is the fact that Aesara's IR lacks the ability to represent selections of specific multi-outputs outputs. All the logic is hidden in Op.__call__ as the graph is being built.

N.B. This is also why a translation to something like e-graphs isn't entirely straightforward.

[rlouf]

Adding nth

I added a nth function that takes an index n, and apply node node and returns node.outputs[n]. I also updated the tests related to etuplization/evaluation because they were not really testing the evaluation (the evaluated object saved during etuplization instead). I added a test for multiple output nodes and nodes that have a default output. Both pass thanks to that nth function.

I believe this nth function should be promoted to a node in Aesara's IR, so that Ops __call__ function becomes:

def __call__(self, *inputs, **kwargs):
    node = self.make_node(*inputs, **kwargs)
    outputs = node.outputs
   
    if len(outputs) == 0:
         return nth(0, outputs)
    elif:
        if self.default_output is not None:
            return nth(self.default_output, outputs)
    else:
         return [nth(i, output) for i, output in enumerate(outputs)]

We could probably rename default_output to selected_output and set its default to 0. The difficulty with this approach is that it changes the variables' owner, and a LOT of logic depends on the direct link variable - apply node, but it surely is worth it long-term. For that reason I would strongly advise against returning the variable directly for single-output nodes, but nth(0, node.outputs).

On the necessity of keeping default_output

I am also starting to think that we should seriously re-evaluate the necessity of keeping this default_output property. I searched through the codebase and it is set only in an experiment module (to store a C object), in nnet (seems to be OpenMP constraint), and in RandomVariable. It does not seem to be something fundamental or even necessary, but present because of the legacy backend and because it seems more convenient for the users in the case of RandomVariable.

The first two use cases will be deprecated, so in the end this will only be used by RandomVariable. From the perspective of aesara-devs/aemcmc#66 as well, I am starting to wonder if that convenience is not causing more problems than it solves. Making the updates of the rng state explicit does not seem so bad to me anymore.

Multiple-output nodes

Regarding multiple output Ops, I added support for the case where the caller calls etuplize of the list that contains both outputs by dispatching car to cons and cdr to tuple(x) + ([],). I added a test to make sure it evaluates back to the list.

Next

• Add tests around unification for multiple outputs and default output node.
• Assess the extent to which this change is going to affect the rewrites.
• It also seems that etuplized is not applied recursively on lists

[rlouf]

This is not immediately relevant to this PR, but let’s consider the single output case and how etuplization, unification and reification currently behave in this case:

import aesara.tensor as at
from etuples import etuplize, etuple
from unification import unify, reify, var

x_at = at.scalar('x')
y_at = at.scalar('y')
z_at = at.add(x_at, y_at)

print(etuplize(z_at))
# oExpressionTuple(ExpressionTuple(at.Add), x, y)

x_lv, y_lv = var('x'), var('y')
z_et = etuple(at.add, x_lv, y_lv)

s = unify(z_et, z_at)
print(s)
# {~x: x, ~y: y}

res = reify(z_et, s)
print(res)
# oExpressionTuple((<aesara.tensor.elemwise.Elemwise object at 0x7f04d3989c30>, x, y))

print(res.evaled_obj)
# Elemwise{add,no_inplace}.0

Now let’s imagine that we have added a nth node to our Aesara graphs. Let’s consider the scenario where z_at corresponds to nth(0, outputs) (where outputs is the outputs of the Apply node). Etuplization following the current rules would give the following:

print(etuplize(z_at))
# ExpressionTuple(nth, 0, oExpressionTuple(ExpressionTuple(at.Add), x, y))

Not knowing about Aesara’s internals, or for convenience, one could try to unify z_at with z_et defined as:

x_lv, y_lv = var('x'), var('y')
z_et = etuple(at.add, x_lv, y_lv)
s = unify(z_at, z_et)

Now, we have two choices: we can either change the logic in unify in Aesara, or z_at directly to oExpressionTuple((..), x, y) but then we would need to change the logic in reify. The former is more appealing to me.

Anyway, this is a question of boundaries and what part of the library we want to do the heavy lifting. The more I work at the interface between the Python IR and its etuplized counterpart the more I am starting to get aware of this. Convenience has a price, and it's good to keep that in mind.

[rlouf]

Well, what I was discussing above has direct consequences today because of default_output. The following fails to unify:

class MyDefaultOutputOp(Op):
    default_output = 1

    def make_node(self, *inputs):
        outputs = [MyType()(), MyType()()]
        return Apply(self, list(inputs), outputs)

    def perform(self, node, inputs, outputs):
        outputs[0] = np.array(np.array(inputs[0]))
        outputs[1] = np.array(np.array(inputs[1]))

x_at = at.vector("x")
y_at = at.vector("y")
op1_np = MyDefaultOutputOp()
q_at = op1_np(x_at, y_at)

x_lv, y_lv = var('x'), var('y')
q_et = etuple(op1_np, x_lv, y_lv)

s = unify(q_et, q_at)
assert s[x_lv] == x_at
assert s[y_lv] == y_at

Relations in AeMCMC can be written using nth and unification will work (we can also special-case the logic for RandomVariables, but we already established we don't want to do that). But to me that seems like a workaround for a much deeper design issue.

This PR has been tightly linked to the design of RandomVariables since the beginning and I need to spend more time thinking about this point. Especially considering the other difficulties encountered in AeMCMC, for instance in aesara-devs/aemcmc#66. Will open a separate discusison.

[rlouf]

I have fixed the unification with the default outputs of a multiple-output Apply nodes. I am now trying to define the behavior of etuplize when passed a list of outputs which can happen with the output of multiple output Op\s:

from aesara.graph.basic import Apply
from aesara.graph.op import Op
from aesara.graph.type import Type

class MyType(Type):
    def filter(self, data):
        return data

    def __eq__(self, other):
        return isinstance(other, MyType)

    def __hash__(self):
        return hash(MyType)

    def __repr__(self):
        return "MyType()"

class MyMultiOutOp(Op):
    def make_node(self, *inputs):
        outputs = [MyType()(), MyType()()]
        return Apply(self, list(inputs), outputs)

    def perform(self, node, inputs, outputs):
        outputs[0] = np.array(inputs[0])
        outputs[1] = np.array(inputs[0])

Etuplizing directly will work, but evaluating the resulting etuple will fail:

import aesara.tensor as at
from etuples import etuplize

x_at = at.vector("x")
y_at = at.vector("y")
op1_np = MyMultiOutOp()
z_at = op1_np()

z_et = etuplize(z_at)
z_et._evaled_obj = z_et.null
try:
    z_et.evaled_obj
except Exception as e:
    print(e)
# ExpressionTuple does not have a callable operator.

This can be solved by etuplizing lists to etuple(cons, *args, []) (which should probably be done in the etuples library directly):

from cons import cons
from cons.core import _car, _cdr
from etuples import etuplize


def car_List(x):
    return cons

_car.add((list,), car_List)

def cdr_List(x):
    return  (*x, [])

_cdr.add((list,), cdr_List)

print(etuplize(z_at))
# e(<class 'cons.core.ConsPair'>, MyMultiOutOp.0, MyMultiOutOp.1, [])

z_et = etuplize(z_at)
z_et._evaled_obj = z_et.null
print(z_et.evaled_obj)
# [MyMultiOutOp.0, MyMultiOutOp.1]

However there is still one issue with the previous snippet: etuplize did not work recursively. That's because after etuplizing the list to e(cons, *args, []) the etuplizing function evaluates the isinstance(x, ConsPair) branch which returns the etuple directly.

[brandonwillard]

On a related note: this is reminding me that we really need hash-consing and/or caching for these etuple conversions.

[rlouf]

This PR got a little out of hand, and I will leave etuplizing a list of multiple outputs for another PR as it seems more complicated than I originally thought, and is non-blocking for now. Let's come back to our original problem, which was being able to have "true" relations in AeMCMC, in the sense that they can work both ways, and merge this PR when we are able to do that.

Summary

The first issue we ran into was that __call__ of RandomVariables does not have the same signature as make_node, which lead to errors during evaluation. This was solved by creating a OpExpressionTuple and modifying some of the etuplization and evalution logic in etuples. This seems fair to me, having Op-specific expression tuples makes sense.

Then we realized that when we evaluated the expression tuples obtained after etuplizing RandomVariables gave us the two outputs that are generated by make_node instead of the just the Variable. That's because Aesara doesn't have an operator to say "This Variable is the nth output of the Apply node". So we faked that behavior by adding a nth operator when etuplizing Variables that are the output of a multiple-output Apply node. This should be ported to Aesara's IR.

And then we realized that we could not unify, say etuple(etuplize(at.random.normal), var(), var(), var(), 0, 1) with etuplize(at.random.normal(0, 1)) since the latter starts with etuple(nth, 1, ...). So we modified the unifying logic to be able to do that with Variables that are a default output (we don't do it with multiple-output nodes as in this case nth must be specified if we want a particular output). This makes me doubt the conceptual validity of the use of default_output, and thus the current design of RandomVariables. Imho we should not have to bend unification to make this happen.

Where we are now

Unification works on simple cases

Let's confirm that unification works despite the fact that RandomVariable\s use default_output:

from etuples import etuple, etuplize
from unification import var, unify

Y_rv = at.random.normal(0, 1)
Y_et = etuple(etuplize(at.random.normal), var(), var(), var(), at.as_tensor(0.), at.as_tensor(1.))

print(etuplize(Y_rv))
# e(<function nth at 0x7fa66e32d360>, 1, oExpressionTuple((ExpressionTuple((<class 'aesara.tensor.random.ba...ormalRV'>, 'normal', 0, (0, 0), 'floatX', False)), RandomGeneratorSharedVariable(<Generator(PCG64) at 0x7FA66B45E500>), TensorConstant{[]}, TensorConstant{11}, TensorConstant{0}, TensorConstant{1})))

print(unify(Y_rv, Y_et))
# {~_95: RandomGeneratorSharedVariable(<Generator(PCG64) at 0x7FA66B45E500>), ~_96: TensorConstant{[]}, ~_97: TensorConstant{11}}

s = unify(Y_et, Y_rv)
print(s)
# {~_95: RandomGeneratorSharedVariable(<Generator(PCG64) at 0x7FA66B45E500>), ~_96: TensorConstant{[]}, ~_97: TensorConstant{11}}

Reification does not work even on simple cases

Reification does not give the expected result:

res = reify(Y_et, s)
print(res)
# e(e(<class 'aesara.tensor.random.basic.NormalRV'>, normal, 0, (0, 0), floatX, False), RandomGeneratorSharedVariable(<Generator(PCG64) at 0x7FA66B45C820>), TensorConstant{[]}, TensorConstant{11}, TensorConstant{0.0}, TensorConstant{1.0})

try:
    res.evaled_obj
except TypeError as e:
    print(e)
# too many positional arguments

which is because reification does not return an OpExpressionTuple, and does not add the nth operator.

Unification does not work on more complex cases

Let's consider now the case, common in AeMCMC, where we have random variables who parameters depend on other random variables. For instance a Poisson distribution with a Gamma prior:

srng = at.random.RandomStream(0)
alpha_tt = at.scalar("alpha")
beta_tt = at.scalar("beta")
z_rv = srng.gamma(alpha_tt, beta_tt)
Y_rv = srng.poisson(z_rv)

Which we try to unify with:

alpha_lv, beta_lv = var(), var()
z_rng_lv = var()
z_size_lv = var()
z_type_idx_lv = var()
z_et = etuple(
    etuplize(at.random.gamma), z_rng_lv, z_size_lv, z_type_idx_lv, alpha_lv, beta_lv
)
Y_et = etuple(etuplize(at.random.poisson), var(), var(), var(), z_et)

Unsurprisingly given the previous section, unifying z_rv and z_et works:

print(unify(z_rv, z_et))
# {~_219: RandomGeneratorSharedVariable(<Generator(PCG64) at 0x7FA66B45E340>), ~_220: TensorConstant{[]}, ~_221: TensorConstant{11}, ~_217: alpha, ~_218: OpExpressionTuple((ExpressionTuple((<class 'aesara.tensor.elemwise....eDiv object at 0x7fa69b0fef20>, <frozendict {}>)), TensorConstant{1.0}, beta))}

But unification fails for Y_rv and Y_et:

print(unify(Y_rv, Y_et))
# False

However unification succeeds if we use the nth operator explicitly in the expression for z_et:

z_et = etuple(nth, 1, etuple(
    etuplize(at.random.gamma), z_rng_lv, z_size_lv, z_type_idx_lv, alpha_lv, beta_lv
))
print(unify(Y_rv, Y_et))
# {~_334: RandomGeneratorSharedVariable(<Generator(PCG64) at 0x7FA66B45C4A0>), ~_335: TensorConstant{[]}, ~_336: TensorConstant{4}, ~_331: RandomGeneratorSharedVariable(<Generator(PCG64) at 0x7FA66B45C040>), ~_332: TensorConstant{[]}, ~_333: TensorConstant{11}, ~_329: alpha, ~_330: OpExpressionTuple((ExpressionTuple((<class 'aesara.tensor.elemwise....eDiv object at 0x7fa69b0fef20>, <frozendict {}>)), TensorConstant{1.0}, beta))}

Next steps before merging:

• Add a test that fails for reification;
• Make the test pass;
• Add a test that fails for unification on nested default-output apply nodes;
• Make the test pass;
• Add a test for reification on nested default-output apply nodes.
• Test on an example taken from AeMCMC