This would do the opposite of expand_dims and is sometimes required when external simulators provide an additional trailing dimension, e.g. in the parameters. Do I see it correctly that there currently is no way to remove those using the adapter (i.e., turn a shape (batch_size, 2, 1) into (batch_size, 2))?
To make it deterministic, I think we would require an axis argument explicitly, just as in expand_dims.