From df2b8d2809c19a28361cdd884ad1b0f02fcc7587 Mon Sep 17 00:00:00 2001
From: "Andreas Koepf (aider)" <andreas.koepf@aleph-alpha.com>
Date: Sat, 25 Jan 2025 00:01:41 +0100
Subject: [PATCH] feat: Register chain_sum dataset with register_dataset
 function

---
 reasoning_gym/arithmetic/chain_sum.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/reasoning_gym/arithmetic/chain_sum.py b/reasoning_gym/arithmetic/chain_sum.py
index 6bfad5f9..9ea7ffc2 100644
--- a/reasoning_gym/arithmetic/chain_sum.py
+++ b/reasoning_gym/arithmetic/chain_sum.py
@@ -3,6 +3,7 @@ from dataclasses import dataclass
 from typing import Optional
 
 from ..dataset import ProceduralDataset
+from ..factory import register_dataset
 
 
 @dataclass
@@ -143,3 +144,6 @@ def chain_sum_dataset(
         size=size,
     )
     return ChainSum(config)
+
+# Register the dataset
+register_dataset("chain_sum", ChainSum, ChainSumConfig)