WIP: Subclass DimTransform to avoid bugs

Author: ricardoV94

pymc/dims/transforms.py

@@ -11,12 +11,17 @@
 #   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 #   See the License for the specific language governing permissions and
 #   limitations under the License.
+import pytensor.tensor as pt
 import pytensor.xtensor as ptx
 
 from pymc.logprob.transforms import Transform
 
 
-class LogTransform(Transform):
+class DimTransform(Transform):
+    """Base class for transforms that are applied to dim distriubtions."""
+
+
+class LogTransform(DimTransform):
     name = "log"
 
     def forward(self, value, *inputs):
@@ -32,7 +37,7 @@ def log_jac_det(self, value, *inputs):
 log_transform = LogTransform()
 
 
-class LogOddsTransform(Transform):
+class LogOddsTransform(DimTransform):
     name = "logodds"
 
     def backward(self, value, *inputs):
@@ -47,3 +52,44 @@ def log_jac_det(self, value, *inputs):
 
 
 log_odds_transform = LogOddsTransform()
+
+
+class ZeroSumTransform(DimTransform):
+    name = "zerosum"
+
+    def __init__(self, dims: tuple[str, ...]):
+        self.dims = dims
+
+    @staticmethod
+    def extend_dim(array, dim):
+        n = (array.sizes[dim] + 1).astype("floatX")
+        sum_vals = array.sum(dim)
+        norm = sum_vals / (pt.sqrt(n) + n)
+        fill_val = norm - sum_vals / pt.sqrt(n)
+
+        out = ptx.concat([array, fill_val], dim=dim)
+        return out - norm
+
+    @staticmethod
+    def reduct_dim(array, dim):
+        n = array.sizes[dim].astype("floatX")
+        last = array.isel({dim: -1})
+
+        sum_vals = -last * pt.sqrt(n)
+        norm = sum_vals / (pt.sqrt(n) + n)
+        return array.isel({dim: slice(None, -1)}) + norm
+
+    def forward(self, value, *rv_inputs):
+        for dim in self.dims:
+            value = self.reduct_dim(value, dim=dim)
+        return value
+
+    def backward(self, value, *rv_inputs):
+        for dim in self.dims:
+            value = self.extend_dim(value, dim=dim)
+        return value
+
+    def log_jac_det(self, value, *rv_inputs):
+        # Use following once broadcast_like is implemented
+        # as_xtensor(0).broadcast_like(value, exclude=self.dims)`
+        return (value * 0).sum(self.dims)