record TrandformedVariable sample

model/src/pyrenew/metaclass.py

@@ -693,14 +693,17 @@ def __init__(
         self.transforms = transforms
         self.validate()
 
-    def sample(self, **kwargs) -> tuple:
+    def sample(self, record=True, **kwargs) -> tuple:
         """
         Sample method. Call self.base_rv.sample()
         and then apply the transforms specified
         in self.transforms.
 
         Parameters
         ----------
+        record : bool, optional
+            Whether to record the value of the deterministic
+            RandomVariable. Defaults to True.
         **kwargs :
             Keyword arguments passed to self.base_rv.sample()
 
@@ -711,8 +714,7 @@ def sample(self, **kwargs) -> tuple:
         """
 
         untransformed_values = self.base_rv.sample(**kwargs)
-
-        return tuple(
+        transformed_values = tuple(
             SampledValue(
                 t(uv.value),
                 t_start=self.t_start,
@@ -721,6 +723,12 @@ def sample(self, **kwargs) -> tuple:
             for t, uv in zip(self.transforms, untransformed_values)
         )
 
+        if record:
+            for i, tv in enumerate(transformed_values):
+                numpyro.deterministic(f"{self.name}_{i}", tv.value)
+
+        return transformed_values
+
     def sample_length(self):
         """
         Sample length for a transformed