last dim.

Also fixed the shape inference code for SpaceToBatchNDOp in order to correctly
handle the given inputs, that contain dynamic dimensions.

PiperOrigin-RevId: 340226617
Change-Id: I054f4ba794b5eb206dba69a67a80eb25514b4a27