Fix to issue raised in https://github.com/danielward27/flowjax/issues/147, that noticed excessive masking from conditioning variables to output variables. Specifically, connections between the conditioning variable and rank 0 output variables were incorrectly masked out.
Fix to issue raised in https://github.com/danielward27/flowjax/issues/147, that noticed excessive masking from conditioning variables to output variables. Specifically, connections between the conditioning variable and rank 0 output variables were incorrectly masked out.