Add shape checks to GPU TridiagonalMatMul.

When given invalid shapes, the GPU TridiagonalMatMul op could read invalid or uninitialized GPU memory.

PiperOrigin-RevId: 401775483
Change-Id: Ib5500aeb8225e50d4ce790b06d2c34751f544ad8