Only focus on axes

Author: ErikQQY

ext/AdvancedHMCComponentArraysExt.jl

@@ -2,38 +2,9 @@ module AdvancedHMCComponentArraysExt
 
 using LinearAlgebra
 
-using AdvancedHMC:
-    AdvancedHMC,
-    Hamiltonian,
-    UnitEuclideanMetric,
-    DiagEuclideanMetric,
-    DenseEuclideanMetric,
-    GaussianKinetic
-using ComponentArrays: ComponentVecOrMat, ComponentVector, ComponentMatrix, getaxes
+using AdvancedHMC: AdvancedHMC, __axes
+using ComponentArrays: ComponentVecOrMat, getaxes
 
-function AdvancedHMC.∂H∂r(
-    h::Hamiltonian{<:DiagEuclideanMetric,<:GaussianKinetic}, r::ComponentVecOrMat
-)
-    (; M⁻¹) = h.metric
-    (getaxes(M⁻¹) !== getaxes(r)) &&
-        throw(ArgumentError("Axes of mass matrix and momentum must match"))
-    return h.metric.M⁻¹ .* r
-end
-function AdvancedHMC.∂H∂r(
-    h::Hamiltonian{<:DenseEuclideanMetric,<:GaussianKinetic}, r::ComponentVector
-)
-    (; M⁻¹) = h.metric
-    (last(getaxes(M⁻¹)) !== first(getaxes(r))) &&
-        throw(ArgumentError("Axes of mass matrix and momentum must match"))
-    return h.metric.M⁻¹ * r
-end
-function AdvancedHMC.∂H∂r(
-    h::Hamiltonian{<:DenseEuclideanMetric,<:GaussianKinetic}, r::ComponentMatrix
-)
-    (; M⁻¹) = h.metric
-    getaxes(M⁻¹) !== getaxes(r) &&
-        throw(ArgumentError("Axes of mass matrix and momentum must match"))
-    return h.metric.M⁻¹ * r
-end
+AdvancedHMC.__axes(r::ComponentVecOrMat) = getaxes(r)
 
 end # module

src/hamiltonian.jl

@@ -44,13 +44,13 @@ end
 ∂H∂r(h::Hamiltonian{<:UnitEuclideanMetric,<:GaussianKinetic}, r::AbstractVecOrMat) = copy(r)
 function ∂H∂r(h::Hamiltonian{<:DiagEuclideanMetric,<:GaussianKinetic}, r::AbstractVecOrMat)
     (; M⁻¹) = h.metric
-    (first(axes(M⁻¹)) !== first(axes(r))) &&
+    (first(__axes(M⁻¹)) !== first(__axes(r))) &&
         throw(ArgumentError("Axes of mass matrix and momentum must match"))
     return h.metric.M⁻¹ .* r
 end
 function ∂H∂r(h::Hamiltonian{<:DenseEuclideanMetric,<:GaussianKinetic}, r::AbstractVecOrMat)
     (; M⁻¹) = h.metric
-    (last(axes(M⁻¹)) !== first(axes(r))) &&
+    (last(__axes(M⁻¹)) !== first(__axes(r))) &&
         throw(ArgumentError("Axes of mass matrix and momentum must match"))
     return h.metric.M⁻¹ * r
 end

src/utilities.jl

@@ -22,6 +22,13 @@ function _randn(
     return out
 end
 
+"""
+    __axes(r::AbstractVecOrMat)
+
+Return the axes of input `r` where `r` can be generic arrays or custom arrays.
+"""
+@inline __axes(r::AbstractVecOrMat) = axes(r)
+
 """ 
 `rand_coupled` produces coupled randomness given a vector of RNGs. For example, 
 when a vector of RNGs is provided, `rand_coupled` peforms a single `rand` call