Fix Jacobian for dynamic-to-fixed transforms

HISTORY.md

@@ -1,3 +1,7 @@
+# 0.42.5
+
+Fixed `apply_transform_strategy` to return only the target transform's forward log-Jacobian when converting a `DynamicLink` value to a `FixedTransform` target. It previously added the source's forward Jacobian as well, so `getlogjac`, `getlogjoint_internal`, and `getlogprior_internal` were wrong whenever linked values were fed into a fixed-transform strategy. See [#1407](https://github.com/TuringLang/DynamicPPL.jl/issues/1407).
+
 # 0.42.4
 
 `arraydist` on a vector of univariate distributions now builds its `Distributions.Product` through the inner constructor instead of `Product(dists)`, which is deprecated. The outer constructor calls `Base.depwarn`, and that walks a backtrace on every call, so models with an `arraydist` likelihood paid it once per evaluation. The return type is unchanged.

Project.toml

@@ -1,6 +1,6 @@
 name = "DynamicPPL"
 uuid = "366bfd00-2699-11ea-058f-f148b4cae6d8"
-version = "0.42.4"
+version = "0.42.5"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"

src/transformed_values.jl

@@ -303,20 +303,23 @@ Specifically, this function does a number of things:
   Otherwise, either linking or unlinking is applied as necessary. Note that this function
   does not perform vectorisation unless it is needed.
 
-  A table summarising the possible transformations is as follows:
+  A table summarising the possible transformations is as follows, writing `target` for
+  `target_transform(strategy, vn)`:
 
-  | tv.transform isa ...| `target_transform(...) isa DynamicLink` | `target_transform(...) isa Unlink` |
-  |---------------------|---------------------------------|------------------------------------|
-  | `DynamicLink`       | -> `DynamicLink`                | -> `NoTransform`                   |
-  | `Unlink`            | -> `DynamicLink`                | -> `Unlink`                        |
-  | `NoTransform`       | -> `DynamicLink`                | -> `NoTransform`                   |
-  | `FixedTransform`    | errors                          | errors                             |
+  | `tv.transform` isa | target `DynamicLink` | target `Unlink` or `NoTransform` | target `FixedTransform` |
+  |--------------------|----------------------|----------------------------------|-------------------------|
+  | `DynamicLink`      | -> `DynamicLink`     | -> `NoTransform`                 | -> `FixedTransform`     |
+  | `Unlink`           | -> `DynamicLink`     | -> `Unlink`                      | -> `FixedTransform`     |
+  | `NoTransform`      | -> `DynamicLink`     | -> `NoTransform`                 | -> `FixedTransform`     |
+  | `FixedTransform`   | -> `DynamicLink`     | -> `NoTransform`                 | -> `FixedTransform`     |
 
-  Note that, for the last row, when using `FixedTransform` we require that `target_transform`
-  exactly matches the fixed transform, otherwise an error is thrown.
+  When `tv.transform` and `target` are both `FixedTransform`s, they must be equal, otherwise
+  an error is thrown.
 
-- If `vn` is supposed to be linked, calculates the associated log-Jacobian adjustment for
-  the **forward** linking transformation (i.e., from unlinked to linked).
+- Calculates the log-Jacobian adjustment for the **forward** transformation from the raw
+  value to the new internal representation (for instance, from unlinked to linked). This
+  depends only on `target`, never on `tv`'s current transform, and is zero when the new
+  representation is untransformed.
 
 This function returns a tuple of `(raw_value, new_tv, logjac)`.
 
@@ -348,10 +351,7 @@ function apply_transform_strategy(
         fwd_transform = inverse(target.transform)
         transformed_value, logjac = with_logabsdet_jacobian(fwd_transform, raw_value)
         transformed_tv = TransformedValue(transformed_value, target)
-        # TODO: https://github.com/TuringLang/DynamicPPL.jl/issues/1407
-        # Likely should return `logjac` rather than `logjac - inv_logjac`; the sibling
-        # branches all return only the target's forward Jacobian.
-        (raw_value, transformed_tv, logjac - inv_logjac)
+        (raw_value, transformed_tv, logjac)
     else
         error("unknown target transform: $target")
     end

test/contexts/init.jl

@@ -322,6 +322,23 @@ using Test
             ifp = InitFromParams(ps)
             @test ifp.params === ps.params
         end
+
+        @testset "fixing the link transforms preserves the log-Jacobian" begin
+            @model function fixed_transform_model()
+                a ~ Exponential()
+                return b ~ Normal()
+            end
+            model = fixed_transform_model()
+            linked = last(DynamicPPL.init!!(model, VarInfo(), InitFromPrior(), LinkAll()))
+            # Feeding linked values into the equivalent fixed transforms must not change the
+            # log-Jacobian: it belongs to the target transforms, not to the representation
+            # the values arrived in.
+            strategy = WithTransforms(get_fixed_transforms(model, LinkAll()), LinkAll())
+            _, vi = DynamicPPL.init!!(
+                model, VarInfo(), InitFromParams(linked.values), strategy
+            )
+            @test DynamicPPL.getlogjac(vi) ≈ DynamicPPL.getlogjac(linked)
+        end
     end
 end
 

test/transformed_values.jl

@@ -229,7 +229,7 @@ end
         dist = Beta(2, 2)
         vn = @varname(x)
         ft = FixedTransform(Bijectors.VectorBijectors.from_linked_vec(dist))
-        wrong_ft = FixedTransform(
+        other_ft = FixedTransform(
             Bijectors.VectorBijectors.from_linked_vec(InverseGamma(2, 3))
         )
 
@@ -241,6 +241,10 @@ end
 
         # Matching transform should work
         strategy_ok = DynamicPPL.WithTransforms(VarNamedTuple(; x=ft), UnlinkAll())
+        _, _, dynamic_to_fixed_logjac = DynamicPPL.apply_transform_strategy(
+            strategy_ok, TransformedValue(linked_val, DynamicLink()), vn, dist
+        )
+        @test dynamic_to_fixed_logjac ≈ logjac
         new_raw, new_tv, new_logjac = DynamicPPL.apply_transform_strategy(
             strategy_ok, tv, vn, dist
         )
@@ -250,10 +254,32 @@ end
         @test new_logjac ≈ logjac
 
         # Mismatched transform should error
-        strategy_bad = DynamicPPL.WithTransforms(VarNamedTuple(; x=wrong_ft), UnlinkAll())
+        strategy_bad = DynamicPPL.WithTransforms(VarNamedTuple(; x=other_ft), UnlinkAll())
         @test_throws ErrorException DynamicPPL.apply_transform_strategy(
             strategy_bad, tv, vn, dist
         )
+
+        # `other_ft` is not `dist`'s own link transform, so unlike the assertion above
+        # the target's forward log-Jacobian here is distinguishable from the source's.
+        @testset "log-Jacobian depends only on the target transform" begin
+            other_logjac = last(
+                Bijectors.with_logabsdet_jacobian(
+                    Bijectors.inverse(other_ft.transform), raw_val
+                ),
+            )
+            @test !isapprox(other_logjac, logjac)
+
+            strategy = DynamicPPL.WithTransforms(VarNamedTuple(; x=other_ft), UnlinkAll())
+            vec_val = Bijectors.VectorBijectors.to_vec(dist)(raw_val)
+            @testset "from $(get_transform(src))" for src in (
+                TransformedValue(linked_val, DynamicLink()),
+                TransformedValue(vec_val, Unlink()),
+                TransformedValue(raw_val, NoTransform()),
+            )
+                _, _, lj = DynamicPPL.apply_transform_strategy(strategy, src, vn, dist)
+                @test lj ≈ other_logjac
+            end
+        end
     end
 end