diff --git a/reasoning_gym/__init__.py b/reasoning_gym/__init__.py
index 6061a9c8..c1b93a3a 100644
--- a/reasoning_gym/__init__.py
+++ b/reasoning_gym/__init__.py
@@ -2,4 +2,7 @@
 Reasoning Gym - A library of procedural dataset generators for training reasoning models
 """
 
+from . import arithmetic
+
 __version__ = "0.1.0"
+__all__ = ["arithmetic"]
diff --git a/reasoning_gym/arithmetic/__init__.py b/reasoning_gym/arithmetic/__init__.py
index a2593e79..ad31381e 100644
--- a/reasoning_gym/arithmetic/__init__.py
+++ b/reasoning_gym/arithmetic/__init__.py
@@ -1,3 +1,9 @@
 from .chain_sum import ChainSum, ChainSumConfig
+from .basic_arithmetic import ArithmeticDataset, ArithmeticDatasetConfig
 
-__all__ = ["ChainSum", "ChainSumConfig"]
+__all__ = [
+    "ChainSum",
+    "ChainSumConfig",
+    "ArithmeticDataset",
+    "ArithmeticDatasetConfig"
+]