Add deterministic (& incremental aware) caching to the static language.

[femtomc]

Adds the ability to store deterministic data in StaticTrace via a new caching primitive.

The caching primitive is aware of incremental computation, and will only recompute the cached function in edit (and update) if changes to the arguments are indicated by UnknownChange.

When holding a StaticTrace, accessing the cache can be performed via StaticTrace.get_cache() (it's stored as a choice map, for now).

TODOs:

• Caching is not compositional currently -- meaning, there's not a compositional trace interface for it, so you can only use it if you have a StaticGenerativeFunction at the toplevel (so e.g. you get to hold a StaticTrace, and not some other trace (e.g. calling a StaticGenerativeFunction with caching inside of a VmapCombinator -- gives you a VmapTrace, and there's not yet a Trace interface which allows you to fetch the cache inside of the VmapTrace).

[femtomc]

@georgematheos See these tests for usage.

[georgematheos]

@femtomc this looks great!

Is there a way to add to this test something that checks if f is run multiple times? Would it work to have f increment a global counter variable, and check that this isn't incremented after update calls that shouldn't cause re-computation?

[femtomc]

Unfortunately JAX may run f abstractly multiple times during tracing -- that doesn't say much about the generated code.

For this one, we might just have to inspect the Jaxprs...

[femtomc]

Just checked the Jaxprs, it works

[femtomc]

Let's try and see to get it working in your codebase

[femtomc]

[codecov]

Codecov Report

Attention: Patch coverage is 76.76056% with 33 lines in your changes missing coverage. Please review.

Project coverage is 87.52%. Comparing base (799c83c) to head (76c5315).

Files with missing lines	Patch %	Lines
src/genjax/_src/generative_functions/static.py	75.38%	32 Missing ⚠️
...genjax/_src/core/generative/generative_function.py	87.50%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1357      +/-   ##
==========================================
- Coverage   87.94%   87.52%   -0.42%     
==========================================
  Files          53       53              
  Lines        3856     3992     +136     
==========================================
+ Hits         3391     3494     +103     
- Misses        465      498      +33     

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

[sritchie]

bound it on TypeVar("T", bound=Callable), then it'll pin the full type sig

[femtomc]

When I do this, I get an error -- that Callable is a variable and it's not allowed in type expressions.

[sritchie]

type this addr, please

[sritchie]

and docs for these new public methods (use Cursor to generate with an example!)

[MathieuHuot]

@georgematheos could you give some more context by providing the original example from ChiSight that motivated this (or a simplified version)?

My understanding from one of our discussions was that you wanted to visualise intermediate results within inference for debugging. If that's the main purpose, we want to tune for that purpose, or at least think about this use case some more.

[femtomc]

@MathieuHuot / @sritchie Let's not merge this in now -- I need to continue working on it to develop an example from George's Notion: https://www.notion.so/chi-mit/Incremental-computation-example-for-McCoy-10701b4b8cde803f9919d308cb0e331c?pvs=4

[femtomc]

Caching a little message I talked over with Mathieu:

# In the static language...
@genjax.gen
def model(x):
    v = submodel.vmap(in_axes=(0, ))(x) @ "v"
    x = cache("x", jnp.sum)(v) # this is new: incremental aware caching
    return some_other_model(x) @ "y"

tr # model trace in hand
# if you change "y", but don't change "v"
# jnp.sum @ "x" doesn't get recomputed.
#
# If we change "v", but don't change the argument x
# then v gets Diff(new_v, UnknownChange)
# the cache primitive sees this, and recomputes jnp.sum @ "x"
#
# Let's say we change "v" using IndexEditRequest(idx = 3, ...)
# NOTE: this is a TODO in the library (we have to implement IndexEditRequest for vmap).
# Now, the value for v -- what is it tagged with?
# 
# We should have: v gets Diff(new_v, IndexChange(idx, old_v_at_idx))
# And then what does cache do?
# 
# That's where we might want to allow users to specify rules.
@genjax.gen
def model(x):
    v = submodel.vmap(in_axes=(0, ))(x) @ "v"
    v = some_other_fn(v)
    x = cache("x", jnp.sum, custom_rule= lambda v: ...)(v) # this is new: incremental aware caching
    return some_other_model(x) @ "y"