Probability: Conditional Chain Bijector

[bkazi]

The Chain bijector currently has support to pass kwargs to the chained bijectors within it, but it is currently a subclass of Bijector and ConditionalBijector so it's not possible to actually supply any kwargs.
Is this by design?

[srvasude]

Hi,
This sounds like a bug.

More generally, we are currently thinking about merging Bijector and Conditional Bijector (as well as TransformedDistribution and ConditionalTransformedDistribution, and Distribution and ConditionalDistribution) in our internal issue tracker. As a side effect, this issue should go away.

[hartikainen]

I raised similar issue on the tensorflow side (https://github.com/tensorflow/tensorflow/issues/21543) some time ago. I think it's currently possible to define:

class ConditionalChain(tfb.ConditionalBijector, tfb.Chain):
    pass

and then use that instead of tfb.Chain.

[jvdillon]

Ill soon be sending out a change which makes kwargs first class citizens, ie, every Bijector (and distribution) will be "Conditional."

[brianwa84]

The merging has happened but there's still a question of how to know which args to pass to which bijector in the chain.

[jvdillon]

Would this be a viable answer?

"""
At least wrt to Chain my hypothesized solution was that all kwarg forwarding would be managed by a user supplied function which splits out the dict elements. For Chain, I envisioned this:

def _default_kwargs_split_fn(bijectors, kwargs):
"""Default kwargs dict getter."""
return kwargs.get("bijector_kwargs", [{}]*len(bijectors))

as an alternative to this:
line 162
tensorflow_probability/python/distributions/transformed_distribution.py

Unfortunately I have not thought deeply about __call__. It is possibly we cannot support kwargs for it, although I remain hopeful that the dependency injection approach would work there too.
"""