Cleanup *_step_sampler.py - fix some interfaces and use keyword arguments consistently

PiperOrigin-RevId: 911378722

Author: ccrepy

hackable_diffusion/lib/sampling/gaussian_step_sampler.py

@@ -90,7 +90,7 @@ def initialize(
     )
 
   @kt.typechecked
-  def update(
+  def _update(
       self,
       prediction: TargetInfo,
       current_step: DiffusionStep,
@@ -137,14 +137,28 @@ def update(
         aux=dict(),
     )
 
+  @kt.typechecked
+  def update(
+      self,
+      prediction: TargetInfo,
+      current_step: DiffusionStep,
+      next_step_info: StepInfo,
+  ) -> DiffusionStep:
+    return self._update(
+        prediction,
+        current_step,
+        next_step_info,
+        stochastic=True,
+    )
+
   @kt.typechecked
   def finalize(
       self,
       prediction: TargetInfo,
       current_step: DiffusionStep,
       last_step_info: StepInfo,
   ) -> DiffusionStep:
-    return self.update(
+    return self._update(
         prediction,
         current_step,
         last_step_info,
@@ -381,7 +395,7 @@ def initialize(
     )
 
   @kt.typechecked
-  def update(
+  def _update(
       self,
       prediction: TargetInfo,
       current_step: DiffusionStep,
@@ -424,14 +438,28 @@ def update(
         aux=dict(),
     )
 
+  @kt.typechecked
+  def update(
+      self,
+      prediction: TargetInfo,
+      current_step: DiffusionStep,
+      next_step_info: StepInfo,
+  ) -> DiffusionStep:
+    return self._update(
+        prediction,
+        current_step,
+        next_step_info,
+        stochastic=True,
+    )
+
   @kt.typechecked
   def finalize(
       self,
       prediction: TargetInfo,
       current_step: DiffusionStep,
       last_step_info: StepInfo,
   ) -> DiffusionStep:
-    return self.update(
+    return self._update(
         prediction,
         current_step,
         last_step_info,

hackable_diffusion/lib/sampling/simplicial_step_sampler.py

@@ -160,9 +160,9 @@ def update(
 
     # Get logits
     logits = self.corruption_process.convert_predictions(
-        prediction,
-        log_xt,
-        time,
+        prediction=prediction,
+        xt=log_xt,
+        time=time,
     )['logits']
 
     # Sample hard token