from tensorflow.python.framework import ops

  last = input_shape[-1].value
  if last is not None and k is not None and last < k:
    raise ValueError("input.shape %s must have last dimension >= k = %d" %
                     (input_shape, k))
  output_shape = input_shape[:-1].concatenate([k])
  return [output_shape, output_shape]


@ops.RegisterShape("BatchNormWithGlobalNormalization")
def _BatchNormShape(op):
  """Shape function for BatchNormWithGlobalNormalization op."""
  input_shape = op.inputs[0].get_shape().with_rank(4)
  mean_shape = op.inputs[1].get_shape().with_rank(1)
  var_shape = op.inputs[2].get_shape().with_rank(1)
  beta_shape = op.inputs[3].get_shape().with_rank(1)
  gamma_shape = op.inputs[4].get_shape().with_rank(1)
