The syntax for constructing a chain is very confusing and inconsistent with most other libraries which build chains of functions.
For example:
bijector = distrax.Chain([
distrax.ScalarAffine(shift = shift, scale = scale),
distrax.Tanh()
])
first applies Tanh and then applies ScalarAffine. Intuitively it should be the other way around.
The syntax for constructing a chain is very confusing and inconsistent with most other libraries which build chains of functions.
For example:
first applies Tanh and then applies ScalarAffine. Intuitively it should be the other way around.