tfp.experimental.distributions.marginal_fns.ps.shape_slice

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