open-thought/reasoning-gym
1
0
Fork 0
mirror of https://github.com/open-thought/reasoning-gym.git synced 2026-04-27 17:23:19 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

feat: Add task augmentation functions mirror, inverse, and identity to arc_1d.py

Browse source
This commit is contained in:
Andreas Koepf (aider) 2025-02-02 22:42:21 +01:00
parent 01cc239746
commit 84e4f1c5bc
2 changed files with 43 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

25
reasoning_gym/cognition/arc_1d.py
View file

@ -1030,6 +1030,31 @@ def task_gravity_weighted_colors(rng: Random, size: int) -> Optional[Dict[str, L
return {"input": question, "output": answer}
def task_mirror(task_result: Optional[Dict[str, List[int]]]) -> Optional[Dict[str, List[int]]]:
"""Mirror the input and output arrays of a task result."""
if task_result is None:
return None
return {
"input": list(reversed(task_result["input"])),
"output": list(reversed(task_result["output"]))
}
def task_inverse(task_result: Optional[Dict[str, List[int]]]) -> Optional[Dict[str, List[int]]]:
"""Swap the input and output arrays of a task result."""
if task_result is None:
return None
return {
"input": task_result["output"],
"output": task_result["input"]
}
def task_identity(task_result: Optional[Dict[str, List[int]]]) -> Optional[Dict[str, List[int]]]:
"""Return the task result unchanged."""
return task_result
def task_color_left_half_blocks(rng: Random, size: int) -> Optional[Dict[str, List[int]]]:
"""Generate a task where left half of blocks are colored differently."""
pos = 0