The Categorical distribution just returns indices on the domain. For what I’m trying to do I need to bind actual domain values to the trace, so I can condition on those values later. For example if my domain is the integers {3, 6, 9}, I want a sample statement to bind those values to the trace instead of {0, 1, 2}.

So I wrote the following discrete uniform version of Categorical.

well for one, your log_prob is completely ignoring the input. why do you need to wrap the categorical distribution instead of using the sample values as indices? something like (from your example):

log_prob is completely ignoring the input
It’s a uniform distribution, should they should all have the same log_prob, so I was trying to pass it the first value in the domain, 0.

why do you need to wrap the categorical distribution instead of using the sample values as indices?
Looks like your solution works, but I’m working a canonical problem where I want to be picky in what names get used in the program trace.

Using your code, I made the following changes, does this look okay?