def assign_relation_splits(self, target_color, target_shape, referred_ref,
                             ref_object):
    """Assign the data example to corresponding splits."""
    ref_color = ref_object.color
    ref_shape = ref_object.shape
    splits = []
    # Experiment 1: visual generalization
    if ((target_color == "red" and target_shape == "square") or
        (ref_color == "red" and ref_shape == "square")):
      splits.append("visual")
    # Experiment 2: novel target and reference combination
    if (target_color == "green" and target_shape == "square" and
        ref_color == "blue" and ref_shape
        == "circle") or (target_color == "blue" and target_shape == "circle" and
                         ref_color == "green" and ref_shape == "square"):
      splits.append("relation")
    # Experiment 3: novel target
    if target_color == "yellow" and target_shape == "square":
      splits.append("referent")
    # Experimentt 4: novel relative position
    if self._vocabulary.translate_word(referred_ref["location"]) == "north of":
      splits.append("relative_position_1")
    # Experiment 5: novel relatiove postiion composition
    if self._vocabulary.translate_word(
        referred_ref["location"]) == "south west of":
      splits.append("relative_position_2")
    return splits
Assign the data example to corresponding splits.
language/gscan/data/dataset.py