tfp.experimental.nn.util.flatten_rightmost

View source on GitHub

Flatten rightmost dims.