The transpose in the example had the same result type as its input, which is incorrect. PiperOrigin-RevId: 275186568