matrix transpose code in Python