import tensorflow as tf

  else:
    # need broadcasting
    target_shape = []
    for axis in range(get_ndim(x)):
      if axis in reduction_axes:
        target_shape.append(1)
      else:
        target_shape.append(tf.shape(x)[axis])
    target_shape = stack(target_shape)

    broadcast_mean = tf.reshape(mean, target_shape)
    broadcast_var = tf.reshape(var, target_shape)
    broadcast_gamma = tf.reshape(gamma, target_shape)
    broadcast_beta = tf.reshape(beta, target_shape)
    normed = tf.nn.batch_normalization(x, broadcast_mean, broadcast_var,
                                       broadcast_beta, broadcast_gamma, epsilon)
  return normed, mean, var


def ones(shape, dtype=None, name=None):
  """Instantiates an all-ones tensor variable and returns it.

  Parameters
  ----------
