[tf.contrib.lookup] Clean up shape inference for lookup ops.

More of the shape inference can be done in C++-land, which may help grappler
do its thing.  Also fix a bug where keys.dim_size(0) was being requested even
when keys.dims() == 0
[this should probably lead to DCHECK failure, but doesn't seem to].

PiperOrigin-RevId: 208529368