Provide an accessor that supports `getitem`-like syntax to hierarchical choicemaps

[bzinberg]

The visualization (Base.show) of choicemaps shows hierarchy, but under the hood, the key space is flat. This means for example that if we have the following choicemap

│
├── :teacher
│   │
│   └── :name : "Bob"
│
└── :students
    │
    ├── 1
    │   │
    │   ├── :name : "Britney"
    │   │
    │   └── :gradepoint : 0.31
    │
    ├── 2
    │   │
    │   ├── :name : "Beatrice"
    │   │
    │   └── :gradepoint : 0.97
    │
    └── 3
        │
        ├── :name : "Bill"
        │
        └── :gradepoint : 0.18

we unfortunately do not get a convenient iterator syntax like

for (studentId, student) in Gen.nested_view(trace)[:students]
  println("Name: $(student[:name]) -- Gradepoint avg: $(student[:gradepoint])")
end

It would be nice to provide an accessor that allows this syntax.

Upon first thought, it seems like this accessor could be implemented by constructing a nested dictionary of Refs, which has cost linear in the size of the choicemap.

(Of course, we could also implement this by restructuring choicemaps themselves to have this hierarchy built-in under the hood, but that design decision has a lot more implications and I'm sure there were good reasons to have a flat key space.)

[marcoct]

Choice maps are indeed hierarchical, and you can achieve what you are asking for:

choices = Gen.get_choices(trace)
for (studentId, submap) in Gen.get_submaps_shallow(choices)
   name = submap[:name]
   gradepoint = submap[:gradepoint]
   # ...
end

Note that the trace Base.getindex syntactic sugar (trace[:students]) only supports getting individual addresses. It could be potentially extended to return the submap when the address refers to a submap and not a single random choice. In fact, it probably should.

Also, choice maps could easily be made iterable, where you iterate over the top-level values and submaps. The reason I did not already make them iterable was that I thought users will in general want to treat submaps and values differently, and returning then both from the iterator could be confusing (therefore I opted for separate function get_submaps_shallow and get_values_shallow). However, maybe it would still be valuable to just provide the iterator, for the common case when the the top-level entries are either all submaps or all individual random choices.

Together these two changes would allow you to write:

for (studentId, submap) in trace[:students]
   name = submap[:name]
   gradepoint = submap[:gradepoint]
   # ...
end

which is not exactly your code snippet, since the iteration over the submap would return key-value pairs and not just the values as assumed in your code snippet.

Choice maps are documented here: https://probcomp.github.io/Gen/dev/ref/choice_maps/

[bzinberg]

Ah whoops -- I think my incorrect recollection of flat key space was based on the syntax choices[:a => :b] which I had seen in the tutorials. My bad.

It could be potentially extended to return the submap when the address refers to a submap and not a single random choice. In fact, it probably should.

Yes, that would be great! Personally, I would a priori expect the choicemap to behave similarly to a dict-of-dicts, including the syntactic sugars for iteration. I didn't realize the subtlety you point out regarding the possible desire to treat values and submaps differently -- though I think if we expose a "chiocemap-of-choicemaps" (either through an accessor function or up front with the values-only and submaps-only modes being the ones that have accessor functions), that would be fine.

I'm happy to take this issue, since it's central to a data layout / code ergonomics problem Vikash and I were discussing earlier (https://github.com/probcomp/Cora/issues/23) which I need to fix this weekend in some way or another.

Thanks!

[bzinberg]

EDIT: This makes no sense -- see correction in https://github.com/probcomp/Gen/issues/106#issuecomment-502414029

I'm currently thinking something like

for (studentId, (name, gradepoint)) in Gen.nested_view(trace)[:students]
  # ...
end

could be good. And/or maybe

for (studentId, (name, gradepoint)) in trace.choices_r[:students]
  # ...
end

[marcoct]

@bzinberg I think this is a pretty core part of Gen's language design, so I want to get it right, and I want to make sure the decision-making isn't rushed for a short-term deadline (i.e. this weekend). I'll take a look at the Cora issue and see if there are other shorter-term solutions.

[bzinberg]

No worries, I already have a short-term solution as mentioned in https://github.com/probcomp/Cora/issues/23#issuecomment-502407025 so that timeline is not blocked on this one. Happy to discuss it fully here and put it into code when we're ready.

[bzinberg]

BTW, If you have an example of a use case where having the choicemap support nested getattrs [Edit: should have said nested getitems] would produce confusing results, I'd be interested to hear it. I get the point abstractly, but don't have a concrete example in mind.

[marcoct]

I think that two separate changes could help with Gen's usability:

1. Allow Base.getindex on a Trace to return a submap, and not just a random choice. (See https://github.com/probcomp/Gen/blob/master/src/gen_fn_interface.jl#L71-L73 and https://github.com/probcomp/Gen/blob/master/src/choice_map.jl#L75.

2. Make choice maps iterable, where we iterate over both submaps and choices (i.e. where the iterator is an iterator over pairs of key and value where the value is either a random choice value or a choicemap).

However, a potential issue with both of these (and the reason why they aren't already present) is that it is technically valid to have random choices that are themselves choice maps (nothing in Gen prevents this). Therefore, there would be a potential ambiguity, where it's not clear whether a value is a random choice or a choice map. It might not matter in the end, but it requires some thought. It also relates to Gen's semantics.

BTW, If you have an example of a use case where having the choicemap support nested getattrs would produce confusing results, I'd be interested to hear it. I get the point abstractly, but don't have a concrete example in mind.

I'm not sure what you mean by choice map supporting nested getattrs. Can you give an example?

[bzinberg]

Right yeah, I did have some concern that as you say, the representation of a choicemap as a dict-of-dicts is not "faithful" (i.e. injective), since the values could themselves be choicemaps. It seems similar to the question of whether one is allowed to use a Pair as the tag for a @traced expression -- is that possible?

I should have set nested getitems -- I meant expressions of the form e.g. trace[:students][1] as an alternative to trace[:students => 1].

[marcoct]

Currently Base.getindex (which is what provides the square-bracket indexing notation) for choice maps is implemented as:

Base.getindex(choices::ChoiceMap, addr) = get_value(choices, addr)

(where get_value means 'get a single random choice').

To get a submap, users need to use get_submap, e.g. :

function get_submap(choices::DynamicChoiceMap, addr)
    if haskey(choices.internal_nodes, addr)
        choices.internal_nodes[addr]
    elseif haskey(choices.leaf_nodes, addr)
        throw(KeyError(addr))
    else
        EmptyChoiceMap()
    end 
end

[marcoct]

It seems similar to the question of whether one is allowed to use a Pair as the tag for a @traced expression -- is that possible?

Yes, you can use a Pair as the tag for a @trace expression (https://probcomp.github.io/Gen/dev/ref/modeling/#Composite-addresses-1).

I think the ambiguity regarding a choice map returned by getindex is actually resolved by whether a Distribution or a GenerativeFunction was responsible for that address. Distributions produce single random choices, and GenerativeFunctions produce choice maps. Whenever the choice map is consumed by a generative function interface method (e.g. as a constraint passed into Gen.generate), the method will already know whether a distribution or a generative function is responsible for the address.

But I'll want to think this through a bit more. The distinction between generative functions and distributions is a core design decision. The short version is that distributions are kind of like 'unboxed' generative functions, and are a type of performance and interface optimization.

[bzinberg]

Yes, you can use a Pair as the tag for a @trace expression (https://probcomp.github.io/Gen/dev/ref/modeling/#Composite-addresses-1).

Oh right, of course you can, I even do it in the code snippet https://github.com/probcomp/Cora/issues/23#issue-456576212.

Whenever the choice map is consumed by a generative function interface method (e.g. as a constraint passed into Gen.generate), the method will already know whether a distribution or a generative function is responsible for the address.

But I'll want to think this through a bit more. The distinction between generative functions and distributions is a core design decision.

Ah, got it.

[alex-lew]

I think the ambiguity regarding a choice map returned by getindex is actually resolved by whether a Distribution or a GenerativeFunction was responsible for that address. Distributions produce single random choices, and GenerativeFunctions produce choice maps. Whenever the choice map is consumed by a generative function interface method (e.g. as a constraint passed into Gen.generate), the method will already know whether a distribution or a generative function is responsible for the address.

Unless I'm missing something, nothing prevents a distribution from returning a choicemap (a valid Julia value), which seems to be the heart of the ambiguity.

EDIT: Or maybe this is what you're saying -- that either a distribution or generative function can produce a choice map, but consumers can disambiguate based on which was the case.

[marcoct]

Right, but I think that the ambiguity might not matter in practice.

The information about whether a given submap in a choice map is a random choice or a choice map need not be maintained, and this ambiguity need never be resolved.

On Sat, Jun 15, 2019, 8:19 PM Alex Lew notifications@github.com wrote:

I think the ambiguity regarding a choice map returned by getindex is actually resolved by whether a Distribution or a GenerativeFunction was responsible for that address. Distributions produce single random choices, and GenerativeFunctions produce choice maps. Whenever the choice map is consumed by a generative function interface method (e.g. as a constraint passed into Gen.generate), the method will already know whether a distribution or a generative function is responsible for the address.

Unless I'm missing something, nothing prevents a distribution from returning a choicemap (a valid Julia value), which seems to be the heart of the ambiguity.

— You are receiving this because you commented. Reply to this email directly, view it on GitHub https://github.com/probcomp/Gen/issues/106?email_source=notifications&email_token=AAMKDNF6YL4MSQZS3OBDJOLP2WBH5A5CNFSM4HYPZMAKYY3PNVWWK3TUL52HS4DFVREXG43VMVBW63LNMVXHJKTDN5WW2ZLOORPWSZGODXZCOLQ#issuecomment-502409006, or mute the thread https://github.com/notifications/unsubscribe-auth/AAMKDNETAMQFUEYUTSHPEOLP2WBH5ANCNFSM4HYPZMAA .

[bzinberg]

It seems to me that we could provide the accessors in https://github.com/probcomp/Gen/issues/106#issuecomment-502406020 regardless of what the core design of choicemaps is? Or does exposing that interface inherently point the user in the wrong direction?

[alex-lew]

Yes, I see. In Metaprob, there is no type distinction between a distribution and a generative function; a distribution is just a generative function that traces a value at its "root" address. In Gen, they are different types.

If Gen reliably makes this distinction (so that any consumer of a choice map could always distinguish the two cases), then choicemaps might as well be dicts-of-dicts. This would make it impossible to implement generic (i.e., not model-specific) choice-map-consuming functions that care about the distinction -- i.e., a pretty printer that knows how to format submaps vs. choices, or an iterator of all random choices made in a choice map. (Though maybe these are not particularly useful.)

@bzinberg I'm not sure I understand what's happening in

for (studentId, (name, gradepoint)) in Gen.nested_view(trace)[:students]

At the top level, Gen.nested_view seems to be a map-like object, but each student-level submap is just a tuple? What order do the fields go in to the tuple? If I write

for (kwd, x) in Gen.nested_view(trace)

do I get two elements, one with kwd == :teacher and the 1-element tuple ("Bob",), and one with kwd == :students and the nested Tuple ((1, ("Brittany", 0.31)), (2, (..., ...)), ...)?

[marcoct]

If Gen reliably makes this distinction (so that any consumer of a choice map could always distinguish the two cases), then choicemaps might as well be dicts-of-dicts. This would make it impossible to implement generic (i.e., not model-specific) choice-map-consuming functions that care about the distinction -- i.e., a pretty printer that knows how to format submaps vs. choices, or an iterator of all random choices made in a choice map. (Though maybe these are not particularly useful.)

Yeah, I think that the distinction between the two cases can probably be elided by the choice map interface, which would result in simplified implementations and user interfaces for choice maps. I would want to think through the implications for JIT-generation in the static modeling language (I don't think there would be issues, but not 100% sure). Also a lot of code would be touched (and in general, simplified I think).

[marcoct]

It seems to me that we could provide the accessors in #106 (comment) regardless of what the core design of choicemaps is? Or does exposing that interface inherently point the user in the wrong direction?

I don't understand the accessors you suggested. What is the specific functionality that you are trying to provide? What is gained beyond this existing already-implemented version:

choices = Gen.get_choices(trace)
for (studentId, submap) in Gen.get_submaps_shallow(choices)
   name = submap[:name]
   gradepoint = submap[:gradepoint]
   # ...
end

or its potential replacement (which I'm in favor of, and which would involve eliding the difference between choice-map-as-choice-map and choice-map-as-random-choice).

for (studentId, submap) in trace[:students] # or get_choices(trace)[:students]
   name = submap[:name]
   gradepoint = submap[:gradepoint]
   # ...
end

It seems like the main functional difference in what you're suggesting is the ability to pattern match to extract the values of random choices? (e.g. (name, gradepoint)) instead of just iterate over key-value pairs? In that case, how would the order of these values be defined?

[marcoct]

Also, we could also reconsider requiring users to call get_choices to get something that implements the choice map interface. The trace could directly implement this interface (i.e. every trace would itself be a choice map). I resisted this in the past, because of the potential confusion between these two very similar types of data (choice maps and traces) that are importantly distinct in Gen (choice maps just represent some values of random choices, whereas traces represent complete executions of specific generative functions along with cached metadata including but not limited to arguments and return value and various probability scores). But the get_choices is starting to seem more silly and unnecessary.

[alex-lew]

Yeah, I think that the distinction between the two cases can probably be elided by the choice map interface, which would result in simplified implementations and user interfaces for choice maps. I would want to think through the implications for JIT-generation in the static modeling language (I don't think there would be issues, but not 100% sure). Also a lot of code would be touched (and in general, simplified I think).

Makes sense!

[bzinberg]

@alex-lew

@bzinberg I'm not sure I understand what's happening

You're totally right, my suggestion in https://github.com/probcomp/Gen/issues/106#issuecomment-502406020 didn't make sense, and I realized that when I just tried to implement it. I should have said:

for (studentId, student) in Gen.nested_view(trace)[:students]
  println("Name: $(student[:name]) -- Gradepoint avg: $(student[:gradepoint])")
end

The iterator on this view of a choicemap provides iteration over (key, value) pairs, and when the value is a submap, the thing we see in place of that submap is itself a view that supports the same iteration and getindex semantics.

@marcoct

or its potential replacement (which I'm in favor of, and which would involve eliding the difference between choice-map-as-choice-map and choice-map-as-random-choice)

Yes, I think that's exactly what I'm going for. I want to be able to iterate over a hierarchical choice map via the same syntax as I would use to iterate over a dict of things-some-of-which-are-themselves-dicts. (Notice that my erroneous suggestion doesn't work on a dict of dicts :slightly_smiling_face:)

I think having that pun available would help decrease a beginning user's cognitive load when thinking about choicemaps. And I think it would also help me achieve a good balance of simple code and decreased data structure complexity as in https://github.com/probcomp/Cora/issues/23.

[bzinberg]

Here is an implementation of the above https://github.com/probcomp/Gen/issues/106#issuecomment-502414029:

https://gist.github.com/bzinberg/89253afc30de37a628b7c9fbdcdeb751/a3b5f11dc8c32729864410be4885a2887664dcac

[bzinberg]

Here's a slightly fancier implementation that gets rid of the duplicated code:

https://gist.github.com/bzinberg/89253afc30de37a628b7c9fbdcdeb751

[bzinberg]

In https://github.com/probcomp/Cora/issues/23#issuecomment-502470207 (and the PR it references, https://github.com/probcomp/Cora/pull/25) I give a concrete example of these code ergonomics at play.

[bzinberg]

Closing this issue and refreshing it as https://github.com/probcomp/Gen/issues/156.