message: Fix warnings in tf.distributions.Categorical
  In tf.distributions.Categorical dimension was used with argmax.
  As dimension has been deprecated this generates a warning.
  This fix fixes the warning by changing to axis.
diff: @@ -311,7 +311,7 @@ class Categorical(distribution.Distribution):
  nn_ops.log_softmax(self.logits) * self.probs, axis=-1)
  def _mode(self):
  - ret = math_ops.argmax(self.logits, dimension=self._batch_rank)
  + ret = math_ops.argmax(self.logits, axis=self._batch_rank)
  ret = math_ops.cast(ret, self.dtype)
  ret.set_shape(self.batch_shape)
  return ret
  