Unsafe variable naming in JointDistributionCoroutine

[chrism0dwk]

The JointDistributionCoroutine class is a fantastic way of constructing a joint probability distribution, particularly with random variables that depend on each other. The class allows you to name your variable, and outputs a joint sample as a StructTuple (which AFAICS is just a namedtuple under the hood).

However, you have to be careful how you name your variables. Consider the trivial example with a variable inadvertently named count:

import tensorflow_probability as tfp
tfd = tfp.distributions

@tfd.JointDistributionCoroutine
def model():
    yield tfd.Uniform(low=0., high=1., name="count")

For this model, calling model.sample() apparently returns a StructTuple with a count field as expected:

>>> sample = model.sample()
>>> sample
StructTuple(
  count=<tf.Tensor: shape=(), dtype=float32, numpy=0.22241592>
)

Great! But what happens if I try to access count?

>>> sample.count
<function structtuple.<locals>.StructTuple.count(value, /)>

Oh dear. It looks like we get the bound tuple.count method back instead. For the same reason, this wrecks any tfd.JointDistributionCoroutine methods that depends on accessing StructTuple attributes. e.g.

>>> model.log_prob(sample)
TypeError: Cannot convert the argument `type_value`: <built-in method count of StructTuple object at 0x78954f1e0140> to a TensorFlow DType.

A fix might just be to document the reserved names in the Name resolution paragraph in the doc for this class. But I wonder if there's something deeper we could do?

Chris

[ColCarroll]

Thanks, Chris -- looks like the StructTuple is a custom class that subclasses the built-in tuple, whose only two (public) attributes are count and index. I guess it is unfortunate that I could imagine using both of these in a graphical model!

I think we should just add those two to this check on validating names. If someone wants a variable named __gt__, they deserve whatever happens down the line 😀

What do you think?

[chrism0dwk]

@ColCarroll sounds a pragmatic way forward. I've suggested a solution in #1828. See what you think...

[SiegeLordEx]

The goal of structtuple was to mimic namedtuple's behavior, which neither forbids these names, nor balks at using them. It sounds like we didn't quite mimic its behavior correctly.

[SiegeLordEx]

I'm guessing we need to use __getattribute__ rather than __getattr__ so that the field names correctly shadow the parent class members.

[ColCarroll]

#1828