Fix bugs in lifting dimensions with inner dimension broadcasting