Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Evaluation of etuplized RandomVariables fails for some distributions #1002

Open
rlouf opened this issue Jun 16, 2022 · 4 comments · May be fixed by #1036
Open

Evaluation of etuplized RandomVariables fails for some distributions #1002

rlouf opened this issue Jun 16, 2022 · 4 comments · May be fixed by #1036
Assignees

Comments

@rlouf
Copy link
Member

rlouf commented Jun 16, 2022

eval_if_etuple does not work when naively evaluating etuplized random variables. Let us consider the following model:

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

srng = at.random.RandomStream(0)
Y_rv = srng.normal(1, 1)

Y_et = etuplize(Y_rv)
Y_et
# e(e(<class 'aesara.tensor.random.basic.NormalRV'>, normal, 0, (0, 0), floatX, False), RandomGeneratorSharedVariable(<Generator(PCG64) at 0x7F25FFC56420>), TensorConstant{[]}, TensorConstant{11}, TensorConstant{1}, TensorConstant{1})

The original object is stored in the ExpressionTuple instance, so the following works:

Y_et.evaled_obj
# normal_rv{0, (0, 0), floatX, False}.out

We can force the ExpressionTuple to re-evaluate instead of returning the saved object; in this case evaluation fails:

Y_et._evaled_obj = Y_et.null
Y_et.evaled_obj
# TypeError: Too many positional arguments

The issue was discussed in aesara-devs/aemcmc#29 and prevents us to write miniKanren goals that represent two-way relations between expressions that contain random variables:

import aesara.tensor as at

from etuples import etuplize, etuple
from kanren import eq, lall, run, var

def normaleo(in_expr, out_expr):
    mu_lv, std_lv = var(), var()
    rng_lv, size_lv, dtype_lv = var(), var(), var()

    norm_in_lv = etuple(etuplize(at.random.normal), rng_lv, size_lv, dtype_lv, mu_lv, std_lv)
    norm_out_lv = etuple(etuplize(at.random.normal), rng_lv, size_lv, dtype_lv, mu_lv, std_lv)

    return lall(eq(in_expr, norm_in_lv), eq(out_expr, norm_out_lv))

srng = at.random.RandomStream(0)
Y_rv = srng.normal(1, 1)

y_lv = var()
(y_expr,) = run(1, y_lv, normaleo(Y_rv, y_lv))
y_expr.evaled_obj
# TypeError: Too many positional arguments

The reason, explained in this comment, is that the evaled_obj method performs op.__call__ when evaluating. However, __call__ does not dispatch to make_node systematically for =RandomVariable=s, hence the error.

Instead of dispatching op.__call__ to op.make_node for every random variable (which would be restrictive according to @brandonwillard), we can wrap RandomVariable Ops with a MakeNodeOp Op that dispatches __call__ to make_node:

from dataclasses import dataclass
from aesara.graph import Op, Variable
import aesara.tensor as at
from cons.core import ConsError, _car, _cdr
from etuples import etuple, etuplize


@dataclass
class MakeNodeOp:
    op: Op

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


def car_Variable(x):
    if x.owner:
        return MakeNodeOp(x.owner.op)
    else:
        raise ConsError("Not a cons pair.")


_car.add((Variable,), car_Variable)


def car_MakeNodeOp(x):
    return type(x)


_car.add((MakeNodeOp,), car_MakeNodeOp)


def cdr_MakeNodeOp(x):
    x_e = etuple(_car(x), x.op, evaled_obj=x)
    return x_e[1:]


_cdr.add((MakeNodeOp,), cdr_MakeNodeOp)


srng = at.random.RandomStream(0)
Y_rv = srng.normal(1, 1)
Y_et = etuplize(Y_rv)
Y_et
# e(e(<class '__main__.MakeNodeOp'>, e(<class 'aesara.tensor.random.basic.NormalRV'>, normal, 0, (0, 0), floatX, False)), RandomGeneratorSharedVariable(<Generator(PCG64) at 0x7FEB800BE260>), TensorConstant{[]}, TensorConstant{11}, TensorConstant{1}, TensorConstant{1})

Y_et._evaled_obj = Y_et.null
Y_et.evaled_obj
# normal_rv{0, (0, 0), floatX, False}(RandomGeneratorSharedVariable(<Generator(PCG64) at 0x7FEB800BE260>), TensorConstant{[]}, TensorConstant{11}, TensorConstant{1}, TensorConstant{1})

The downside to this approach is that defining kanren goals becomes more cumbersome as we need to wrap random variables with MakeNodeOp like:

norm_et = etuple(etuplize(at.random.normal), *Y_et[1:])
norm_et.evaled_obj
# TypeError: too many position arguments

I suggest instead to dispatch etuplize for RandomVariable, which handles both the etuplization of random variables and of RandomVariables ops:

from dataclasses import dataclass

from aesara.graph import Op, Variable
import aesara.tensor as at
from aesara.tensor.random.basic import RandomVariable

from etuples import etuplize, etuple


@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))

srng = at.random.RandomStream(0)
Y_rv = srng.normal(1, 1)
Y_et = etuplize(Y_rv)
Y_et
# e(e(<class 'aesara.tensor.random.basic.NormalRV'>, normal, 0, (0, 0), floatX, False), RandomGeneratorSharedVariable(<Generator(PCG64) at 0x7FF5DDA0E340>), TensorConstant{[]}, TensorConstant{11}, TensorConstant{1}, TensorConstant{1})

norm_et = etuple(etuplize(at.random.normal), *Y_et[1:])
norm_et.evaled_obj
# normal_rv{0, (0, 0), floatX, False}(RandomGeneratorSharedVariable(<Generator(PCG64) at 0x7FF5DDA0E340>), TensorConstant{[]}, TensorConstant{11}, TensorConstant{1}, TensorConstant{1})

And now the above kanren goal works both ways:

from kanren import run, lall, eq, var

def normaleo(in_expr, out_expr):
    mu_lv, std_lv = var(), var()
    rng_lv, size_lv, dtype_lv = var(), var(), var()

    norm_in_lv = etuple(etuplize(at.random.normal), rng_lv, size_lv, dtype_lv, mu_lv, std_lv)
    norm_out_lv = etuple(etuplize(at.random.normal), rng_lv, size_lv, dtype_lv, mu_lv, std_lv)

    return lall(eq(in_expr, norm_in_lv), eq(out_expr, norm_out_lv))


y_lv = var()
(y_expr,) = run(1, y_lv, normaleo(Y_rv, y_lv))
y_expr
# normal_rv{0, (0, 0), floatX, False}(RandomGeneratorSharedVariable(<Generator(PCG64) at 0x7FF5DDA0E340>), TensorConstant{[]}, TensorConstant{11}, TensorConstant{1}, TensorConstant{1})

y_lv = var()
(y_expr,) = run(1, y_lv, normaleo(y_lv, Y_rv))
y_expr
y_expr.evaled_obj
# normal_rv{0, (0, 0), floatX, False}(RandomGeneratorSharedVariable(<Generator(PCG64) at 0x7FF5DDA0E340>), TensorConstant{[]}, TensorConstant{11}, TensorConstant{1}, TensorConstant{1})

Note that the following code will still fail, however:

Y_et._evaled_obj = Y_et.null
Y_et.evaled_obj
# TypeError: Too many positional arguments

So we wil need a mix of the two approches. I will open a PR soon.

@rlouf rlouf added the bug Something isn't working label Jun 16, 2022
@rlouf rlouf self-assigned this Jun 16, 2022
@rlouf rlouf changed the title Reification of etuplized RandomVariables fail for some distributions Reification of etuplized RandomVariables fails for some distributions Jun 16, 2022
@rlouf rlouf changed the title Reification of etuplized RandomVariables fails for some distributions Evaluation of etuplized RandomVariables fails for some distributions Jun 29, 2022
@brandonwillard
Copy link
Member

Is your etuplize_random working as expected? I don't see a MakeRVNodeOp in the printed string output.

@rlouf
Copy link
Member Author

rlouf commented Jun 30, 2022

Yes, if you print norm_et defined above you get:

norm_et
# e(e(<class '__main__.MakeRVNodeOp'>, e(<class 'aesara.tensor.random.basic.NormalRV'>, normal, 0, (0, 0), floatX, False)), RandomGeneratorSharedVariable(<Generator(PCG64) at 0x7F65BBBE9A80>), TensorConstant{[]}, TensorConstant{11}, TensorConstant{1}, TensorConstant{1})

i.e. etuplize(at.random.normal) does wrap the RandomVariable pp with MakeRVNode. The problem is that MakeRVNodeOp does not appear in Y_et.

My understanding is that while we dispatch etuplize, it is the inner function etuplize_step that is applied recursively (see here) and this one is not dispatched. Is that expected?

The only way is through dispatching _car the way you did above. Still, it would be nice if the code above with etuplize_random just worked.

@rlouf
Copy link
Member Author

rlouf commented Jul 4, 2022

The following works:

from dataclasses import dataclass

from aesara.graph import Op, Variable
import aesara.tensor as at
from aesara.tensor.random.basic import RandomVariable

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

@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))

def car_RVariable(x):
    if x.owner:
        return MakeRVNodeOp(x.owner.op)
    else:
        raise ConsError("Not a cons pair.")

_car.add((Variable,), car_RVariable)

def car_MakeNodeOp(x):
    return type(x)

_car.add((MakeRVNodeOp,), car_MakeNodeOp)

def cdr_MakeNodeOp(x):
    x_e = etuple(_car(x), x.op, evaled_obj=x)
    return x_e[1:]


_cdr.add((MakeRVNodeOp,), cdr_MakeNodeOp)

srng = at.random.RandomStream(0)
Y_rv = srng.normal(1, 1)
Y_et = etuplize(Y_rv)
print(Y_et)
# e(e(<class '__main__.MakeRVNodeOp'>, e(<class 'aesara.tensor.random.basic.NormalRV'>, normal, 0, (0, 0), floatX, False)), RandomGeneratorSharedVariable(<Generator(PCG64) at 0x7F07F31FD2A0>), TensorConstant{[]}, TensorConstant{11}, TensorConstant{1}, TensorConstant{1})


norm_et = etuple(etuplize(at.random.normal), *Y_et[1:])
print(norm_et)
# e(e(<class '__main__.MakeRVNodeOp'>, e(<class 'aesara.tensor.random.basic.NormalRV'>, normal, 0, (0, 0), floatX, False)), RandomGeneratorSharedVariable(<Generator(PCG64) at 0x7F07F31FD2A0>), TensorConstant{[]}, TensorConstant{11}, TensorConstant{1}, TensorConstant{1})

which does feel a bit hacky; ideally dispatching etuplize would be enough to get the expected behavior. The current behavior, without knowing the internals of etuplize, is actually confusing.

@rlouf
Copy link
Member Author

rlouf commented Jul 4, 2022

Finally, just making sure this solves the motivating issue (aesara-devs/aemcmc#29) with running kanren goals forward and backward:

from kanren import eq, lall, run, var

def normaleo(in_expr, out_expr):
    mu_lv, std_lv = var(), var()
    rng_lv, size_lv, dtype_lv = var(), var(), var()

    norm_in_lv = etuple(etuplize(at.random.normal), rng_lv, size_lv, dtype_lv, mu_lv, std_lv)
    norm_out_lv = etuple(etuplize(at.random.normal), rng_lv, size_lv, dtype_lv, mu_lv, std_lv)

    return lall(eq(in_expr, norm_in_lv), eq(out_expr, norm_out_lv))

srng = at.random.RandomStream(0)
Y_rv = srng.normal(1, 1)
Y_et = etuplize(Y_rv)

y_lv = var()
(y_expr,) = run(1, y_lv, normaleo(Y_et, y_lv))
y_expr.evaled_obj
# normal_rv{0, (0, 0), floatX, False}(RandomGeneratorSharedVariable(<Generator(PCG64) at 0x7FC8E9B784A0>), TensorConstant{[]}, TensorConstant{11}, TensorConstant{1}, TensorConstant{1})


y_lv_rev = var()
(y_expr_rev,) = run(1, y_lv_rev, normaleo(y_lv_rev, Y_et))
y_expr_rev.evaled_obj
# normal_rv{0, (0, 0), floatX, False}(RandomGeneratorSharedVariable(<Generator(PCG64) at 0x7FC8E9B784A0>), TensorConstant{[]}, TensorConstant{11}, TensorConstant{1}, TensorConstant{1})

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
None yet
Development

Successfully merging a pull request may close this issue.

2 participants