tfp.experimental.distributions.marginal_fns.ps.dimension_size

Equivalent to shape(x)[idx], but robust to partially-known shapes.