Update calling of expand_dims with axis

This fix updates calling of `expand_dims` with `dim -> axis`
as the `dim=` in `tf.expand_dims` has been deprecated and was
generating unnecessary warnings.

Signed-off-by: Yong Tang <yong.tang.github@outlook.com>