def sum_rightmost_ndims_preserving_shape(x, ndims):
  """
  """
  x = tf.convert_to_tensor(value=x)
  if x.shape.ndims is not None:
    axes = tf.range(x.shape.ndims - ndims, x.shape.ndims)
  else:
    axes = tf.range(tf.rank(x) - ndims, tf.rank(x))
  return tf.reduce_sum(input_tensor=x, axis=axes)