import tensorflow as tf


  Parameters
  ----------
  x: input tensor.
  axis: axis along which to perform normalization.

  Returns
  -------
  A tensor.
  """
  if axis < 0:
    axis = axis % len(x.get_shape())
  return tf.nn.l2_normalize(x, dim=axis)


def categorical_crossentropy(output, target, from_logits=False):
  """Categorical crossentropy between an output tensor
  and a target tensor, where the target is a tensor of the same
  shape as the output.

  # TODO(rbharath): Should probably swap this over to tf mode.
  """
  # Note: tf.nn.softmax_cross_entropy_with_logits
  # expects logits, Keras expects probabilities.
