Allow gradient on LDF to be called with e.g. views

Author: penelopeysm

HISTORY.md

@@ -1,5 +1,9 @@
 # DynamicPPL Changelog
 
+## 0.39.1
+
+`LogDensityFunction` now allows you to call `logdensity_and_gradient(ldf, x)` with `AbstractVector`s `x` that are not plain Vectors (they will be converted internally before calculating the gradient).
+
 ## 0.39.0
 
 ### Breaking changes

Project.toml

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

src/logdensityfunction.jl

@@ -148,6 +148,8 @@ struct LogDensityFunction{
     F<:Function,
     N<:NamedTuple,
     ADP<:Union{Nothing,DI.GradientPrep},
+    # type of the vector passed to logdensity functions
+    X<:AbstractVector,
 }
     model::M
     adtype::AD
@@ -202,12 +204,17 @@ struct LogDensityFunction{
             typeof(getlogdensity),
             typeof(all_iden_ranges),
             typeof(prep),
+            typeof(x),
         }(
             model, adtype, getlogdensity, all_iden_ranges, all_ranges, prep, dim
         )
     end
 end
 
+function _get_input_vector_type(::LogDensityFunction{T,M,A,G,I,P,X}) where {T,M,A,G,I,P,X}
+    return X
+end
+
 ###################################
 # LogDensityProblems.jl interface #
 ###################################
@@ -265,6 +272,7 @@ end
 function LogDensityProblems.logdensity_and_gradient(
     ldf::LogDensityFunction{Tlink}, params::AbstractVector{<:Real}
 ) where {Tlink}
+    params = convert(_get_input_vector_type(ldf), params)
     return DI.value_and_gradient(
         LogDensityAt{Tlink}(
             ldf.model, ldf._getlogdensity, ldf._iden_varname_ranges, ldf._varname_ranges

test/logdensityfunction.jl

@@ -186,6 +186,25 @@ end
         end
     end
 
+    @testset "logdensity_and_gradient with views" begin
+        # This test ensures that you can call `logdensity_and_gradient` with an array 
+        # type that isn't the same as the one used in the gradient preparation.
+        @model function f()
+            x ~ Normal()
+            return y ~ Normal()
+        end
+        @testset "$adtype" for adtype in test_adtypes
+            x = randn(2)
+            ldf = LogDensityFunction(f(); adtype)
+            logp, grad = LogDensityProblems.logdensity_and_gradient(ldf, x)
+            logp_view, grad_view = LogDensityProblems.logdensity_and_gradient(
+                ldf, (@view x[:])
+            )
+            @test logp == logp_view
+            @test grad == grad_view
+        end
+    end
+
     # Test that various different ways of specifying array types as arguments work with all
     # ADTypes.
     @testset "Array argument types" begin