//  [python]
def render_constellations(pred_points,
                          capsule_num,
                          canvas_size,
                          gt_points=None,
                          n_caps=2,
                          gt_presence=None,
                          pred_presence=None,
                          caps_presence_prob=None):
  """Renderes predicted and ground-truth points as gaussian blobs.

  Args:
    pred_points: [B, m, 2].
    capsule_num: [B, m] tensor indicating which capsule the corresponding point
      comes from. Plots from different capsules are plotted with different
      colors. Currently supported values: {0, 1, ..., 11}.
    canvas_size: tuple of ints
    gt_points: [B, k, 2]; plots ground-truth points if present.
    n_caps: integer, number of capsules.
    gt_presence: [B, k] binary tensor.
    pred_presence: [B, m] binary tensor.
    caps_presence_prob: [B, m], a tensor of presence probabilities for caps.

  Returns:
    [B, *canvas_size] tensor with plotted points
  """

  # convert coords to be in [0, side_length]
  pred_points = denormalize_coords(pred_points, canvas_size, rounded=True)

  # render predicted points
  batch_size, n_points = pred_points.shape[:2].as_list()
  capsule_num = tf.to_float(tf.one_hot(capsule_num, depth=n_caps))
  capsule_num = tf.reshape(capsule_num, [batch_size, n_points, 1, 1, n_caps, 1])

  color = tf.convert_to_tensor(_COLORS[:n_caps])
  color = tf.reshape(color, [1, 1, 1, 1, n_caps, 3]) * capsule_num
  color = tf.reduce_sum(color, -2)
  color = tf.squeeze(tf.squeeze(color, 3), 2)

  colored = render_by_scatter(canvas_size, pred_points, color, pred_presence)

  # Prepare a vertical separator between predicted and gt points.
  # Separator is composed of all supported colors and also serves as
  # a legend.
  # [b, h, w, 3]
  n_colors = _COLORS.shape[0]
  sep = tf.reshape(tf.convert_to_tensor(_COLORS), [1, 1, n_colors, 3])
  n_tiles = int(colored.shape[2]) // n_colors
  sep = snt.TileByDim([0, 1, 3], [batch_size, 3, n_tiles])(sep)
  sep = tf.reshape(sep, [batch_size, 3, n_tiles * n_colors, 3])

  pad = int(colored.shape[2]) - n_colors * n_tiles
  pad, r = pad // 2, pad % 2

  if caps_presence_prob is not None:
    n_caps = int(caps_presence_prob.shape[1])
    prob_pads = ([0, 0], [0, n_colors - n_caps])
    caps_presence_prob = tf.pad(caps_presence_prob, prob_pads)
    zeros = tf.zeros([batch_size, 3, n_colors, n_tiles, 3], dtype=tf.float32)

    shape = [batch_size, 1, n_colors, 1, 1]
    caps_presence_prob = tf.reshape(caps_presence_prob, shape)

    prob_vals = snt.MergeDims(2, 2)(caps_presence_prob + zeros)
    sep = tf.concat([sep, tf.ones_like(sep[:, :1]), prob_vals], 1)

  sep = tf.pad(sep, [(0, 0), (1, 1), (pad, pad + r), (0, 0)],
               constant_values=1.)

  # render gt points
  if gt_points is not None:
    gt_points = denormalize_coords(gt_points, canvas_size, rounded=True)

    gt_rendered = render_by_scatter(canvas_size, gt_points, colors=None,
                                    gt_presence=gt_presence)

    colored = tf.where(tf.cast(colored, bool), colored, gt_rendered)
    colored = tf.concat([gt_rendered, sep, colored], 1)

  res = tf.clip_by_value(colored, 0., 1.)
  return res