Fix broadcasting issue in ScaleTransform jacobian

Author: ricardoV94

pymc/logprob/transforms.py

@@ -677,7 +677,7 @@ def backward(self, value, *inputs):
 
     def log_jac_det(self, value, *inputs):
         scale = self.transform_args_fn(*inputs)
-        return -pt.log(pt.abs(scale))
+        return -pt.log(pt.abs(pt.broadcast_to(scale, value.shape)))
 
 
 class LogTransform(RVTransform):

tests/logprob/test_transforms.py

@@ -626,8 +626,8 @@ def test_chained_transform():
 
     log_jac_det = ch.log_jac_det(x_val_forward, *x.owner.inputs, scale, loc)
     assert np.isclose(
-        log_jac_det.eval(),
-        -np.log(scale) - np.sum(np.log(x_val_forward - loc)),
+        pt.sum(log_jac_det).eval(),
+        np.sum(-np.log(scale) - np.log(x_val_forward - loc)),
     )