Use `PartialTensorShape` instead of `TensorShape`.

`TensorShape` constructor throws a CHECK-fail if shape is partial/overflows which the other doesn't. We are only determining the number of elements in the shape and partial shape should be used as it returns negative number when needed.

PiperOrigin-RevId: 409205384
Change-Id: Ia56542ff9ec758f2c9ffc7e4dcc9fa7eecd86e7b