Consolidate core API definitions into hd_api.py.

PiperOrigin-RevId: 952108946

Author: ccrepy

gemma/diffusion/hackable_diffusion_adapter/hd/hd_gemma_ar_state_handler.py

@@ -36,7 +36,7 @@
 Tokens = Int["*B L"]
 Conditioning = hd_typing.Conditioning
 DataArray = hd_typing.DataArray
-DiffusionStepTree = hd.sampling.DiffusionStepTree
+DiffusionStep = hd.sampling.DiffusionStep
 
 
 ################################################################################
@@ -50,7 +50,7 @@ class PropagateSelfConditioningFn(hd.sampling.UpdateConditioningFn):
   def __call__(
       self,
       conditioning: Conditioning,
-      step_carry: DiffusionStepTree,
+      step_carry: DiffusionStep,
   ) -> Conditioning:
     """Update conditioning based on the current diffusion step.
 
@@ -275,7 +275,7 @@ def init_ar_state(
 
   def update_ar_state(
       self,
-      canvas_last_step: DiffusionStepTree,
+      canvas_last_step: DiffusionStep,
       sampler_state: ar_diffusion_sampler.SamplerState,
   ) -> ar_diffusion_sampler.SamplerState:
     """Post-processes a sampled canvas and updates the sampler state.