k_squeeze function

Removes a 1-dimension from the tensor at index axis.