Support scan markov with history > 1

[fehiepsi]

Resolves #702.

Tasks:

• Revise log_density implementation to incorporate this feature.
• Add hmm example with history > 1
• Add test to verify the consistency of density, output carry, ys with unrolled scan

[fehiepsi]

@eb8680 Could you help me review this PR? Most of the changes are generalized from history=1 -> history>1. There is a complicated issue to return the correct shape for the last carry: assume history=2, then we run step 0, step 1 outside of lax.scan, then use lax.scan for steps 2 -> (T - 1). Under enumeration, we have 3 different shapes for carry, but we can only interpret those shapes using the input carry and output carry at step 2. So one of them is missing.

We can run step 0, step 1, step 2, then using lax.scan for steps 3 -> (T - 1). This way we can interpret all carry shapes (using infos at steps 2 and 3) but need to unroll history * 2 steps (in general). So unless it is really needed, I would like to unroll the number of steps as small as possible. :) Anyway, we already added a note previously that: there should not have any site outside of scan depend on the first output of scan (the last carry value)

[eb8680]

For clarification, after changing this to for i in range(2 * history) what other changes would be necessary to get the correct final carry_shape?

[eb8680]

@fehiepsi I think I understand this issue, thanks for explaining. How much compilation overhead does unrolling 2 * history steps add in the HMM examples? Unrolling history steps in exchange for incorrect output carry shapes is a very subtle optimization that only produces constant-time savings in program length, and may cause maintenance difficulties down the line, so it's important to be sure that it's worth it.

[fehiepsi]

Thinking about it again, I guess it won't take much compiling time because history is usually small (1 or 2). Let me address the issue in this PR.

[eb8680]

LGTM after latest changes.

[fehiepsi]

Thanks for reviewing, @eb8680! Looking like there is a bug when I unrolled more steps. Looking into it...

[fehiepsi]

we only want to do this right before getting the trace.

[fehiepsi]

The usage of enum(), markov() here still gives correct log density but does not recycle the dimensions at this step. To resolve this, I have moved block right before the generator markov(unroll_steps + 1, history) and remove markov handler here.

[fehiepsi]

This is the minimum number of unroll_steps to get a correct last carry shape. But it is safe to use any number greater than this number.

[fehiepsi]

@eb8680 Though the previous implementation gives correct log density, the carry shapes are not recycled with

with markov(history=history):
    for i in range(unroll_steps):
        transition_fn()
    lax.scan(with markov(): transition_fn())

logic. I have revised it to be

for i in markov(range(unroll_steps + 1)):
    if i < unroll_steps:
        transition_fn()
    else:
        lax.scan(transition_fn())

Previously, this issue is not detected due to a typo in the test, where I computed expected and actual values on the same model:

    expected_x_prev, expected_x_curr = enum(config_enumerate(model))()
    actual_x_prev, actual_x_curr = enum(config_enumerate(model))()  # typo: should be fun_model

[eb8680]

LGTM. Now that sarkka_bilmes_product is actually being used, I guess we should probably try to move away from the P prefix naming scheme upstream in sarkka_bilmes_product to a safer interface so no non-Markov names from user code with P in them get incorrectly mangled. Maybe the simplest thing to do would be to make the prefixes longer (e.g. _PREV_ instead of P).

[fehiepsi]

no non-Markov names from user code with P in them get incorrectly mangled

Good point! _PREV_ should work well here.