Remove the data_sample model constructor argument

HYRAX_GUIDE.md

@@ -162,6 +162,8 @@ Hyrax discovers components through three registries:
 - Decorator: `@hyrax_model`
 - Models must inherit from `torch.nn.Module` and implement `__init__`, `forward`,
   `train_batch`, and `prepare_inputs`.
+- Hyrax constructs a model as `Model(config)` and passes nothing else. Models size
+  themselves from config; the built-in image models read `[model.<Name>].input_shape`.
 - The decorator wires up save/load, optimizer, and criterion handling.
 - Every value in the metric dict returned by `train_batch`, `validate_batch`, or `test_batch`
   must be a per-sample mean over the batch (not a sum or count); epoch metrics weight each

docs/model_class_reference.rst

@@ -12,10 +12,8 @@ How Hyrax uses your model class
 
 Hyrax does this sequence:
 
-1. Builds one sample batch from your dataset.
-2. Calls ``YourModel.prepare_inputs(sample_batch)``.
-3. Creates your model with ``YourModel(config, data_sample=prepared_sample)``.
-4. During runs, calls:
+1. Creates your model with ``YourModel(config)``.
+2. During runs, calls:
 
    * ``train_batch`` for training
    * ``infer_batch`` for inference
@@ -33,7 +31,7 @@ Your class must have these:
    find it by name in the :doc:`configuration </configuration>`).
 3. Implement:
 
-   * ``__init__(self, config, data_sample=None)``
+   * ``__init__(self, config)``
    * ``forward(self, ...)`` (required by PyTorch modules)
    * ``infer_batch(self, batch)``
    * ``train_batch(self, batch)``
@@ -44,15 +42,15 @@ Your class must have these:
 Method-by-method guide
 ----------------------
 
-``__init__(self, config, data_sample=None)``
-^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+``__init__(self, config)``
+^^^^^^^^^^^^^^^^^^^^^^^^^^
 
 What to do in this method:
 
 1. Save ``config`` to ``self.config``.
-2. Read model-specific settings from config.
+2. Read model-specific settings from config, including any input sizes your
+   layers need. Hyrax passes nothing but ``config``, so shapes come from it.
 3. Build your layers.
-4. Optionally use ``data_sample`` to infer shapes.
 
 The model-specific config values are typically defined in your external library
 config layout. See :doc:`/external_library_package` for packaging/config setup.
@@ -61,26 +59,23 @@ Example (read settings from config):
 
 .. code-block:: python
 
-   def __init__(self, config, data_sample=None):
+   def __init__(self, config):
        super().__init__()
        self.config = config
 
+       in_dim = config["my_package"]["MyModel"]["input_size"]
        hidden = config["my_package"]["MyModel"]["hidden_size"]
        out_dim = config["my_package"]["MyModel"]["num_classes"]
 
-       in_dim = 128
-       if data_sample is not None:
-           x, _ = data_sample
-           in_dim = x.shape[-1]
-
        self.fc1 = nn.Linear(in_dim, hidden)
        self.fc2 = nn.Linear(hidden, out_dim)
 
 To set those config values in a notebook, one simple pattern is:
 
 .. code-block:: python
 
-   h.config.setdefault("my_package", {}).setdefault("MyModel", {})["hidden_size"] = 64
+   h.config["my_package"]["MyModel"]["input_size"] = 128
+   h.config["my_package"]["MyModel"]["hidden_size"] = 64
    h.config["my_package"]["MyModel"]["num_classes"] = 10
 
 PyTorch reference for building ``nn.Module`` layers:
@@ -366,7 +361,7 @@ when you use it always gives you the run that is actually underway.
 
    @hyrax_model
    class MyModel(nn.Module):
-       def __init__(self, config, data_sample=None):
+       def __init__(self, config):
            super().__init__()
            self.config = config
            self.activations = []
@@ -409,7 +404,7 @@ Complete minimal class
 
    @hyrax_model
    class MyModel(nn.Module):
-       def __init__(self, config, data_sample=None):
+       def __init__(self, config):
            super().__init__()
            self.config = config
            self.fc = nn.Linear(128, 10)

docs/notebooks/custom_training_metrics.ipynb

@@ -57,7 +57,7 @@
     "class MetricsDemoAutoencoder(nn.Module):\n",
     "    \"\"\"Tiny autoencoder that demonstrates custom training metrics.\"\"\"\n",
     "\n",
-    "    def __init__(self, config, data_sample=None):\n",
+    "    def __init__(self, config):\n",
     "        super().__init__()\n",
     "        self.encoder = nn.Linear(4, 2)\n",
     "        self.decoder = nn.Linear(2, 4)\n",
@@ -135,7 +135,7 @@
     "class EpochMetricsDemoAutoencoder(nn.Module):\n",
     "    \"\"\"Autoencoder that reports both batch and epoch-level metrics.\"\"\"\n",
     "\n",
-    "    def __init__(self, config, data_sample=None):\n",
+    "    def __init__(self, config):\n",
     "        super().__init__()\n",
     "        self.encoder = nn.Linear(4, 2)\n",
     "        self.decoder = nn.Linear(2, 4)\n",

docs/notebooks/list_available_models_and_datasets.ipynb

@@ -95,7 +95,7 @@
     "class AAA_MyCustomModel(nn.Module):\n",
     "    \"\"\"A simple toy model for demonstration purposes.\"\"\"\n",
     "\n",
-    "    def __init__(self, config, data_sample=None):\n",
+    "    def __init__(self, config):\n",
     "        super().__init__()\n",
     "        self.linear = nn.Linear(10, 2)\n",
     "\n",
@@ -166,7 +166,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.13.12"
+   "version": "3.13.15"
   }
  },
  "nbformat": 4,

docs/pre_executed/external_model_class.ipynb

@@ -115,7 +115,7 @@
     "\n",
     "The `@hyrax_model` decorator registers the class by name so Hyrax can find it when you set `model.name`.\n",
     "\n",
-    "The constructor receives `config` (the Hyrax config dict) and an optional `data_sample` — a single batch produced by `prepare_inputs`. The sample lets you infer input dimensions automatically instead of hardcoding them."
+    "The constructor receives `config`, the Hyrax config dict. Any sizes your layers need (image shape, number of classes, sequence length) come from a config entry you add for your model, as shown below with `h.config[\"model\"][\"ImageOnlyCNN\"]`."
    ]
   },
   {
@@ -132,12 +132,11 @@
     "\n",
     "@hyrax_model  # registers the class so h.set_config(\"model.name\", \"ImageOnlyCNN\") works\n",
     "class ImageOnlyCNN(nn.Module):\n",
-    "    def __init__(self, config, data_sample=None):\n",
+    "    def __init__(self, config):\n",
     "        super().__init__()\n",
     "        self.config = config\n",
-    "        # Unpack the sample to derive input dimensions automatically\n",
-    "        image_sample, _ = data_sample\n",
-    "        c, h, w = image_sample.shape[1:]  # channels, height, width\n",
+    "        # Input dimensions come from the config entry set in the training cell below\n",
+    "        c, h, w = config[\"model\"][\"ImageOnlyCNN\"][\"input_shape\"]  # channels, height, width\n",
     "        self.net = nn.Sequential(\n",
     "            nn.Conv2d(c, 8, kernel_size=3, padding=1),\n",
     "            nn.ReLU(),\n",
@@ -201,6 +200,8 @@
     "\n",
     "h = Hyrax()\n",
     "h.set_config(\"model.name\", \"ImageOnlyCNN\")\n",
+    "# The image shape the model is built for; the dataset above produces 3x32x32 images.\n",
+    "h.set_config(\"model.ImageOnlyCNN.input_shape\", [3, 32, 32])\n",
     "h.set_config(\"train.epochs\", 1)\n",
     "# These are set to default values, uncomment to change.\n",
     "# h.set_config(\"data_loader.batch_size\", 512)\n",
@@ -255,14 +256,13 @@
    "source": [
     "@hyrax_model\n",
     "class ImageAndLightCurveModel(nn.Module):\n",
-    "    def __init__(self, config, data_sample=None):\n",
+    "    def __init__(self, config):\n",
     "        super().__init__()\n",
     "        self.config = config\n",
-    "        image, light_curve, light_curve_mask, _ = data_sample\n",
+    "        model_config = config[\"model\"][\"ImageAndLightCurveModel\"]\n",
     "\n",
-    "        c, h, w = image.shape[1:]\n",
-    "        lc_dim = light_curve.shape[1]  # padded sequence length\n",
-    "        print(f\"LC_DIM={lc_dim}\")\n",
+    "        c, h, w = model_config[\"input_shape\"]\n",
+    "        lc_dim = model_config[\"light_curve_length\"]  # padded sequence length\n",
     "\n",
     "        # Image branch: simple CNN → 32-d embedding\n",
     "        self.image_encoder = nn.Sequential(\n",
@@ -332,6 +332,9 @@
    "source": [
     "h2 = Hyrax()\n",
     "h2.set_config(\"model.name\", \"ImageAndLightCurveModel\")\n",
+    "# The sizes the model is built for: 3x32x32 images, and light curves padded to max_len (100).\n",
+    "h2.set_config(\"model.ImageAndLightCurveModel.input_shape\", [3, 32, 32])\n",
+    "h2.set_config(\"model.ImageAndLightCurveModel.light_curve_length\", 100)\n",
     "h2.set_config(\"train.epochs\", 1)\n",
     "# These are set to default values, uncomment to change.\n",
     "# h.set_config(\"data_loader.batch_size\", 512)\n",
@@ -403,7 +406,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "kbmod",
+   "display_name": "hyrax",
    "language": "python",
    "name": "python3"
   },
@@ -417,7 +420,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.12.11"
+   "version": "3.13.15"
   },
   "nbsphinx": {
    "execute": "never"

docs/pre_executed/lsdb_stream_dataset.ipynb

@@ -277,16 +277,19 @@
    },
    "outputs": [],
    "source": [
+    "# The number of catalog columns prepare_inputs stacks into each input row.\n",
+    "N_FEATURES = 2\n",
+    "\n",
+    "\n",
     "@hyrax_model\n",
     "class CatalogAutoencoder(nn.Module):\n",
     "    \"\"\"A tiny dense autoencoder over a handful of catalog columns.\"\"\"\n",
     "\n",
-    "    def __init__(self, config, data_sample=None):\n",
+    "    def __init__(self, config):\n",
     "        super().__init__()\n",
     "        self.config = config\n",
-    "        n_features = data_sample.shape[1]\n",
-    "        self.encoder = nn.Sequential(nn.Linear(n_features, 8), nn.GELU(), nn.Linear(8, 2))\n",
-    "        self.decoder = nn.Sequential(nn.Linear(2, 8), nn.GELU(), nn.Linear(8, n_features))\n",
+    "        self.encoder = nn.Sequential(nn.Linear(N_FEATURES, 8), nn.GELU(), nn.Linear(8, 2))\n",
+    "        self.decoder = nn.Sequential(nn.Linear(2, 8), nn.GELU(), nn.Linear(8, N_FEATURES))\n",
     "\n",
     "    def forward(self, batch):\n",
     "        return self.encoder(batch)\n",
@@ -304,7 +307,7 @@
     "\n",
     "    @staticmethod\n",
     "    def prepare_inputs(data_dict) -> tuple:\n",
-    "        \"\"\"Stack the requested catalog columns into one (batch, n_features) array.\"\"\"\n",
+    "        \"\"\"Stack the requested catalog columns into one (batch, N_FEATURES) array.\"\"\"\n",
     "        import numpy as np\n",
     "\n",
     "        data = data_dict[\"data\"]\n",
@@ -651,7 +654,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "hyrax-lsdb-dataset",
+   "display_name": "hyrax",
    "language": "python",
    "name": "python3"
   },
@@ -665,7 +668,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.13.14"
+   "version": "3.13.15"
   }
  },
  "nbformat": 4,

docs/pre_executed/supervised_lightcurve_transients.ipynb

@@ -514,7 +514,7 @@
     "class LightCurveGRU(nn.Module):\n",
     "    \"\"\"A bidirectional GRU for variable-length multi-band light-curve classification.\"\"\"\n",
     "\n",
-    "    def __init__(self, config, data_sample=None):\n",
+    "    def __init__(self, config):\n",
     "        super().__init__()\n",
     "        self.config = config\n",
     "        model_config = config[\"model\"][\"LightCurveGRU\"]\n",
@@ -1623,7 +1623,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "kbmod",
+   "display_name": "hyrax",
    "language": "python",
    "name": "python3"
   },
@@ -1637,7 +1637,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.12.11"
+   "version": "3.13.15"
   },
   "widgets": {
    "application/vnd.jupyter.widget-state+json": {