tfp.experimental.distributions.marginal_fns.ps.non_negative_axis

Make (possibly negatively indexed) axis argument non-negative.