Merge pull request #10 from mctigger/unify-tensor_distribution-implementations

Unify tensor distribution implementations

Author: mctigger

src/tensorcontainer/tensor_annotated.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import Any, Dict, Iterable, List, Tuple, TypeVar, Union, get_args
+from typing import Any, Iterable, TypeVar, Union, get_args
 
 import torch
 from torch import Tensor
@@ -20,7 +20,7 @@
 class TensorAnnotated(TensorContainer, PytreeRegistered):
     def __init__(
         self,
-        shape: torch.Size | List[int] | Tuple[int],
+        shape: torch.Size | list[int] | tuple[int, ...],
         device: str | torch.device | int | None,
     ):
         super().__init__(shape, device, True)
@@ -76,14 +76,14 @@ def _get_meta_attributes(self):
         return meta_attributes
 
     def _get_pytree_context(
-        self, flat_names: List[str], flat_leaves: List[TDCompatible], meta_data
-    ) -> Tuple:
+        self, flat_names: list[str], flat_leaves: list[TDCompatible], meta_data
+    ) -> tuple:
         batch_ndim = len(self.shape)
         event_ndims = tuple(leaf.ndim - batch_ndim for leaf in flat_leaves)
 
         return flat_names, event_ndims, meta_data, self.device
 
-    def _pytree_flatten(self) -> Tuple[List[Any], Any]:
+    def _pytree_flatten(self) -> tuple[list[Any], Any]:
         tensor_attributes = self._get_tensor_attributes()
         flat_names = list(tensor_attributes.keys())
         flat_values = list(tensor_attributes.values())
@@ -132,8 +132,8 @@ def _pytree_unflatten(cls, leaves: Iterable[Any], context: pytree.Context) -> Se
     @classmethod
     def _init_from_reconstructed(
         cls,
-        tensor_attributes: Dict[str, TDCompatible],
-        meta_attributes: Dict[str, Any],
+        tensor_attributes: dict[str, TDCompatible],
+        meta_attributes: dict[str, Any],
         device,
         shape,
     ):

src/tensorcontainer/tensor_distribution/base.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from abc import abstractmethod
-from typing import Any, Dict, List, Tuple
+from typing import Any
 
 from torch import Size, Tensor
 from torch._C import device
@@ -47,10 +47,10 @@ class TensorDistribution(TensorAnnotated):
     --------
     ```python
     class TensorNormal(TensorDistribution):
-        _loc: Optional[Tensor] = None
-        _scale: Optional[Tensor] = None
+        _loc: Tensor | None = None
+        _scale: Tensor | None = None
 
-        def __init__(self, loc: Tensor, scale: Tensor, validate_args: Optional[bool] = None):
+        def __init__(self, loc: Tensor, scale: Tensor, validate_args: bool | None = None):
             self._loc = loc
             self._scale = scale
             super().__init__(loc.shape, loc.device, validate_args)
@@ -64,7 +64,7 @@ def dist(self) -> Distribution:
 
     def __init__(
         self,
-        shape: Size | List[int] | Tuple[int],
+        shape: Size | list[int] | tuple[int, ...],
         device: str | device | int | None,
         validate_args: bool | None = None,
     ):
@@ -89,8 +89,8 @@ def __init__(
     @classmethod
     def _init_from_reconstructed(
         cls,
-        tensor_attributes: Dict[str, TDCompatible],
-        meta_attributes: Dict[str, Any],
+        tensor_attributes: dict[str, TDCompatible],
+        meta_attributes: dict[str, Any],
         device,
         shape,
     ):
@@ -112,7 +112,7 @@ def _init_from_reconstructed(
         return cls._unflatten_distribution({**tensor_attributes, **meta_attributes})
 
     @classmethod
-    def _unflatten_distribution(cls, attributes: Dict[str, Any]):
+    def _unflatten_distribution(cls, attributes: dict[str, Any]):
         """
         Reconstruct a distribution from flattened tensor and metadata attributes.
 

src/tensorcontainer/tensor_distribution/bernoulli.py

@@ -1,10 +1,10 @@
 from __future__ import annotations
 
-from typing import Any, Dict, Optional, Union
+from typing import Any
 
 from torch import Size, Tensor
 from torch.distributions import Bernoulli
-from torch.distributions.utils import broadcast_all
+from .utils import broadcast_all
 from torch.types import Number
 
 from .base import TensorDistribution
@@ -14,14 +14,14 @@ class TensorBernoulli(TensorDistribution):
     """Tensor-aware Bernoulli distribution."""
 
     # Annotated tensor parameters
-    _probs: Optional[Tensor]
-    _logits: Optional[Tensor]
+    _probs: Tensor | None
+    _logits: Tensor | None
 
     def __init__(
         self,
-        probs: Optional[Union[Number, Tensor]] = None,
-        logits: Optional[Union[Number, Tensor]] = None,
-        validate_args: Optional[bool] = None,
+        probs: Number | Tensor | None = None,
+        logits: Number | Tensor | None = None,
+        validate_args: bool | None = None,
     ):
         if (probs is None) == (logits is None):
             raise ValueError(
@@ -43,7 +43,7 @@ def __init__(
         super().__init__(shape, device, validate_args)
 
     @classmethod
-    def _unflatten_distribution(cls, attributes: Dict[str, Any]) -> TensorBernoulli:
+    def _unflatten_distribution(cls, attributes: dict[str, Any]) -> TensorBernoulli:
         """Reconstruct distribution from tensor attributes."""
         return cls(
             probs=attributes.get("_probs"),
@@ -56,21 +56,6 @@ def dist(self) -> Bernoulli:
             probs=self._probs, logits=self._logits, validate_args=self._validate_args
         )
 
-    def log_prob(self, value: Tensor) -> Tensor:
-        return self.dist().log_prob(value)
-
-    @property
-    def mean(self) -> Tensor:
-        return self.dist().mean
-
-    @property
-    def variance(self) -> Tensor:
-        return self.dist().variance
-
-    @property
-    def mode(self) -> Tensor:
-        return self.dist().mode
-
     @property
     def logits(self) -> Tensor:
         return self.dist().logits

src/tensorcontainer/tensor_distribution/beta.py

@@ -1,10 +1,10 @@
 from __future__ import annotations
 
-from typing import Any, Dict, Optional
+from typing import Any
 
 from torch import Tensor
 from torch.distributions import Beta
-from torch.distributions.utils import broadcast_all
+from .utils import broadcast_all
 
 from .base import TensorDistribution
 
@@ -29,9 +29,9 @@ class TensorBeta(TensorDistribution):
 
     def __init__(
         self,
-        concentration1: Tensor,
-        concentration0: Tensor,
-        validate_args: Optional[bool] = None,
+        concentration1: float | Tensor,
+        concentration0: float | Tensor,
+        validate_args: bool | None = None,
     ):
         self._concentration1, self._concentration0 = broadcast_all(
             concentration1, concentration0
@@ -45,12 +45,12 @@ def __init__(
     @classmethod
     def _unflatten_distribution(
         cls,
-        attributes: Dict[str, Any],
+        attributes: dict[str, Any],
     ) -> TensorBeta:
         """Reconstruct distribution from tensor attributes."""
         return cls(
-            concentration1=attributes.get("_concentration1"),  # type: ignore
-            concentration0=attributes.get("_concentration0"),  # type: ignore
+            concentration1=attributes["_concentration1"],
+            concentration0=attributes["_concentration0"],
             validate_args=attributes.get("_validate_args"),
         )
 
@@ -62,9 +62,6 @@ def dist(self) -> Beta:
             validate_args=self._validate_args,
         )
 
-    def log_prob(self, value: Tensor) -> Tensor:
-        return self.dist().log_prob(value)
-
     @property
     def concentration1(self) -> Tensor:
         """Returns the concentration1 parameter of the distribution."""
@@ -74,23 +71,3 @@ def concentration1(self) -> Tensor:
     def concentration0(self) -> Tensor:
         """Returns the concentration0 parameter of the distribution."""
         return self.dist().concentration0
-
-    @property
-    def mean(self) -> Tensor:
-        """Returns the mean of the distribution."""
-        return self.dist().mean
-
-    @property
-    def variance(self) -> Tensor:
-        """Returns the variance of the distribution."""
-        return self.dist().variance
-
-    @property
-    def mode(self) -> Tensor:
-        """Returns the mode of the distribution."""
-        return self.dist().mode
-
-    @property
-    def stddev(self) -> Tensor:
-        """Returns the standard deviation of the distribution."""
-        return self.dist().stddev

src/tensorcontainer/tensor_distribution/binomial.py

@@ -1,11 +1,10 @@
 from __future__ import annotations
 
-from typing import Any, Dict, Optional, Union
+from typing import Any
 
-import torch
 from torch import Size, Tensor
 from torch.distributions import Binomial
-from torch.distributions.utils import broadcast_all
+from .utils import broadcast_all
 
 from .base import TensorDistribution
 
@@ -24,16 +23,16 @@ class TensorBinomial(TensorDistribution):
     """
 
     # Annotated tensor parameters
-    _total_count: Union[int, Tensor]
-    _probs: Optional[Tensor] = None
-    _logits: Optional[Tensor] = None
+    _total_count: Tensor
+    _probs: Tensor | None = None
+    _logits: Tensor | None = None
 
     def __init__(
         self,
-        total_count: Union[int, Tensor] = 1,
-        probs: Optional[Tensor] = None,
-        logits: Optional[Tensor] = None,
-        validate_args: Optional[bool] = None,
+        total_count: int | Tensor = 1,
+        probs: Tensor | None = None,
+        logits: Tensor | None = None,
+        validate_args: bool | None = None,
     ):
         if (probs is None) == (logits is None):
             raise ValueError(
@@ -42,36 +41,19 @@ def __init__(
 
         if probs is not None:
             self._total_count, self._probs = broadcast_all(total_count, probs)
-            param = self._probs
-            assert param is not None
         else:
             self._total_count, self._logits = broadcast_all(total_count, logits)
-            param = self._logits
-            assert param is not None
 
-        # Ensure total_count has the same dtype as the parameter tensor if it's a Tensor
-        if isinstance(self._total_count, Tensor):
-            self._total_count = self._total_count.type_as(param)
-
-        shape = param.shape
-        device = param.device
+        shape = self._total_count.shape
+        device = self._total_count.device
 
         super().__init__(shape, device, validate_args)
 
     @classmethod
-    def _unflatten_distribution(cls, attributes: Dict[str, Any]):
-        """Reconstruct distribution from tensor attributes."""
+    def _unflatten_distribution(cls, attributes: dict[str, Any]):
         total_count = attributes["_total_count"]
-        if isinstance(total_count, Tensor):
-            total_count = total_count.clone()
-
-        probs = attributes.get("_probs")
-        if probs is not None:
-            probs = probs.clone()
-
-        logits = attributes.get("_logits")
-        if logits is not None:
-            logits = logits.clone()
+        probs = attributes["_probs"]
+        logits = attributes["_logits"]
 
         return cls(
             total_count=total_count,
@@ -81,63 +63,29 @@ def _unflatten_distribution(cls, attributes: Dict[str, Any]):
         )
 
     def dist(self) -> Binomial:
-        total_count = self._total_count
-        if isinstance(total_count, int):
-            # Convert int total_count to a tensor with the correct device and dtype
-            # The device and dtype should match the probs/logits tensor
-            if self._probs is not None:
-                total_count = torch.tensor(
-                    total_count, device=self._probs.device, dtype=self._probs.dtype
-                )
-            elif self._logits is not None:
-                total_count = torch.tensor(
-                    total_count, device=self._logits.device, dtype=self._logits.dtype
-                )
-            else:
-                # Fallback if neither probs nor logits are set (should not happen with current init logic)
-                total_count = torch.tensor(total_count)
-
         return Binomial(
-            total_count=total_count,
+            total_count=self._total_count,
             probs=self._probs,
             logits=self._logits,
             validate_args=self._validate_args,
         )
 
-    def log_prob(self, value: Tensor) -> Tensor:
-        return self.dist().log_prob(value)
-
     @property
-    def total_count(self) -> Union[int, Tensor]:
+    def total_count(self) -> Tensor:
         """Returns the total_count parameter of the distribution."""
         return self._total_count
 
     @property
-    def probs(self) -> Optional[Tensor]:
+    def probs(self) -> Tensor | None:
         """Returns the probs parameter of the distribution."""
         return self.dist().probs
 
     @property
-    def logits(self) -> Optional[Tensor]:
+    def logits(self) -> Tensor | None:
         """Returns the logits parameter of the distribution."""
         return self.dist().logits
 
     @property
     def param_shape(self) -> Size:
         """Returns the shape of the underlying parameter."""
         return self.dist().param_shape
-
-    @property
-    def mean(self) -> Tensor:
-        """Returns the mean of the Binomial distribution."""
-        return self.dist().mean
-
-    @property
-    def variance(self) -> Tensor:
-        """Returns the variance of the Binomial distribution."""
-        return self.dist().variance
-
-    @property
-    def mode(self) -> Tensor:
-        """Returns the mode of the Binomial distribution."""
-        return self.dist().mode