Add more conjugates

[xjing76]

Trying to add more conjugates.

will be addition with current commit Uniform -Pareto.

@brandonwillard @rlouf Please let me know which of the conjugates should be more prioritized.

[xjing76]

I am having most trouble right here, where I am not sure how to get the at.max(obs, x_m) @brandonwillard

[brandonwillard]

Does etuple(at.max, observed_val, x_lv) work?

[brandonwillard]

@brandonwillard @rlouf Please let me know which of the conjugates should be more prioritized.

It would probably be best for us to get all the exponential family models out of the way first; otherwise, I have no particular order/preference.

[brandonwillard]

Currently, the test graph looks like this:

>>> aesara.dprint(Y_rv)
pareto_rv{0, (0, 0), floatX, False}.1 [id A]
 |RandomGeneratorSharedVariable(<Generator(PCG64) at 0x7FC4162BFF20>) [id B]
 |TensorConstant{[]} [id C]
 |TensorConstant{11} [id D]
 |MaxAndArgmax{axis=()}.0 [id E] 'max'
 | |y [id F]
 |Elemwise{add,no_inplace} [id G]
   |k [id H]
   |n [id I]

and that MaxAndArgmax Op isn't same as the at.math.max used in the etuple graph. at.math.max is a function that constructs a MaxAndArgmax Op and uses it to further construct a graph for the max of its argument. In other words, we need an etuple form/"pattern" that matches the types of graphs output by the helper function at.math.max.

Often the easiest way to find etuple forms for the graphs constructed by helper functions is to etuplize said graphs and spot their generalities.
For example:

>>> from etuples import etuplize
>>> etuplize(at.math.max(at.vector("x")))
e(e(aesara.tensor.math.MaxAndArgmax, (0,)), x)
>>> etuplize(at.math.max(at.matrix("x")))
e(e(aesara.tensor.math.MaxAndArgmax, (0, 1)), x)

As we can see, the axis property in the MaxAndArgmax Op will change according to the dimensions of the input (i.e. it computes the max across all dimensions), so we don't want to use a very specific value for the matching form. Instead, we can use another logic variable in place of those values.

Here's a general testing setup for that part of the problem:

import aesara
import aesara.tensor as at

from etuples import etuplize, etuple


srng = at.random.RandomStream(0)

k_tt = at.scalar("k")
y_vv = at.iscalar("y")
n_tt = at.scalar("n")

Y_rv = srng.pareto(at.max(y_vv), k_tt + n_tt)

# This is what we need to match/unify:
etuplize(Y_rv)
# e(
#     e(aesara.tensor.random.basic.ParetoRV, 'pareto', 0, (0, 0), 'floatX', False),
#     RandomGeneratorSharedVariable(<Generator(PCG64) at 0x7FA1F9B3D9E0>),
#     TensorConstant{[]},
#     TensorConstant{11},
#     e(e(aesara.tensor.math.MaxAndArgmax, ()), y),
#     e(
#         e(
#             aesara.tensor.elemwise.Elemwise,
#             <aesara.scalar.basic.Add at 0x7fa1fd3823d0>,
#             <frozendict {}>),
#         k,
#         n))

from unification import var
from kanren import run, eq
from aesara.tensor.math import MaxAndArgmax


observed_val = var()
axis_lv = var()
new_x_et = etuple(etuple(MaxAndArgmax, axis_lv), observed_val)

k_lv, n_lv = var(), var()
new_k_et = etuple(etuplize(at.add), k_lv, n_lv)

theta_rng_lv = var()
theta_size_lv = var()
theta_type_idx_lv = var()
theta_posterior_et = etuple(
    etuplize(at.random.pareto),
    theta_rng_lv,
    theta_size_lv,
    theta_type_idx_lv,
    new_x_et,
    new_k_et,
)


run(0, (new_x_et, new_k_et), eq(Y_rv, theta_posterior_et))
# ((e(e(aesara.tensor.math.MaxAndArgmax, ()), y),
#   e(
#       e(
#           aesara.tensor.elemwise.Elemwise,
#           <aesara.scalar.basic.Add at 0x7fa1fd3823d0>,
#           <frozendict {}>),
#       k,
#       n)),)

[codecov]

Codecov Report

❌ Patch coverage is 50.00000% with 28 lines in your changes missing coverage. Please review.
✅ Project coverage is 94.75%. Comparing base (64b0e50) to head (ae0af70).
⚠️ Report is 19 commits behind head on main.

Files with missing lines	Patch %	Lines
aemcmc/conjugates.py	50.00%	28 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #113      +/-   ##
==========================================
- Coverage   98.50%   94.75%   -3.76%     
==========================================
  Files          10       11       +1     
  Lines         737      820      +83     
  Branches       63       69       +6     
==========================================
+ Hits          726      777      +51     
- Misses          4       33      +29     
- Partials        7       10       +3     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[xjing76]

Sorry, have not been on top of this. (After Covid and rounds of strep ), I will try to add at least one conjugate a day from now on.

[brandonwillard]

Sorry, have not been on top of this. (After Covid and rounds of strep ), I will try to add at least one conjugate a day from now on.

Ouch! No problem at all; we really appreciate what you've done already. I still need to provide a solution for the default updates on the re-used RNG states anyway, so no rush.

N.B. This comment describes the issue.

[xjing76]

Great! I just realized that the coverage testing suggesting that I have entire local_uniform_pareto_posterior and local_beta_bernoulli_posterior missed. I probably got something wrong.. https://app.codecov.io/gh/aesara-devs/aemcmc/pull/113/blob/aemcmc/conjugates.py#L332