Consolidate core API definitions into hd_api.py.

docs/corruption.md

@@ -28,7 +28,7 @@ The main components are:
 
 ## `CorruptionProcess` Protocol
 
-(`lib/corruption/base.py`)
+(`lib/hd_api.py`)
 
 The `CorruptionProcess` is a protocol (an interface) that all corruption
 processes must implement. It defines the following key methods:

docs/inference.md

@@ -31,26 +31,26 @@ The main components are:
 
 ## `InferenceFn` Protocol
 
-(`lib/inference/base.py`)
+(`lib/hd_api.py`)
 
 The `InferenceFn` is a protocol that defines a callable with the following
 signature:
 
 ```python
 def __call__(
-    self, time: TimeTree, xt: DataTree, conditioning: Conditioning | None
-) -> TargetInfoTree:
+    self, time: TimeArray, xt: DataArray, conditioning: Conditioning | None
+) -> TargetInfo:
   ...
 ```
 
   * **Inputs**:
       * `time`: The current timestep(s).
       * `xt`: The noisy data at time `t`.
       * `conditioning`: A dictionary of conditioning signals.
-  * **Output**:
-      * A `TargetInfoTree` (a dictionary or pytree of dictionaries) containing
-        the model's predictions. For a `GaussianProcess`, this would include
-        keys like `'x0'`, `'epsilon'`, `'score'`, etc.
+*   **Output**:
+    *   A `TargetInfo` dictionary containing the model's predictions. For a
+        `GaussianProcess`, this would include keys like `'x0'`, `'epsilon'`,
+        `'score'`, etc.
 
 ## Classifier-Free Guidance
 

docs/sampling.md

@@ -39,7 +39,7 @@ The overall flow for `N` steps is:
 
 ## Core Data Structures
 
-(`lib/sampling/base.py`)
+(`lib/hd_api.py`)
 
 Two main data structures manage the state of the sampling loop:
 

docs/training.md

@@ -25,7 +25,7 @@ The main components are:
 
 ## `DiffusionLoss` Protocol
 
-(`lib/training/base.py`)
+(`lib/hd_api.py`)
 
 This protocol defines the basic interface for a diffusion loss. It's a callable
 that takes the model's predictions, the ground truth targets, and the time, and
@@ -34,10 +34,10 @@ returns a batch-wise loss.
 ```python
 def __call__(
     self,
-    preds: TargetInfoTree,
-    targets: TargetInfoTree,
-    time: TimeTree,
-) -> LossOutputTree:
+    preds: TargetInfo,
+    targets: TargetInfo,
+    time: TimeArray,
+) -> LossOutput:
   ...
 ```
 

hackable_diffusion/hd.py

@@ -26,12 +26,13 @@
 from hackable_diffusion.lib import corruption
 from hackable_diffusion.lib import diffusion_network
 from hackable_diffusion.lib import fast_random
+from hackable_diffusion.lib import hd_api
 from hackable_diffusion.lib import hd_typing
 from hackable_diffusion.lib import inference
+from hackable_diffusion.lib import jax_helpers
 from hackable_diffusion.lib import manifolds
 from hackable_diffusion.lib import multimodal
 from hackable_diffusion.lib import sampling
 from hackable_diffusion.lib import training
-from hackable_diffusion.lib import jax_helpers
 
 # pylint: enable=g-importing-member, unused-import

hackable_diffusion/kdiff/core.py

@@ -177,7 +177,7 @@ class KauldronLossWrapper(kd.losses.Loss):
 
   # Implicitly supports `weight` and `mask` as well (see `kd.losses.Loss`).
 
-  loss: hd.training.DiffusionLoss
+  loss: hd.hd_api.DiffusionLoss
 
   @typechecked
   def get_values(

hackable_diffusion/lib/corruption/__init__.py

@@ -15,7 +15,6 @@
 """API for corruption processes."""
 
 # pylint: disable=g-importing-member
-from hackable_diffusion.lib.corruption.base import CorruptionProcess
 from hackable_diffusion.lib.corruption.discrete import CategoricalProcess
 from hackable_diffusion.lib.corruption.discrete import IdentityPostCorruptionFn
 from hackable_diffusion.lib.corruption.discrete import PostCorruptionFn
@@ -43,5 +42,6 @@
 from hackable_diffusion.lib.corruption.simplicial import SimplicialProcess
 from hackable_diffusion.lib.corruption.simplicial import SimplicialSchedule
 from hackable_diffusion.lib.corruption.simplicial import SymmetricSimplicialPostCorruptionFn
+from hackable_diffusion.lib.hd_api import CorruptionProcess
 
 # pylint: enable=g-importing-member

hackable_diffusion/lib/corruption/discrete.py

@@ -20,9 +20,9 @@
 import enum
 from typing import Protocol, Sequence
 
+from hackable_diffusion.lib import hd_api
 from hackable_diffusion.lib import hd_typing
 from hackable_diffusion.lib import jax_helpers
-from hackable_diffusion.lib.corruption import base
 from hackable_diffusion.lib.corruption import schedules
 import jax
 import jax.numpy as jnp
@@ -46,7 +46,7 @@
 TargetInfo = hd_typing.TargetInfo
 TimeArray = hd_typing.TimeArray
 
-CorruptionProcess = base.CorruptionProcess
+CorruptionProcess = hd_api.CorruptionProcess
 DiscreteSchedule = schedules.DiscreteSchedule
 
 ################################################################################

hackable_diffusion/lib/corruption/gaussian.py

@@ -144,9 +144,9 @@
 
 import dataclasses
 
+from hackable_diffusion.lib import hd_api
 from hackable_diffusion.lib import hd_typing
 from hackable_diffusion.lib import jax_helpers
-from hackable_diffusion.lib.corruption import base
 from hackable_diffusion.lib.corruption import schedules
 import immutabledict
 import jax
@@ -165,7 +165,7 @@
 TargetInfo = hd_typing.TargetInfo
 
 GaussianSchedule = schedules.GaussianSchedule
-CorruptionProcess = base.CorruptionProcess
+CorruptionProcess = hd_api.CorruptionProcess
 
 ################################################################################
 # MARK: GaussianProcess

hackable_diffusion/lib/corruption/riemannian.py

@@ -17,10 +17,10 @@
 import dataclasses
 from typing import Any
 
+from hackable_diffusion.lib import hd_api
 from hackable_diffusion.lib import hd_typing
-from hackable_diffusion.lib import manifolds
 from hackable_diffusion.lib import jax_helpers
-from hackable_diffusion.lib.corruption import base
+from hackable_diffusion.lib import manifolds
 from hackable_diffusion.lib.corruption import schedules
 import kauldron.ktyping as kt
 
@@ -31,7 +31,7 @@
 
 
 @dataclasses.dataclass(kw_only=True, frozen=True)
-class RiemannianProcess(base.CorruptionProcess):
+class RiemannianProcess(hd_api.CorruptionProcess):
   """Riemannian Flow Matching corruption process.
 
   This is based on https://arxiv.org/abs/2302.03660.

hackable_diffusion/lib/corruption/schedules_test.py

@@ -434,4 +434,4 @@ def test_evaluate(self):
 
 
 if __name__ == '__main__':
-  absltest.main()
+  absltest.main()

hackable_diffusion/lib/corruption/simplicial.py

@@ -21,9 +21,9 @@
 from typing import Protocol, Sequence
 
 from hackable_diffusion.lib import fast_random
+from hackable_diffusion.lib import hd_api
 from hackable_diffusion.lib import hd_typing
 from hackable_diffusion.lib import jax_helpers
-from hackable_diffusion.lib.corruption import base
 from hackable_diffusion.lib.corruption import schedules
 import jax
 import jax.numpy as jnp
@@ -47,7 +47,7 @@
 TargetInfo = hd_typing.TargetInfo
 TimeArray = hd_typing.TimeArray
 
-CorruptionProcess = base.CorruptionProcess
+CorruptionProcess = hd_api.CorruptionProcess
 SimplicialSchedule = schedules.SimplicialSchedule
 
 ################################################################################

hackable_diffusion/lib/diffusion_network.py

@@ -50,6 +50,7 @@
 Shape = hd_typing.Shape
 ShapeTree = hd_typing.ShapeTree
 
+
 ################################################################################
 # MARK: Rescalers
 ################################################################################
@@ -79,11 +80,11 @@ class ConditionalBackbone(Protocol):
 
   def __call__(
       self,
-      x: DataTree,  # pyrefly: ignore[not-a-type]
+      x: DataArray,  # pyrefly: ignore[not-a-type]
       conditioning_embeddings: ConditioningEmbeddings,
       *,
       is_training: bool,
-  ) -> DataTree:  # pyrefly: ignore[not-a-type]
+  ) -> DataArray:  # pyrefly: ignore[not-a-type]
     ...
 
 
@@ -92,11 +93,11 @@ class BaseDiffusionNetwork(Protocol):
 
   def __call__(
       self,
-      time: TimeTree,  # pyrefly: ignore[not-a-type]
-      xt: DataTree,  # pyrefly: ignore[not-a-type]
+      time: TimeArray,  # pyrefly: ignore[not-a-type]
+      xt: DataArray,  # pyrefly: ignore[not-a-type]
       conditioning: Conditioning | None,
       is_training: bool,
-  ) -> TargetInfoTree:  # pyrefly: ignore[not-a-type]
+  ) -> TargetInfo:  # pyrefly: ignore[not-a-type]
     ...