From 1b7fa53b16b6d41941f36d2a0ede528a189ba921 Mon Sep 17 00:00:00 2001
From: xraymemory <me.anzuoni@gmail.com>
Date: Mon, 20 Jul 2026 12:33:25 -0400
Subject: [PATCH 1/2] refactor(models): remove unused x_init

---
 AGENTS.md                                     |  5 +-
 src/sampleworks/models/boltz/wrapper.py       | 86 +++----------------
 src/sampleworks/models/protenix/wrapper.py    | 44 +---------
 src/sampleworks/models/protocol.py            | 13 ---
 src/sampleworks/models/rf3/wrapper.py         | 30 +------
 .../utils/guidance_script_utils.py            |  6 +-
 .../integration/test_no_guidance_geometry.py  |  4 +-
 .../integration/test_pipeline_integration.py  | 13 +--
 tests/mocks/model_wrappers.py                 | 16 ++--
 tests/models/boltz/test_boltz_wrapper.py      | 72 ++++++----------
 tests/models/test_model_wrapper_protocol.py   | 59 ++++---------
 tests/models/test_rf3_atom_ordering.py        |  6 +-
 tests/models/test_rf3_chiral_features.py      |  6 --
 13 files changed, 83 insertions(+), 277 deletions(-)

diff --git a/AGENTS.md b/AGENTS.md
index dae7d5d2..b60e44d8 100644
--- a/AGENTS.md
+++ b/AGENTS.md
@@ -397,8 +397,7 @@ class MyModelWrapper:
         atom_array = structure["asym_unit"]
         # ... model-specific featurization ...
         conditioning = my_model_features(atom_array)
-        x_init = torch.zeros(n_atoms, 3, device=self.device)
-        return GenerativeModelInput(x_init=x_init, conditioning=conditioning)
+        return GenerativeModelInput(conditioning=conditioning)
 
     def step(
         self,
@@ -418,7 +417,7 @@ class MyModelWrapper:
         shape: tuple[int, ...] | None = None,
     ) -> Tensor:
         """Sample from the prior (typically Gaussian noise)."""
-        n_atoms = features.x_init.shape[-2] if features else shape[-2]
+        n_atoms = features.conditioning.num_atoms if features else shape[-2]
         return torch.randn(batch_size, n_atoms, 3, device=self.device)
 ```
 
diff --git a/src/sampleworks/models/boltz/wrapper.py b/src/sampleworks/models/boltz/wrapper.py
index 273c3476..a02f846c 100644
--- a/src/sampleworks/models/boltz/wrapper.py
+++ b/src/sampleworks/models/boltz/wrapper.py
@@ -314,8 +314,6 @@ class BoltzConfig:
         Output directory for processed Boltz files.
     num_workers : int
         Number of parallel workers for preprocessing.
-    ensemble_size : int
-        Number of samples to generate (batch dimension of x_init).
     recycling_steps : int
         Number of recycling steps to perform during featurization Pairformer pass.
 
@@ -323,7 +321,6 @@ class BoltzConfig:
 
     out_dir: str | Path | None = None
     num_workers: int = 0
-    ensemble_size: int = 1
     recycling_steps: int = 3
 
 
@@ -332,7 +329,6 @@ def process_structure_for_boltz(
     *,
     out_dir: str | Path | None = None,
     num_workers: int = 0,
-    ensemble_size: int = 1,
     recycling_steps: int | None = 3,
 ) -> dict:
     """Annotate an Atomworks structure with Boltz-specific configuration.
@@ -346,8 +342,6 @@ def process_structure_for_boltz(
         Defaults to structure metadata ID or "boltz_output".
     num_workers : int
         Number of parallel workers for preprocessing.
-    ensemble_size : int
-        Number of samples to generate (batch dimension of x_init).
     recycling_steps : int | None
         Number of recycling steps to perform during featurization Pairformer pass.
         Will set to 3 if None.
@@ -372,7 +366,6 @@ def process_structure_for_boltz(
     config = BoltzConfig(
         out_dir=out_dir or structure.get("metadata", {}).get("id", "boltz_output"),
         num_workers=num_workers,
-        ensemble_size=ensemble_size,
         recycling_steps=recycling_steps,
     )
     return {**structure, "_boltz_config": config}
@@ -649,7 +642,7 @@ def _setup_data_module(
     def featurize(self, structure: dict) -> GenerativeModelInput[BoltzConditioning]:
         """From an Atomworks structure, calculate Boltz-2 input features.
 
-        Runs Pairformer pass and initializes x_init from prior distribution.
+        Runs the Pairformer pass to produce conditioning features.
 
         NOTE: Has side effect of creating Boltz input YAML and initial processed files with
         create_boltz_input_from_structure() and the data module setup.
@@ -665,7 +658,7 @@ def featurize(self, structure: dict) -> GenerativeModelInput[BoltzConditioning]:
         Returns
         -------
         GenerativeModelInput[BoltzConditioning]
-            Model input with x_init and Pairformer conditioning.
+            Model input with Pairformer conditioning.
         """
         self.cached_representations.clear()
 
@@ -681,8 +674,6 @@ def featurize(self, structure: dict) -> GenerativeModelInput[BoltzConditioning]:
             else base_dir
         )
         num_workers = config.num_workers
-        ensemble_size = config.ensemble_size
-
         input_path = create_boltz_input_from_structure(  # side effect of creating Boltz input YAML
             structure,
             out_dir,
@@ -721,32 +712,7 @@ def featurize(self, structure: dict) -> GenerativeModelInput[BoltzConditioning]:
             model_atom_array=model_atom_array,
         )
 
-        # Use atom count from Boltz featurization in atom_pad_mask (via conditioning) to ensure
-        # consistency between x_init shape and atom_pad_mask used in step(). Note: x_init=None is
-        # a temporary placeholder; initialize_from_prior derives shape from conditioning.
-        feats = pairformer_out["feats"]
-        atom_mask = cast(Tensor, feats.get("atom_pad_mask"))
-        num_atoms = int(atom_mask.sum())
-
-        # x_init should be the reference coordinates for alignment purposes.
-        if true_atom_array is not None and len(true_atom_array) == num_atoms:
-            x_init = torch.tensor(true_atom_array.coord, device=self.device, dtype=torch.float32)
-            x_init = match_batch(x_init.unsqueeze(0), target_batch_size=ensemble_size).clone()
-        else:
-            # TODO: The temp features situation is not ideal and I think we can do better probably
-            # not sure exactly what the best way to handle x_init is, should define it a bit better
-            # most likely
-            logger.warning(
-                "True structure not available or atom count mismatch; initializing "
-                "x_init from prior. This means align_to_input will not work properly,"
-                " and reward functions dependent on this won't be accurate."
-            )
-            temp_features = GenerativeModelInput[BoltzConditioning](
-                x_init=cast(Any, None), conditioning=conditioning
-            )
-            x_init = self.initialize_from_prior(batch_size=ensemble_size, features=temp_features)
-
-        return GenerativeModelInput(x_init=x_init, conditioning=conditioning)
+        return GenerativeModelInput(conditioning=conditioning)
 
     def _pairformer_pass(
         self, features: dict[str, Any], recycling_steps: int = 3
@@ -951,8 +917,7 @@ def initialize_from_prior(
         if shape is not None:
             if len(shape) != 2 or shape[1] != 3:
                 raise ValueError("shape must be of the form (num_atoms, 3)")
-            x_init = torch.randn((batch_size, *shape), device=self.device)
-            return x_init
+            return torch.randn((batch_size, *shape), device=self.device)
         if features is None or features.conditioning is None:
             raise ValueError("Either features or shape must be provided to initialize_from_prior()")
 
@@ -963,9 +928,7 @@ def initialize_from_prior(
 
         num_atoms = int(atom_mask.sum())
 
-        x_init = torch.randn((batch_size, num_atoms, 3), device=self.device)
-
-        return x_init
+        return torch.randn((batch_size, num_atoms, 3), device=self.device)
 
 
 class Boltz1Wrapper:
@@ -1109,7 +1072,7 @@ def _setup_data_module(
     def featurize(self, structure: dict) -> GenerativeModelInput[BoltzConditioning]:
         """From an Atomworks structure, calculate Boltz-1 input features.
 
-        Runs Pairformer pass and initializes x_init from prior distribution.
+        Runs the Pairformer pass to produce conditioning features.
 
         NOTE: Has side effect of creating Boltz input YAML and initial processed files with
         create_boltz_input_from_structure() and the data module setup.
@@ -1125,7 +1088,7 @@ def featurize(self, structure: dict) -> GenerativeModelInput[BoltzConditioning]:
         Returns
         -------
         GenerativeModelInput[BoltzConditioning]
-            Model input with x_init and Pairformer conditioning.
+            Model input with Pairformer conditioning.
         """
         self.cached_representations.clear()
 
@@ -1141,8 +1104,6 @@ def featurize(self, structure: dict) -> GenerativeModelInput[BoltzConditioning]:
             else base_dir
         )
         num_workers = config.num_workers
-        ensemble_size = config.ensemble_size
-
         input_path = create_boltz_input_from_structure(
             structure,
             out_dir,
@@ -1180,31 +1141,7 @@ def featurize(self, structure: dict) -> GenerativeModelInput[BoltzConditioning]:
             model_atom_array=model_atom_array,
         )
 
-        # x_init should be the reference coordinates for alignment purposes.
-        # The guidance scalers call initialize_from_prior() separately for starting noise.
-        # Use atom count from Boltz featurization to ensure shape consistency with model.
-        feats = pairformer_out["feats"]
-        atom_mask = cast(Tensor, feats.get("atom_pad_mask"))
-        num_atoms = int(atom_mask.sum())
-
-        if true_atom_array is not None and len(true_atom_array) == num_atoms:
-            x_init = torch.tensor(true_atom_array.coord, device=self.device, dtype=torch.float32)
-            x_init = match_batch(x_init.unsqueeze(0), target_batch_size=ensemble_size).clone()
-        else:
-            # TODO: The temp features situation is not ideal and I think we can do better probably
-            # not sure exactly what the best way to handle x_init is, should define it a bit better
-            # most likely
-            logger.warning(
-                "True structure not available or atom count mismatch; initializing "
-                "x_init from prior. This means align_to_input will not work properly,"
-                " and reward functions dependent on this won't be accurate."
-            )
-            temp_features = GenerativeModelInput[BoltzConditioning](
-                x_init=cast(Any, None), conditioning=conditioning
-            )
-            x_init = self.initialize_from_prior(batch_size=ensemble_size, features=temp_features)
-
-        return GenerativeModelInput(x_init=x_init, conditioning=conditioning)
+        return GenerativeModelInput(conditioning=conditioning)
 
     def step(
         self,
@@ -1312,8 +1249,7 @@ def initialize_from_prior(
         if shape is not None:
             if len(shape) != 2 or shape[1] != 3:
                 raise ValueError("shape must be of the form (num_atoms, 3)")
-            x_init = torch.randn((batch_size, *shape), device=self.device)
-            return x_init
+            return torch.randn((batch_size, *shape), device=self.device)
         if features is None or features.conditioning is None:
             raise ValueError("Either features or shape must be provided to initialize_from_prior()")
 
@@ -1324,9 +1260,7 @@ def initialize_from_prior(
 
         num_atoms = int(atom_mask.sum())
 
-        x_init = torch.randn((batch_size, num_atoms, 3), device=self.device)
-
-        return x_init
+        return torch.randn((batch_size, num_atoms, 3), device=self.device)
 
     def _pairformer_pass(
         self, features: dict[str, Any], recycling_steps: int = 3
diff --git a/src/sampleworks/models/protenix/wrapper.py b/src/sampleworks/models/protenix/wrapper.py
index ef94c537..fa307a12 100644
--- a/src/sampleworks/models/protenix/wrapper.py
+++ b/src/sampleworks/models/protenix/wrapper.py
@@ -58,7 +58,7 @@ class ProtenixConditioning:
         Raw feature dict for diffusion module.
     num_atoms : int
         Number of atoms in the Protenix model's representation. This is the authoritative
-        atom count for the diffusion module and should be used for x_init shape.
+        atom count for the diffusion module and prior initialization.
     pair_z : Tensor | None
         Cached pair representation for diffusion conditioning.
     p_lm : Tensor | None
@@ -93,8 +93,6 @@ class ProtenixConfig:
     ----------
     out_dir : str | Path | None
         Output directory for intermediate JSON file.
-    ensemble_size : int
-        Number of ensemble members to generate.
     recycling_steps : int | None
         Number of recycling steps to perform. If None, uses model default.
     use_msa : bool
@@ -104,7 +102,6 @@ class ProtenixConfig:
     """
 
     out_dir: str | Path | None = None
-    ensemble_size: int = 1
     recycling_steps: int | None = None
     use_msa: bool = True
     enable_diffusion_shared_vars_cache: bool = True
@@ -114,7 +111,6 @@ def annotate_structure_for_protenix(
     structure: dict,
     *,
     out_dir: str | Path | None = None,
-    ensemble_size: int = 1,
     recycling_steps: int | None = None,
     use_msa: bool = True,
     enable_diffusion_shared_vars_cache: bool = True,
@@ -128,8 +124,6 @@ def annotate_structure_for_protenix(
     out_dir : str | Path | None
         Output directory for intermediate files.
         Defaults to structure metadata ID or "protenix_output".
-    ensemble_size : int
-        Number of ensemble members to generate.
     recycling_steps : int | None
         Number of recycling steps to perform. If None, uses model default.
     use_msa : bool
@@ -144,7 +138,6 @@ def annotate_structure_for_protenix(
     """
     config = ProtenixConfig(
         out_dir=out_dir or structure.get("metadata", {}).get("id", "protenix_output"),
-        ensemble_size=ensemble_size,
         recycling_steps=recycling_steps,
         use_msa=use_msa,
         enable_diffusion_shared_vars_cache=enable_diffusion_shared_vars_cache,
@@ -327,7 +320,7 @@ def __init__(
     def featurize(self, structure: dict) -> GenerativeModelInput[ProtenixConditioning]:
         """From an Atomworks structure, calculate Protenix input features.
 
-        Runs Pairformer pass and initializes x_init from prior distribution.
+        Runs the Pairformer pass to produce conditioning features.
 
         Parameters
         ----------
@@ -340,14 +333,13 @@ def featurize(self, structure: dict) -> GenerativeModelInput[ProtenixConditionin
         Returns
         -------
         GenerativeModelInput[ProtenixConditioning]
-            Model input with x_init and Pairformer conditioning.
+            Model input with Pairformer conditioning.
         """
         config = structure.get("_protenix_config", ProtenixConfig())
         if isinstance(config, dict):
             config = ProtenixConfig(**config)
 
         out_dir = config.out_dir or structure.get("metadata", {}).get("id", "protenix_output")
-        ensemble_size = config.ensemble_size
         use_msa = config.use_msa
         recycling_steps = config.recycling_steps
         enable_diffusion_shared_vars_cache = config.enable_diffusion_shared_vars_cache
@@ -505,35 +497,7 @@ def featurize(self, structure: dict) -> GenerativeModelInput[ProtenixConditionin
             model_atom_array=model_aa,
         )
 
-        # x_init should be the reference coordinates for alignment purposes.
-        # Must match num_atoms_protenix so downstream samplers/scalers see
-        # consistent shapes with the model representation.
-        if "asym_unit" in structure:
-            n_input = len(atom_array)
-            if n_input != num_atoms_protenix:
-                logger.warning(
-                    f"Atom-count mismatch: atom_array has {n_input} atoms, "
-                    f"atom_array_protenix has {num_atoms_protenix} atoms. "
-                    "Using atom_array_protenix coords for x_init to match model "
-                    "atom count.",
-                )
-                x_init = torch.as_tensor(
-                    atom_array_protenix.coord, device=self.device, dtype=torch.float32
-                )
-            else:
-                x_init = torch.as_tensor(atom_array.coord, device=self.device, dtype=torch.float32)
-            x_init = match_batch(x_init.unsqueeze(0), target_batch_size=ensemble_size).clone()
-        else:
-            logger.warning(
-                "True structure not available, so initializing "
-                "x_init from prior. This means align_to_input will not work properly,"
-                " and reward functions dependent on this won't be accurate."
-            )
-            x_init = self.initialize_from_prior(
-                batch_size=ensemble_size, shape=(num_atoms_protenix, 3)
-            )
-
-        return GenerativeModelInput(x_init=x_init, conditioning=conditioning)
+        return GenerativeModelInput(conditioning=conditioning)
 
     def _pairformer_pass(
         self, features: dict[str, Any], grad_needed: bool = False, **kwargs
diff --git a/src/sampleworks/models/protocol.py b/src/sampleworks/models/protocol.py
index e24411a4..fd3df734 100644
--- a/src/sampleworks/models/protocol.py
+++ b/src/sampleworks/models/protocol.py
@@ -26,25 +26,12 @@
 class GenerativeModelInput(Generic[C]):  # noqa: UP046 (for Python 3.11 compatibility)
     """Container for inputs to generative models.
 
-    The x_init tensor is typically sampled from a prior distribution,
-    with shape determined by the input data (e.g., sequence length
-    determines atom count).
-
     Attributes
     ----------
-    x_init : Float[Array, "*batch atoms 3"]
-        Initial structure coordinates.
-        This can be a reference structure (e.g. for alignment during sampling)
-        or a noisy sample from a prior distribution. This should have the proper
-        shape expected for the given ensemble being sampled, e.g. (4, atoms, 3) for
-        ensemble size 4.
     conditioning : C | None
         Model-specific conditioning features, or None.
     """
 
-    # TODO: make x_init more general (not just Float),
-    # relate this to StateT in Sampler protocol?
-    x_init: Float[Array, "*batch atoms 3"]
     conditioning: C | None
 
 
diff --git a/src/sampleworks/models/rf3/wrapper.py b/src/sampleworks/models/rf3/wrapper.py
index f2ecbc5b..1ff27750 100644
--- a/src/sampleworks/models/rf3/wrapper.py
+++ b/src/sampleworks/models/rf3/wrapper.py
@@ -74,8 +74,6 @@ class RF3Config:
         - str/Path to .json: JSON file with chain_id -> MSA path mapping
         - str/Path to .a3m: Single MSA file applied to all protein chains
         - None (default): No MSA information is used
-    ensemble_size : int
-        Number of samples to generate (batch dimension of x_init). Default is 1.
     recycling_steps : int | None
         Number of recycling steps to perform. Default is None, uses model default.
     disable_chiral_features : bool
@@ -93,7 +91,6 @@ class RF3Config:
     """
 
     msa_path: str | Path | dict | None = None
-    ensemble_size: int = 1
     recycling_steps: int | None = None
     disable_chiral_features: bool = False
     track_chiral_features: bool = False
@@ -103,7 +100,6 @@ def annotate_structure_for_rf3(
     structure: dict,
     *,
     msa_path: str | Path | dict | None = None,
-    ensemble_size: int = 1,
     recycling_steps: int | None = None,
     disable_chiral_features: bool = False,
     track_chiral_features: bool = False,
@@ -120,8 +116,6 @@ def annotate_structure_for_rf3(
         - str/Path to .json: JSON file with chain_id -> MSA path mapping
         - str/Path to .a3m: Single MSA file applied to all protein chains
         - None (default): No MSA information is used
-    ensemble_size : int
-        Number of samples to generate (batch dimension of x_init). Default is 1.
     recycling_steps : int | None
         Number of recycling steps to perform. Default is None, uses model default.
     disable_chiral_features : bool
@@ -136,7 +130,6 @@ def annotate_structure_for_rf3(
     """
     config = RF3Config(
         msa_path=msa_path,
-        ensemble_size=ensemble_size,
         recycling_steps=recycling_steps,
         disable_chiral_features=disable_chiral_features,
         track_chiral_features=track_chiral_features,
@@ -322,7 +315,7 @@ def _inner_model(self) -> RF3WithConfidence:
     def featurize(self, structure: dict) -> GenerativeModelInput[RF3Conditioning]:
         """From an Atomworks structure, calculate RF3 input features.
 
-        Runs trunk forward pass and initializes x_init from prior distribution.
+        Runs the trunk forward pass to produce conditioning features.
 
         Parameters
         ----------
@@ -335,7 +328,7 @@ def featurize(self, structure: dict) -> GenerativeModelInput[RF3Conditioning]:
         Returns
         -------
         GenerativeModelInput[RF3Conditioning]
-            Model input with ``x_init`` and trunk conditioning.
+            Model input with trunk conditioning.
         """
         config = structure.get("_rf3_config", RF3Config())
         if isinstance(config, dict):
@@ -346,7 +339,6 @@ def featurize(self, structure: dict) -> GenerativeModelInput[RF3Conditioning]:
         self._chiral_grad_stats = []
 
         msa_path = config.msa_path
-        ensemble_size = config.ensemble_size
         recycling_steps = config.recycling_steps
 
         if "asym_unit" not in structure:
@@ -520,23 +512,7 @@ def featurize(self, structure: dict) -> GenerativeModelInput[RF3Conditioning]:
             )
             logger.info("Chiral features disabled: zeroed out chiral_centers in features dict")
 
-        # x_init here is a shape-compatible reference carried with the featurized
-        # model input. During guided sampling, alignment/reference coordinates are
-        # built later via process_structure_to_trajectory_input() and AtomReconciler.
-        # TODO: figure out if this is necessary or if we should just remove x_init completely.
-        if len(true_atom_array) == num_atoms:
-            x_init = torch.tensor(true_atom_array.coord, device=self.device, dtype=torch.float32)
-            x_init = match_batch(x_init.unsqueeze(0), target_batch_size=ensemble_size)
-        else:
-            logger.info(
-                f"Input structure has {len(true_atom_array)} atoms, but RF3 operates on "
-                f"{num_atoms} atoms after model-specific preprocessing. Initializing "
-                "x_init from prior to match model shape; guided sampling will "
-                "reconcile alignment and reward inputs later on the common atom subset."
-            )
-            x_init = self.initialize_from_prior(batch_size=ensemble_size, shape=(num_atoms, 3))
-
-        return GenerativeModelInput(x_init=x_init, conditioning=conditioning)
+        return GenerativeModelInput(conditioning=conditioning)
 
     def _pairformer_pass(
         self, features: dict[str, Any], grad_needed: bool = False, recycling_steps: int = 10
diff --git a/src/sampleworks/utils/guidance_script_utils.py b/src/sampleworks/utils/guidance_script_utils.py
index 86b4802c..1598618c 100644
--- a/src/sampleworks/utils/guidance_script_utils.py
+++ b/src/sampleworks/utils/guidance_script_utils.py
@@ -458,15 +458,12 @@ def _run_guidance(args: GuidanceConfig, guidance_type: str, model_wrapper, devic
     if "Protenix" in wrapper_class_name:
         from sampleworks.models.protenix.wrapper import annotate_structure_for_protenix
 
-        structure = annotate_structure_for_protenix(
-            structure, ensemble_size=args.ensemble_size, recycling_steps=recycling_steps
-        )
+        structure = annotate_structure_for_protenix(structure, recycling_steps=recycling_steps)
     elif "RF3" in wrapper_class_name:
         from sampleworks.models.rf3.wrapper import annotate_structure_for_rf3
 
         structure = annotate_structure_for_rf3(
             structure,
-            ensemble_size=args.ensemble_size,
             recycling_steps=recycling_steps,
             msa_path=getattr(args, "msa_path", None),
             disable_chiral_features=getattr(args, "disable_chiral_features", False),
@@ -481,7 +478,6 @@ def _run_guidance(args: GuidanceConfig, guidance_type: str, model_wrapper, devic
         structure = process_structure_for_boltz(
             structure,
             out_dir=args.output_dir,
-            ensemble_size=args.ensemble_size,
             recycling_steps=recycling_steps,
         )
     else:
diff --git a/tests/integration/test_no_guidance_geometry.py b/tests/integration/test_no_guidance_geometry.py
index 18ad737a..41310756 100644
--- a/tests/integration/test_no_guidance_geometry.py
+++ b/tests/integration/test_no_guidance_geometry.py
@@ -52,9 +52,7 @@ def test_no_guidance_produces_valid_peptide_geometry(
         structure = request.getfixturevalue(structure_fixture)
         device = wrapper.device if hasattr(wrapper, "device") else torch.device("cpu")
 
-        annotated = annotate_structure_for_wrapper_type(
-            wrapper_type, structure, temp_output_dir, ensemble_size=1
-        )
+        annotated = annotate_structure_for_wrapper_type(wrapper_type, structure, temp_output_dir)
         features = wrapper.featurize(annotated)
 
         sampler = AF3EDMSampler(
diff --git a/tests/integration/test_pipeline_integration.py b/tests/integration/test_pipeline_integration.py
index e956ac16..7073d005 100644
--- a/tests/integration/test_pipeline_integration.py
+++ b/tests/integration/test_pipeline_integration.py
@@ -510,7 +510,7 @@ class TestTrajectoryScalerMatrixMock:
         "trajectory_scaler_type", get_all_trajectory_scalers(), ids=lambda s: s.value
     )
     @pytest.mark.parametrize("sampler_type", get_all_trajectory_samplers(), ids=lambda s: s.value)
-    def test_trajectory_scaler_returns_guidance_output(
+    def test_trajectory_scaler_initializes_from_prior_and_returns_guidance_output(
         self,
         trajectory_scaler_type: TrajectoryScalers,
         sampler_type: TrajectorySamplers,
@@ -520,7 +520,7 @@ def test_trajectory_scaler_returns_guidance_output(
         mock_step_scaler: MockStepScaler,
         mock_gradient_reward: MockGradientRewardFunction,
     ):
-        """TrajectoryScaler.sample() returns valid GuidanceOutput."""
+        """TrajectoryScaler.sample() initializes from the prior and returns valid output."""
         sampler = create_sampler_from_type(sampler_type, device=device)
         num_steps = 5
         ensemble_size = 2
@@ -546,6 +546,9 @@ def test_trajectory_scaler_returns_guidance_output(
         assert result.trajectory is not None
         assert len(result.trajectory) == num_steps
         assert torch.isfinite(torch.as_tensor(result.final_state)).all()
+        assert len(mock_wrapper.prior_initialization_features) == 1
+        assert mock_wrapper.prior_initialization_features[0] is not None
+        assert not hasattr(mock_wrapper.prior_initialization_features[0], "x_init")
 
     @pytest.mark.parametrize(
         "trajectory_scaler_type", get_all_trajectory_scalers(), ids=lambda s: s.value
@@ -929,7 +932,7 @@ def test_featurize_produces_valid_input(
         temp_output_dir,
         request,
     ):
-        """featurize() should produce valid GenerativeModelInput."""
+        """featurize() should produce valid conditioning for prior initialization."""
         wrapper = request.getfixturevalue(get_fixture_name_for_wrapper_type(wrapper_type))
         structure = request.getfixturevalue(structure_fixture)
 
@@ -937,8 +940,8 @@ def test_featurize_produces_valid_input(
         features = wrapper.featurize(annotated)
 
         assert features is not None
-        assert features.x_init is not None
-        assert features.x_init.ndim >= 2
+        assert features.conditioning is not None
+        assert not hasattr(features, "x_init")
 
     def test_single_step_runs(
         self,
diff --git a/tests/mocks/model_wrappers.py b/tests/mocks/model_wrappers.py
index 1084a462..487c5c7d 100644
--- a/tests/mocks/model_wrappers.py
+++ b/tests/mocks/model_wrappers.py
@@ -86,7 +86,7 @@ class MockFlowModelWrapper:
     """Mock FlowModelWrapper that satisfies the protocol without a real model.
 
     This mock provides deterministic behavior for testing:
-    - featurize() returns a GenerativeModelInput with random but reproducible data
+    - featurize() returns a GenerativeModelInput with conditioning data
     - step() returns a slightly denoised version of the input
     - initialize_from_prior() returns random noise
 
@@ -110,12 +110,12 @@ def __init__(
         self.num_atoms = num_atoms
         self.device = device or torch.device("cpu")
         self.target = target
+        self.prior_initialization_features: list[GenerativeModelInput[MockConditioning] | None] = []
 
     def featurize(self, structure: dict, **kwargs: Any) -> GenerativeModelInput[MockConditioning]:
         """Featurize a structure dict into model inputs."""
-        x_init = torch.randn(1, self.num_atoms, 3, device=self.device)
         conditioning = MockConditioning(sequence_length=10, num_atoms=self.num_atoms)
-        return GenerativeModelInput(x_init=x_init, conditioning=conditioning)
+        return GenerativeModelInput(conditioning=conditioning)
 
     def step(
         self,
@@ -147,13 +147,14 @@ def initialize_from_prior(
         shape: tuple[int, ...] | None = None,
     ) -> Float[Tensor, "batch atoms 3"]:
         """Initialize coordinates from prior distribution (random noise)."""
+        self.prior_initialization_features.append(features)
         if features is None and shape is None:
             raise ValueError("Either features or shape must be provided")
         if shape is not None:
             return torch.randn(batch_size, *shape, device=self.device)
-        assert features is not None
-        num_atoms = features.x_init.shape[-2]
-        return torch.randn(batch_size, num_atoms, 3, device=self.device)
+        if features is None or features.conditioning is None:
+            raise ValueError("features with conditioning required")
+        return torch.randn(batch_size, features.conditioning.num_atoms, 3, device=self.device)
 
 
 class MismatchCaseWrapper:
@@ -178,13 +179,12 @@ def __init__(
 
     def featurize(self, structure: dict, **kwargs: Any) -> GenerativeModelInput[MockConditioning]:
         """Return conditioning with case model atom array for reconciliation."""
-        x_init = torch.randn(1, self.n_model, 3, device=self.device)
         conditioning = MockConditioning(
             sequence_length=10,
             num_atoms=self.n_model,
             model_atom_array=self.case.model_atom_array,
         )
-        return GenerativeModelInput(x_init=x_init, conditioning=conditioning)
+        return GenerativeModelInput(conditioning=conditioning)
 
     def step(
         self,
diff --git a/tests/models/boltz/test_boltz_wrapper.py b/tests/models/boltz/test_boltz_wrapper.py
index 94c7d085..c0afa182 100644
--- a/tests/models/boltz/test_boltz_wrapper.py
+++ b/tests/models/boltz/test_boltz_wrapper.py
@@ -24,7 +24,6 @@
 )
 from sampleworks.utils.guidance_constants import StructurePredictor
 from tests.conftest import get_fixture_name_for_wrapper_type, STRUCTURES
-from torch import Tensor
 
 
 BOLTZ_WRAPPER_TYPES = [StructurePredictor.BOLTZ_1, StructurePredictor.BOLTZ_2]
@@ -113,7 +112,6 @@ def test_annotate_default_values(self, structure_6b8x: dict, temp_output_dir: Pa
         result = process_structure_for_boltz(structure_6b8x, out_dir=temp_output_dir)
         config = result["_boltz_config"]
         assert config.num_workers == 0
-        assert config.ensemble_size == 1
         assert config.recycling_steps == 3
 
     def test_annotate_custom_values(self, structure_6b8x: dict, temp_output_dir: Path):
@@ -121,12 +119,10 @@ def test_annotate_custom_values(self, structure_6b8x: dict, temp_output_dir: Pat
             structure_6b8x,
             out_dir=temp_output_dir,
             num_workers=4,
-            ensemble_size=2,
             recycling_steps=5,
         )
         config = result["_boltz_config"]
         assert config.num_workers == 4
-        assert config.ensemble_size == 2
         assert config.recycling_steps == 5
 
     def test_annotate_out_dir_from_metadata(self, structure_6b8x: dict):
@@ -143,19 +139,16 @@ def test_boltz_config_default_values(self):
         config = BoltzConfig()
         assert config.out_dir is None
         assert config.num_workers == 0
-        assert config.ensemble_size == 1
         assert config.recycling_steps == 3
 
     def test_boltz_config_custom_values(self, temp_output_dir: Path):
         config = BoltzConfig(
             out_dir=temp_output_dir,
             num_workers=4,
-            ensemble_size=5,
             recycling_steps=2,
         )
         assert config.out_dir == temp_output_dir
         assert config.num_workers == 4
-        assert config.ensemble_size == 5
         assert config.recycling_steps == 2
 
 
@@ -256,9 +249,8 @@ def test_featurize_returns_generative_model_input(
         annotated = process_structure_for_boltz(structure, out_dir=temp_output_dir)
         features = wrapper.featurize(annotated)
         assert isinstance(features, GenerativeModelInput)
-        assert features.x_init is not None
         assert isinstance(features.conditioning, BoltzConditioning)
-        assert torch.isfinite(torch.as_tensor(features.x_init)).all()
+        assert not hasattr(features, "x_init")
 
     def test_featurize_creates_data_module(
         self,
@@ -274,22 +266,6 @@ def test_featurize_creates_data_module(
         assert hasattr(wrapper, "data_module")
         assert wrapper.data_module is not None
 
-    def test_featurize_with_ensemble_size(
-        self,
-        wrapper_type: StructurePredictor,
-        structure_fixture: str,
-        temp_output_dir: Path,
-        request,
-    ):
-        wrapper = request.getfixturevalue(get_fixture_name_for_wrapper_type(wrapper_type))
-        structure = request.getfixturevalue(structure_fixture)
-        ensemble_size = 3
-        annotated = process_structure_for_boltz(
-            structure, out_dir=temp_output_dir, ensemble_size=ensemble_size
-        )
-        features = wrapper.featurize(annotated)
-        assert features.x_init.shape[0] == ensemble_size
-
 
 @pytest.mark.gpu
 @pytest.mark.slow
@@ -309,10 +285,10 @@ def test_step_returns_tensor(
         structure = request.getfixturevalue(structure_fixture)
         annotated = process_structure_for_boltz(structure, out_dir=temp_output_dir)
         features = wrapper.featurize(annotated)
-        x_init = cast(Tensor, features.x_init)
+        state = wrapper.initialize_from_prior(batch_size=1, features=features)
 
         t = torch.tensor([1.0])
-        result = wrapper.step(x_init, t, features=features)
+        result = wrapper.step(state, t, features=features)
 
         assert torch.is_tensor(result)
 
@@ -327,12 +303,12 @@ def test_step_output_shape_matches_input(
         structure = request.getfixturevalue(structure_fixture)
         annotated = process_structure_for_boltz(structure, out_dir=temp_output_dir)
         features = wrapper.featurize(annotated)
-        x_init = cast(Tensor, features.x_init)
+        state = wrapper.initialize_from_prior(batch_size=1, features=features)
 
         t = torch.tensor([1.0])
-        result = wrapper.step(x_init, t, features=features)
+        result = wrapper.step(state, t, features=features)
 
-        assert result.shape == x_init.shape
+        assert result.shape == state.shape
         assert result.shape[-1] == 3
 
     def test_step_with_high_t(
@@ -346,13 +322,13 @@ def test_step_with_high_t(
         structure = request.getfixturevalue(structure_fixture)
         annotated = process_structure_for_boltz(structure, out_dir=temp_output_dir)
         features = wrapper.featurize(annotated)
-        x_init = cast(Tensor, features.x_init)
+        state = wrapper.initialize_from_prior(batch_size=1, features=features)
 
         t = torch.tensor([160.0])
-        result = wrapper.step(x_init, t, features=features)
+        result = wrapper.step(state, t, features=features)
 
         assert torch.is_tensor(result)
-        assert result.shape == x_init.shape
+        assert result.shape == state.shape
 
     def test_step_with_low_t(
         self,
@@ -365,13 +341,13 @@ def test_step_with_low_t(
         structure = request.getfixturevalue(structure_fixture)
         annotated = process_structure_for_boltz(structure, out_dir=temp_output_dir)
         features = wrapper.featurize(annotated)
-        x_init = cast(Tensor, features.x_init)
+        state = wrapper.initialize_from_prior(batch_size=1, features=features)
 
         t = torch.tensor([0.01])
-        result = wrapper.step(x_init, t, features=features)
+        result = wrapper.step(state, t, features=features)
 
         assert torch.is_tensor(result)
-        assert result.shape == x_init.shape
+        assert result.shape == state.shape
 
     def test_step_with_float_t(
         self,
@@ -384,9 +360,9 @@ def test_step_with_float_t(
         structure = request.getfixturevalue(structure_fixture)
         annotated = process_structure_for_boltz(structure, out_dir=temp_output_dir)
         features = wrapper.featurize(annotated)
-        x_init = cast(Tensor, features.x_init)
+        state = wrapper.initialize_from_prior(batch_size=1, features=features)
 
-        result = wrapper.step(x_init, 1.0, features=features)
+        result = wrapper.step(state, 1.0, features=features)
 
         assert torch.is_tensor(result)
 
@@ -401,11 +377,11 @@ def test_step_requires_features(
         structure = request.getfixturevalue(structure_fixture)
         annotated = process_structure_for_boltz(structure, out_dir=temp_output_dir)
         features = wrapper.featurize(annotated)
-        x_init = cast(Tensor, features.x_init)
+        state = wrapper.initialize_from_prior(batch_size=1, features=features)
 
         t = torch.tensor([1.0])
         with pytest.raises(ValueError, match="features"):
-            wrapper.step(x_init, t, features=None)
+            wrapper.step(state, t, features=None)
 
     def test_step_denoises_input(
         self,
@@ -418,12 +394,12 @@ def test_step_denoises_input(
         structure = request.getfixturevalue(structure_fixture)
         annotated = process_structure_for_boltz(structure, out_dir=temp_output_dir)
         features = wrapper.featurize(annotated)
-        x_init = cast(Tensor, features.x_init)
+        state = wrapper.initialize_from_prior(batch_size=1, features=features)
 
         t = torch.tensor([1.0])
-        result = wrapper.step(x_init, t, features=features)
+        result = wrapper.step(state, t, features=features)
 
-        assert not torch.allclose(result, x_init)
+        assert not torch.allclose(result, state)
 
 
 @pytest.mark.gpu
@@ -545,14 +521,14 @@ def test_full_pipeline(
         assert isinstance(features, GenerativeModelInput)
         assert isinstance(features.conditioning, BoltzConditioning)
 
-        x_init = wrapper.initialize_from_prior(batch_size=2, features=features)
-        assert torch.is_tensor(x_init)
-        assert x_init.shape[0] == 2
+        state = wrapper.initialize_from_prior(batch_size=2, features=features)
+        assert torch.is_tensor(state)
+        assert state.shape[0] == 2
 
         t = torch.tensor([1.0])
-        output = wrapper.step(x_init, t, features=features)
+        output = wrapper.step(state, t, features=features)
         assert torch.is_tensor(output)
-        assert output.shape == x_init.shape
+        assert output.shape == state.shape
         assert torch.isfinite(output).all()
 
     def test_multiple_step_calls(
diff --git a/tests/models/test_model_wrapper_protocol.py b/tests/models/test_model_wrapper_protocol.py
index 86fe109c..b9268fd3 100644
--- a/tests/models/test_model_wrapper_protocol.py
+++ b/tests/models/test_model_wrapper_protocol.py
@@ -21,6 +21,17 @@ def get_slow_wrapper_infos() -> list[ComponentInfo]:
     return [info for info in MODEL_WRAPPER_REGISTRY.values() if info.requires_checkpoint]
 
 
+def test_generative_model_input_carries_only_conditioning():
+    """GenerativeModelInput leaves state initialization to the model wrapper."""
+    from sampleworks.models.protocol import GenerativeModelInput
+
+    conditioning = {"num_atoms": 3}
+    features = GenerativeModelInput(conditioning=conditioning)
+
+    assert features.conditioning is conditioning
+    assert not hasattr(features, "x_init")
+
+
 @pytest.mark.parametrize("wrapper_info", get_slow_wrapper_infos(), ids=lambda w: w.name)
 class TestFlowModelWrapperProtocol:
     """Test that wrappers implement FlowModelWrapper protocol correctly.
@@ -49,7 +60,7 @@ def test_isinstance_flow_model_wrapper(self, wrapper_info: ComponentInfo, reques
     def test_featurize_returns_generative_model_input(
         self, wrapper_info: ComponentInfo, structure_fixture: str, temp_output_dir, request
     ):
-        """Test featurize returns GenerativeModelInput with x_init and conditioning."""
+        """Test featurize returns GenerativeModelInput with conditioning."""
         from sampleworks.models.protocol import GenerativeModelInput
 
         fixture_name = get_fixture_name_for_wrapper(wrapper_info)
@@ -63,9 +74,6 @@ def test_featurize_returns_generative_model_input(
         assert isinstance(features, GenerativeModelInput), (
             f"{wrapper_info.name}.featurize must return GenerativeModelInput, got {type(features)}"
         )
-        assert features.x_init is not None, (
-            f"{wrapper_info.name}.featurize returned None for x_init"
-        )
         assert features.conditioning is not None, (
             f"{wrapper_info.name}.featurize returned None for conditioning"
         )
@@ -74,37 +82,6 @@ def test_featurize_returns_generative_model_input(
             f"got {type(features.conditioning)}"
         )
 
-    @pytest.mark.gpu
-    @pytest.mark.slow
-    @pytest.mark.parametrize(
-        "structure_fixture", STRUCTURES, ids=lambda s: s.replace("structure_", "")
-    )
-    def test_featurize_x_init_shape(
-        self, wrapper_info: ComponentInfo, structure_fixture: str, temp_output_dir, request
-    ):
-        """Test featurize x_init has correct shape (batch, atoms, 3)."""
-        fixture_name = get_fixture_name_for_wrapper(wrapper_info)
-        wrapper = request.getfixturevalue(fixture_name)
-        structure = request.getfixturevalue(structure_fixture)
-
-        ensemble_size = 2
-        annotated = annotate_structure_for_wrapper(
-            wrapper_info, structure, temp_output_dir, ensemble_size=ensemble_size
-        )
-        features = wrapper.featurize(annotated)
-
-        assert features.x_init.ndim == 3, (
-            f"{wrapper_info.name}.featurize x_init should be 3D, got {features.x_init.ndim}D"
-        )
-        assert features.x_init.shape[0] == ensemble_size, (
-            f"{wrapper_info.name}.featurize x_init batch should be {ensemble_size}, "
-            f"got {features.x_init.shape[0]}"
-        )
-        assert features.x_init.shape[2] == 3, (
-            f"{wrapper_info.name}.featurize x_init last dim should be 3, "
-            f"got {features.x_init.shape[2]}"
-        )
-
     @pytest.mark.gpu
     @pytest.mark.slow
     @pytest.mark.parametrize(
@@ -122,7 +99,8 @@ def test_step_returns_tensor(
         features = wrapper.featurize(annotated)
 
         t = torch.tensor([1.0])
-        result = wrapper.step(features.x_init, t, features=features)
+        state = wrapper.initialize_from_prior(batch_size=1, features=features)
+        result = wrapper.step(state, t, features=features)
 
         assert torch.is_tensor(result), (
             f"{wrapper_info.name}.step must return Tensor, got {type(result)}"
@@ -130,9 +108,8 @@ def test_step_returns_tensor(
         assert result.shape[-1] == 3, (
             f"{wrapper_info.name}.step result last dim should be 3, got {result.shape[-1]}"
         )
-        assert result.shape == features.x_init.shape, (
-            f"{wrapper_info.name}.step output shape {result.shape} != input shape "
-            f"{features.x_init.shape}"
+        assert result.shape == state.shape, (
+            f"{wrapper_info.name}.step output shape {result.shape} != input shape {state.shape}"
         )
 
     @pytest.mark.gpu
@@ -152,7 +129,8 @@ def test_step_with_float_t(
         features = wrapper.featurize(annotated)
 
         t = 1.0
-        result = wrapper.step(features.x_init, t, features=features)
+        state = wrapper.initialize_from_prior(batch_size=1, features=features)
+        result = wrapper.step(state, t, features=features)
 
         assert torch.is_tensor(result), f"{wrapper_info.name}.step must return Tensor with float t"
 
@@ -172,6 +150,7 @@ def test_initialize_from_prior_with_features(
         annotated = annotate_structure_for_wrapper(wrapper_info, structure, temp_output_dir)
         features = wrapper.featurize(annotated)
 
+        assert not hasattr(features, "x_init")
         batch_size = 3
         result = wrapper.initialize_from_prior(batch_size, features=features)
 
diff --git a/tests/models/test_rf3_atom_ordering.py b/tests/models/test_rf3_atom_ordering.py
index 5cdfa677..184f4423 100644
--- a/tests/models/test_rf3_atom_ordering.py
+++ b/tests/models/test_rf3_atom_ordering.py
@@ -37,7 +37,7 @@ def test_model_atom_array_matches_feature_count(self, rf3_wrapper, structure_fix
         internal atom accounting, causing coordinate misalignment.
         """
         structure = request.getfixturevalue(structure_fixture)
-        annotated = annotate_structure_for_rf3(structure, ensemble_size=1)
+        annotated = annotate_structure_for_rf3(structure)
         features = rf3_wrapper.featurize(annotated)
         cond = features.conditioning
 
@@ -66,7 +66,7 @@ def test_no_oxt_atoms_in_model_atom_array(self, rf3_wrapper, structure_fixture,
         The pipeline's ``RemoveTerminalOxygen`` transform removes these atoms.
         """
         structure = request.getfixturevalue(structure_fixture)
-        annotated = annotate_structure_for_rf3(structure, ensemble_size=1)
+        annotated = annotate_structure_for_rf3(structure)
         features = rf3_wrapper.featurize(annotated)
         cond = features.conditioning
 
@@ -93,7 +93,7 @@ def test_no_hydrogen_atoms_in_model_atom_array(self, rf3_wrapper, structure_fixt
         The pipeline's ``RemoveHydrogens`` transform removes these.
         """
         structure = request.getfixturevalue(structure_fixture)
-        annotated = annotate_structure_for_rf3(structure, ensemble_size=1)
+        annotated = annotate_structure_for_rf3(structure)
         features = rf3_wrapper.featurize(annotated)
         cond = features.conditioning
 
diff --git a/tests/models/test_rf3_chiral_features.py b/tests/models/test_rf3_chiral_features.py
index f691cc6d..04d04293 100644
--- a/tests/models/test_rf3_chiral_features.py
+++ b/tests/models/test_rf3_chiral_features.py
@@ -42,7 +42,6 @@ def test_track_chiral_features_populates_stats(self, rf3_wrapper, structure_1vme
         entry per denoising step, each containing 't' and 'l2_norm' keys."""
         annotated = annotate_structure_for_rf3(
             structure_1vme,
-            ensemble_size=1,
             track_chiral_features=True,
             disable_chiral_features=False,
         )
@@ -73,7 +72,6 @@ def test_disable_chiral_features_zeros_features(self, rf3_wrapper, structure_1vm
         conditioning features dict."""
         annotated = annotate_structure_for_rf3(
             structure_1vme,
-            ensemble_size=1,
             disable_chiral_features=True,
             track_chiral_features=False,
         )
@@ -104,7 +102,6 @@ def test_tracking_consistent_across_disabled_and_enabled(self, rf3_wrapper, stru
         def featurize_and_run(*, disable_chiral: bool):
             annotated = annotate_structure_for_rf3(
                 structure_1vme,
-                ensemble_size=1,
                 disable_chiral_features=disable_chiral,
                 track_chiral_features=True,
             )
@@ -149,7 +146,6 @@ def test_no_tracking_when_disabled(self, rf3_wrapper, structure_1vme):
         """With track_chiral_features=False, _chiral_grad_stats stays empty."""
         annotated = annotate_structure_for_rf3(
             structure_1vme,
-            ensemble_size=1,
             track_chiral_features=False,
             disable_chiral_features=False,
         )
@@ -166,7 +162,6 @@ def test_baseline_chiral_features_present(self, rf3_wrapper, structure_1vme):
         disable_chiral_features=False."""
         annotated = annotate_structure_for_rf3(
             structure_1vme,
-            ensemble_size=1,
             disable_chiral_features=False,
             track_chiral_features=False,
         )
@@ -186,7 +181,6 @@ def test_featurize_resets_tracking_state(self, rf3_wrapper, structure_1vme):
 
         annotated = annotate_structure_for_rf3(
             structure_1vme,
-            ensemble_size=1,
             track_chiral_features=True,
         )
         features = rf3_wrapper.featurize(annotated)

From 0f53cca323c6d8effdb343fcf6b015cad704117f Mon Sep 17 00:00:00 2001
From: xraymemory <me.anzuoni@gmail.com>
Date: Tue, 21 Jul 2026 00:03:04 -0400
Subject: [PATCH 2/2] test(models): cover batched prior steps

---
 tests/models/test_model_wrapper_protocol.py | 16 +++++++++++++---
 1 file changed, 13 insertions(+), 3 deletions(-)

diff --git a/tests/models/test_model_wrapper_protocol.py b/tests/models/test_model_wrapper_protocol.py
index b9268fd3..131acce1 100644
--- a/tests/models/test_model_wrapper_protocol.py
+++ b/tests/models/test_model_wrapper_protocol.py
@@ -87,8 +87,14 @@ def test_featurize_returns_generative_model_input(
     @pytest.mark.parametrize(
         "structure_fixture", STRUCTURES, ids=lambda s: s.replace("structure_", "")
     )
+    @pytest.mark.parametrize("batch_size", [1, 2])
     def test_step_returns_tensor(
-        self, wrapper_info: ComponentInfo, structure_fixture: str, temp_output_dir, request
+        self,
+        wrapper_info: ComponentInfo,
+        structure_fixture: str,
+        batch_size: int,
+        temp_output_dir,
+        request,
     ):
         """Test step(x_t, t, features) returns coordinates tensor."""
         fixture_name = get_fixture_name_for_wrapper(wrapper_info)
@@ -98,8 +104,8 @@ def test_step_returns_tensor(
         annotated = annotate_structure_for_wrapper(wrapper_info, structure, temp_output_dir)
         features = wrapper.featurize(annotated)
 
-        t = torch.tensor([1.0])
-        state = wrapper.initialize_from_prior(batch_size=1, features=features)
+        t = torch.ones(batch_size)
+        state = wrapper.initialize_from_prior(batch_size=batch_size, features=features)
         result = wrapper.step(state, t, features=features)
 
         assert torch.is_tensor(result), (
@@ -111,6 +117,10 @@ def test_step_returns_tensor(
         assert result.shape == state.shape, (
             f"{wrapper_info.name}.step output shape {result.shape} != input shape {state.shape}"
         )
+        assert result.shape[0] == batch_size, (
+            f"{wrapper_info.name}.step output batch should be {batch_size}, got {result.shape[0]}"
+        )
+        assert torch.isfinite(result).all(), f"{wrapper_info.name}.step returned non-finite values"
 
     @pytest.mark.gpu
     @pytest.mark.slow