diff --git a/MParT/Distributions/GaussianSamplerDensity.h b/MParT/Distributions/GaussianSamplerDensity.h
index bad4a8f5..b7add81e 100644
--- a/MParT/Distributions/GaussianSamplerDensity.h
+++ b/MParT/Distributions/GaussianSamplerDensity.h
@@ -54,6 +54,8 @@ class GaussianSamplerDensity: public SampleGenerator<MemorySpace>, public Densit
 
     unsigned int Dim() const override { return dim_; };
 
+    bool IsStandardNormal() const { return idCov_ && mean_.size() == 0; }
+
     protected:
     using GeneratorType = typename SampleGenerator<MemorySpace>::PoolType::generator_type;
     using SampleGenerator<MemorySpace>::rand_pool;
diff --git a/MParT/MapObjective.h b/MParT/MapObjective.h
index 86ce36e3..5938a82d 100644
--- a/MParT/MapObjective.h
+++ b/MParT/MapObjective.h
@@ -178,12 +178,44 @@ class KLObjective: public MapObjective<MemorySpace> {
     double ObjectiveImpl(StridedMatrix<const double, MemorySpace> data, std::shared_ptr<ConditionalMapBase<MemorySpace>> map) const override;
     void CoeffGradImpl(StridedMatrix<const double, MemorySpace> data, StridedVector<double, MemorySpace> grad, std::shared_ptr<ConditionalMapBase<MemorySpace>> map) const override;
     unsigned int MapOutputDim() const override {return density_->Dim();}
-    private:
+
     /**
      * @brief Density \f$\mu\f$ to calculate the KL with respect to (i.e. \f$D(\cdot||\mu)\f$ )
      *
      */
-    std::shared_ptr<DensityBase<MemorySpace>> density_;
+    const std::shared_ptr<DensityBase<MemorySpace>> density_;
+};
+
+template<typename MemorySpace>
+class FastGaussianReverseKLObjective: public MapObjective<MemorySpace> {
+    public:
+    FastGaussianReverseKLObjective(StridedMatrix<const double, MemorySpace> train, unsigned int numCoeffs);
+    FastGaussianReverseKLObjective(StridedMatrix<const double, MemorySpace> train, StridedMatrix<const double, MemorySpace> test, unsigned int numCoeffs);
+
+    double ObjectivePlusCoeffGradImpl(StridedMatrix<const double, MemorySpace> data, StridedVector<double, MemorySpace> grad, std::shared_ptr<ConditionalMapBase<MemorySpace>> map) const override;
+
+    double ObjectiveImpl(StridedMatrix<const double, MemorySpace> data, std::shared_ptr<ConditionalMapBase<MemorySpace>> map) const override;
+
+    void CoeffGradImpl(StridedMatrix<const double, MemorySpace> data, StridedVector<double, MemorySpace> grad, std::shared_ptr<ConditionalMapBase<MemorySpace>> map) const override;
+
+    private:
+    using ExecSpace = typename MemoryToExecution<MemorySpace>::Space;
+    enum class ObjectiveType {Eval = 0, Grad = 1, EvalGrad = 2};
+    
+    template<unsigned int Type_idx>
+    void FillSpaces(std::shared_ptr<ConditionalMapBase<MemorySpace>> map, StridedMatrix<const double, MemorySpace> data) const;
+
+    template<unsigned int Type_idx>
+    double CommonEval(StridedMatrix<const double, MemorySpace> data, StridedVector<double, MemorySpace> grad, std::shared_ptr<ConditionalMapBase<MemorySpace>> map) const;
+
+    template<unsigned int Type_idx>
+    void ClearSpaces() const;
+
+    mutable Kokkos::View<double**, MemorySpace> eval_space_;
+    mutable Kokkos::View<double*, MemorySpace> logdet_space_;
+
+    mutable Kokkos::View<double**, MemorySpace> grad_space_;
+    mutable Kokkos::View<double**, MemorySpace> logdet_grad_space_;
 };
 
 namespace ObjectiveFactory {
diff --git a/MParT/TrainMap.h b/MParT/TrainMap.h
index b54758df..e538aec1 100644
--- a/MParT/TrainMap.h
+++ b/MParT/TrainMap.h
@@ -5,6 +5,7 @@
 #include <nlopt.hpp>
 #include <iostream>
 #include "MParT/ConditionalMapBase.h"
+#include "MParT/TriangularMap.h"
 #include "MParT/MapObjective.h"
 
 namespace mpart {
diff --git a/MParT/TriangularMap.h b/MParT/TriangularMap.h
index 27d96877..ef79d6a1 100644
--- a/MParT/TriangularMap.h
+++ b/MParT/TriangularMap.h
@@ -64,6 +64,7 @@ class TriangularMap : public ConditionalMapBase<MemorySpace>{
     void WrapCoeffs(Kokkos::View<double*, MemorySpace> coeffs) override;
 
     virtual std::shared_ptr<ConditionalMapBase<MemorySpace>> GetComponent(unsigned int i){ return comps_.at(i);}
+    unsigned int NumComponents() const { return comps_.size();}
 
     /** @brief Computes the log determinant of the Jacobian matrix of this map.
 
diff --git a/MParT/Utilities/ArrayConversions.h b/MParT/Utilities/ArrayConversions.h
index cc9cde94..24be2ac1 100644
--- a/MParT/Utilities/ArrayConversions.h
+++ b/MParT/Utilities/ArrayConversions.h
@@ -148,9 +148,10 @@ namespace mpart{
     template<typename ScalarType, class... ViewTraits>
     std::vector<typename std::remove_const<ScalarType>::type> KokkosToStd(Kokkos::View<ScalarType*,ViewTraits...> const& view)
     {
+        auto view_copy = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace(), view);
         std::vector<typename std::remove_const<ScalarType>::type> output(view.extent(0));
         for(unsigned int i=0; i<view.extent(0); ++i)
-            output[i] = view(i);
+            output[i] = view_copy(i);
         return output;
     }
 
diff --git a/src/Distributions/GaussianSamplerDensity.cpp b/src/Distributions/GaussianSamplerDensity.cpp
index 22d05196..38120d02 100644
--- a/src/Distributions/GaussianSamplerDensity.cpp
+++ b/src/Distributions/GaussianSamplerDensity.cpp
@@ -30,9 +30,7 @@ void GaussianSamplerDensity<MemorySpace>::LogDensityImpl(StridedMatrix<const dou
     Kokkos::View<double**, Kokkos::LayoutLeft, MemorySpace> diff ("diff", M, N);
 
     if(mean_.extent(0) == 0){
-        Kokkos::parallel_for(policy, KOKKOS_CLASS_LAMBDA(const int& j, const int& i) {
-            diff(i,j) = pts(i,j);
-        });
+        Kokkos::deep_copy(diff, pts);
     }
     else {
         Kokkos::parallel_for(policy, KOKKOS_CLASS_LAMBDA(const int& j, const int& i) {
@@ -75,7 +73,17 @@ void GaussianSamplerDensity<MemorySpace>::LogDensityInputGradImpl(StridedMatrix<
     }
 
     if(!idCov_) {
-        covChol_.solveInPlace(output);
+        
+        if(output.stride_0()==1){
+            covChol_.solveInPlace(output);
+        }else{
+            StridedMatrix<double,MemorySpace> outLeft;
+            outLeft = Kokkos::View<double**, Kokkos::LayoutLeft, MemorySpace>("OutLeft", output.extent(0), output.extent(1));
+            Kokkos::deep_copy(outLeft, output);
+            covChol_.solveInPlace(outLeft);
+            Kokkos::deep_copy(output, outLeft);
+        }
+        
     }
 }
 
diff --git a/src/MapObjective.cpp b/src/MapObjective.cpp
index 50431688..0946a0b3 100644
--- a/src/MapObjective.cpp
+++ b/src/MapObjective.cpp
@@ -6,6 +6,7 @@ template<typename MemorySpace>
 double MapObjective<MemorySpace>::operator()(unsigned int n, const double* coeffs, double* grad, std::shared_ptr<ConditionalMapBase<MemorySpace>> map) {
 
     Kokkos::View<const double*, MemorySpace> coeffView = ToConstKokkos<double,MemorySpace>(coeffs, n);
+    if (grad == nullptr) return ObjectiveImpl(train_, map);
     StridedVector<double, MemorySpace> gradView = ToKokkos<double,MemorySpace>(grad, n);
     map->SetCoeffs(coeffView);
     return ObjectivePlusCoeffGradImpl(train_, gradView, map);
@@ -89,6 +90,149 @@ double KLObjective<MemorySpace>::ObjectivePlusCoeffGradImpl(StridedMatrix<const
     return sumDensity/N_samps;
 }
 
+template<typename MemorySpace>
+FastGaussianReverseKLObjective<MemorySpace>::FastGaussianReverseKLObjective(StridedMatrix<const double, MemorySpace> train, StridedMatrix<const double, MemorySpace> test, unsigned int numCoeffs) : MapObjective<MemorySpace>(train, test) {
+    unsigned int N = std::max(train.extent(1), test.extent(1));
+    eval_space_ = Kokkos::View<double**, MemorySpace>("eval_space", 1, N);
+    logdet_space_ = Kokkos::View<double*, MemorySpace>("logdet_space", N);
+    grad_space_ = Kokkos::View<double**, MemorySpace>("grad_space", numCoeffs, N);
+    logdet_grad_space_ = Kokkos::View<double**, MemorySpace>("logdet_grad_space", numCoeffs, N);
+}
+
+template<typename MemorySpace>
+FastGaussianReverseKLObjective<MemorySpace>::FastGaussianReverseKLObjective(StridedMatrix<const double, MemorySpace> train, unsigned int numCoeffs) : FastGaussianReverseKLObjective<MemorySpace>(train, Kokkos::View<const double**, MemorySpace>(), numCoeffs) {}
+
+template<typename MemorySpace>
+template<unsigned int Type_idx>
+void FastGaussianReverseKLObjective<MemorySpace>::FillSpaces(std::shared_ptr<ConditionalMapBase<MemorySpace>> map, StridedMatrix<const double, MemorySpace> data) const {
+    constexpr ObjectiveType Type = static_cast<ObjectiveType>(Type_idx);
+    unsigned int N_points = data.extent(1);
+    if(N_points > eval_space_.extent(1)){
+        std::stringstream ss;
+        ss << "FastGaussianReverseKLObjective: Not enough space allocated for evaluation!" <<
+            "Need " << N_points << " points, storing " << eval_space_.extent(1) << " points.";
+        throw std::invalid_argument(ss.str().c_str());
+    }
+    StridedMatrix<double, MemorySpace> eval_space_view = Kokkos::subview(eval_space_, Kokkos::ALL(), Kokkos::make_pair(0u, N_points));
+    map->EvaluateImpl(data, eval_space_view);
+    if constexpr((Type == ObjectiveType::Eval) || (Type == ObjectiveType::EvalGrad)) {
+        StridedVector<double, MemorySpace> logdet_space_view = Kokkos::subview(logdet_space_, Kokkos::make_pair(0u, N_points));
+        map->LogDeterminantImpl(data, logdet_space_view);
+    }
+    if constexpr((Type == ObjectiveType::Grad) || (Type == ObjectiveType::EvalGrad)) {
+        StridedMatrix<double, MemorySpace> grad_space_view = Kokkos::subview(grad_space_, Kokkos::ALL(), Kokkos::make_pair(0u, N_points));
+        StridedMatrix<double, MemorySpace> logdet_grad_space_view = Kokkos::subview(logdet_grad_space_, Kokkos::ALL(), Kokkos::make_pair(0u, N_points));
+        Kokkos::fence();
+        map->CoeffGradImpl(data, eval_space_view, grad_space_view);
+        map->LogDeterminantCoeffGradImpl(data, logdet_grad_space_view);
+    }
+    Kokkos::fence();
+}
+
+template<typename MemorySpace>
+template<unsigned int Type_idx>
+void FastGaussianReverseKLObjective<MemorySpace>::ClearSpaces() const {
+    constexpr ObjectiveType Type = static_cast<ObjectiveType>(Type_idx);
+    if constexpr((Type == ObjectiveType::Eval) || (Type == ObjectiveType::EvalGrad)) {
+        Kokkos::deep_copy(eval_space_, 0.);
+        Kokkos::deep_copy(logdet_space_, 0.);
+    }
+    if constexpr((Type == ObjectiveType::Grad) || (Type == ObjectiveType::EvalGrad)) {
+        Kokkos::deep_copy(grad_space_, 0.);
+        Kokkos::deep_copy(logdet_grad_space_, 0.);
+    }
+}
+
+template<typename EvalSpaceType, typename LogDetSpaceType>
+class EvaluationFunctor {
+    public:
+    EvaluationFunctor(const EvalSpaceType& eval_space_view, const LogDetSpaceType& logdet_space_view): eval_space_view_(eval_space_view), logdet_space_view_(logdet_space_view){}
+    KOKKOS_FUNCTION void operator()(const unsigned int j, double& loss) const {
+        loss += 0.5*eval_space_view_(0, j)*eval_space_view_(0, j) - logdet_space_view_(j);
+    }
+    private:
+    const EvalSpaceType eval_space_view_;
+    const LogDetSpaceType logdet_space_view_;
+};
+
+template<typename KLGradSpaceType, typename GradSpaceType, typename GradLogDetSpaceType>
+class GradientFunctor {
+    using team_handle=typename Kokkos::TeamPolicy<typename GradSpaceType::execution_space>::member_type;
+    public:
+    GradientFunctor(KLGradSpaceType& kl_grad, unsigned int N_points, const GradSpaceType& grad_space_view, const GradLogDetSpaceType& logdet_grad_space_view):
+        kl_grad_(kl_grad), N_points_(N_points), grad_space_view_(grad_space_view), logdet_grad_space_view_(logdet_grad_space_view) {
+    }
+    KOKKOS_FUNCTION void operator()(const team_handle& team) const {
+        int d = team.league_rank();
+        double grad_d_ = 0.;
+        Kokkos::parallel_reduce(Kokkos::TeamThreadRange(team, N_points_), [&](const unsigned int j, double& grad_d_tmp) {
+            grad_d_tmp += grad_space_view_(d, j) - logdet_grad_space_view_(d, j);
+        }, grad_d_);
+        kl_grad_(d) = grad_d_/N_points_;
+    }
+    
+    private:
+    mutable KLGradSpaceType kl_grad_;
+    const unsigned int N_points_;
+    const GradSpaceType grad_space_view_;
+    const GradLogDetSpaceType logdet_grad_space_view_;
+};
+
+template<typename MemorySpace>
+template<unsigned int Type_idx>
+double FastGaussianReverseKLObjective<MemorySpace>::CommonEval(StridedMatrix<const double, MemorySpace> points, StridedVector<double, MemorySpace> kl_grad, std::shared_ptr<ConditionalMapBase<MemorySpace>> map) const {
+    if(map->outputDim != 1){
+        throw std::invalid_argument("FastGaussianReverseKLObjective: Map output dimension must be 1! Found dimension " + std::to_string(map->outputDim) + ".");
+    }
+    constexpr ObjectiveType Type = static_cast<ObjectiveType>(Type_idx);
+    unsigned int N_points = points.extent(1);
+    unsigned int numCoeffs = map->numCoeffs;
+    if(Type != ObjectiveType::Eval && numCoeffs != kl_grad.extent(0)){
+        std::stringstream ss;
+        ss << "FastGaussianReverseKLObjective: Gradient vector has incorrect size!"
+           << "Given size " << kl_grad.extent(0) << ", expected size " << numCoeffs
+           << ".";
+        throw std::invalid_argument(ss.str().c_str());
+    }
+    FillSpaces<Type_idx>(map, points);
+
+    double kl_loss = 0.;
+    if constexpr((Type == ObjectiveType::Eval) || (Type == ObjectiveType::EvalGrad)) {
+        Kokkos::RangePolicy<ExecSpace> eval_policy(0, N_points);
+        EvaluationFunctor eval_kernel(eval_space_, logdet_space_);
+        Kokkos::parallel_reduce("FastGaussianReverseKL Evaluate", eval_policy, eval_kernel, kl_loss);
+        kl_loss /= N_points;
+    }
+    if constexpr((Type == ObjectiveType::Grad) || (Type == ObjectiveType::EvalGrad)) {
+        Kokkos::TeamPolicy<ExecSpace> grad_policy(numCoeffs, Kokkos::AUTO());
+        GradientFunctor gradient_functor(kl_grad, N_points, grad_space_, logdet_grad_space_);
+
+        Kokkos::parallel_for("FastGaussianReverseKL Gradient", grad_policy, gradient_functor);
+        Kokkos::fence();
+    }
+    ClearSpaces<Type_idx>();
+    return kl_loss;
+}
+
+template<typename MemorySpace>
+double FastGaussianReverseKLObjective<MemorySpace>::ObjectivePlusCoeffGradImpl(StridedMatrix<const double, MemorySpace> data, StridedVector<double, MemorySpace> grad, std::shared_ptr<ConditionalMapBase<MemorySpace>> map) const {
+    constexpr unsigned int EvalGrad_idx = static_cast<unsigned int>(ObjectiveType::EvalGrad);
+    return CommonEval<EvalGrad_idx>(data, grad, map);
+}
+
+template<typename MemorySpace>
+double FastGaussianReverseKLObjective<MemorySpace>::ObjectiveImpl(StridedMatrix<const double, MemorySpace> data, std::shared_ptr<ConditionalMapBase<MemorySpace>> map) const {
+    Kokkos::View<double*, MemorySpace> grad_holder;
+    constexpr unsigned int Eval_idx = static_cast<unsigned int>(ObjectiveType::Eval);
+    return CommonEval<Eval_idx>(data, grad_holder, map);
+}
+
+template<typename MemorySpace>
+void FastGaussianReverseKLObjective<MemorySpace>::CoeffGradImpl(StridedMatrix<const double, MemorySpace> data, StridedVector<double, MemorySpace> grad, std::shared_ptr<ConditionalMapBase<MemorySpace>> map) const {
+    constexpr unsigned int Grad_idx = static_cast<unsigned int>(ObjectiveType::Grad);
+    CommonEval<Grad_idx>(data, grad, map);
+}
+
 template<typename MemorySpace>
 double KLObjective<MemorySpace>::ObjectiveImpl(StridedMatrix<const double, MemorySpace> data, std::shared_ptr<ConditionalMapBase<MemorySpace>> map) const {
     unsigned int N_samps = data.extent(1);
@@ -107,7 +251,7 @@ void KLObjective<MemorySpace>::CoeffGradImpl(StridedMatrix<const double, MemoryS
     using ExecSpace = typename MemoryToExecution<MemorySpace>::Space;
     unsigned int N_samps = data.extent(1);
     unsigned int grad_dim = grad.extent(0);
-    PullbackDensity<MemorySpace> pullback {map, density_};
+    PullbackDensity<MemorySpace> pullback {map, density_}; 
     StridedMatrix<double, MemorySpace> densityGradX = pullback.LogDensityCoeffGrad(data);
 
     double scale = -1.0/((double) N_samps);
@@ -131,11 +275,13 @@ void KLObjective<MemorySpace>::CoeffGradImpl(StridedMatrix<const double, MemoryS
 // Explicit template instantiation
 template class mpart::MapObjective<Kokkos::HostSpace>;
 template class mpart::KLObjective<Kokkos::HostSpace>;
+template class mpart::FastGaussianReverseKLObjective<Kokkos::HostSpace>;
 template std::shared_ptr<MapObjective<Kokkos::HostSpace>> mpart::ObjectiveFactory::CreateGaussianKLObjective<Kokkos::HostSpace>(StridedMatrix<const double, Kokkos::HostSpace>, unsigned int);
 template std::shared_ptr<MapObjective<Kokkos::HostSpace>> mpart::ObjectiveFactory::CreateGaussianKLObjective<Kokkos::HostSpace>(StridedMatrix<const double, Kokkos::HostSpace>, StridedMatrix<const double, Kokkos::HostSpace>, unsigned int);
 #if defined(MPART_ENABLE_GPU)
     template class mpart::MapObjective<DeviceSpace>;
     template class mpart::KLObjective<DeviceSpace>;
+    template class mpart::FastGaussianReverseKLObjective<DeviceSpace>;
     template std::shared_ptr<MapObjective<DeviceSpace>> mpart::ObjectiveFactory::CreateGaussianKLObjective<DeviceSpace>(StridedMatrix<const double, DeviceSpace>, unsigned int);
     template std::shared_ptr<MapObjective<DeviceSpace>> mpart::ObjectiveFactory::CreateGaussianKLObjective<DeviceSpace>(StridedMatrix<const double, DeviceSpace>, StridedMatrix<const double, DeviceSpace>, unsigned int);
 #endif
diff --git a/src/TrainMap.cpp b/src/TrainMap.cpp
index af74a58b..92d79ccc 100644
--- a/src/TrainMap.cpp
+++ b/src/TrainMap.cpp
@@ -56,25 +56,136 @@ nlopt::opt SetupOptimization(unsigned int dim, TrainOptions options) {
     return opt;
 }
 
-template<>
-double mpart::TrainMap(std::shared_ptr<ConditionalMapBase<Kokkos::HostSpace>> map, std::shared_ptr<MapObjective<Kokkos::HostSpace>> objective, TrainOptions options) {
+template<typename MemorySpace>
+bool isMapObjectiveStandardGaussianKL(std::shared_ptr<MapObjective<MemorySpace>> objective) {
+    std::shared_ptr<KLObjective<MemorySpace>> klObjective = std::dynamic_pointer_cast<KLObjective<MemorySpace>>(objective);
+    if(klObjective == nullptr) return false;
+    std::shared_ptr<DensityBase<MemorySpace>> density = klObjective->density_;
+    std::shared_ptr<GaussianSamplerDensity<MemorySpace>> gaussianDensity = std::dynamic_pointer_cast<GaussianSamplerDensity<MemorySpace>>(density);
+    return gaussianDensity != nullptr && gaussianDensity->IsStandardNormal();
+}
+
+template<typename MemorySpace>
+class NLOptFunctor {
+    public:
+    NLOptFunctor(std::shared_ptr<MapObjective<MemorySpace>> objective, std::shared_ptr<ConditionalMapBase<MemorySpace>> map): objective_(objective), map_(map) {
+        unsigned int numCoeffs = map->numCoeffs;
+        coeff_d_ = Kokkos::View<double*, MemorySpace>("coeff", numCoeffs);
+        grad_d_ = Kokkos::View<double*, MemorySpace>("grad", numCoeffs);
+    }
+
+    double operator()(unsigned n, const double* coeff_ptr, double* grad_ptr) {
+        assert(n == coeff_d_.extent(0));
+        Kokkos::View<const double*, Kokkos::HostSpace> coeff = ToKokkos<const double, Kokkos::HostSpace>(coeff_ptr, n);
+        Kokkos::View<double*, Kokkos::HostSpace> grad = ToKokkos<double, Kokkos::HostSpace>(grad_ptr, n);
+        Kokkos::deep_copy(coeff_d_, coeff);
+        double error;
+        if(grad_ptr != nullptr) {
+            error = (*objective_)(n, coeff_d_.data(), grad_d_.data(), map_);
+            Kokkos::deep_copy(grad, grad_d_);
+            Kokkos::deep_copy(grad_d_, 0);
+        } else {
+            error = (*objective_)(n, coeff_d_.data(), nullptr, map_);
+        }
+        return error;
+    }
+
+    private:
+    Kokkos::View<double*, MemorySpace> coeff_d_;
+    Kokkos::View<double*, MemorySpace> grad_d_;
+    std::shared_ptr<MapObjective<MemorySpace>> objective_;
+    std::shared_ptr<ConditionalMapBase<MemorySpace>> map_;
+};
+
+template<typename MemorySpace>
+std::function<double(unsigned, const double*, double*)> CreateNLOptObjective(std::shared_ptr<MapObjective<MemorySpace>> objective, std::shared_ptr<ConditionalMapBase<MemorySpace>> map) {
+    if constexpr(std::is_same_v<MemorySpace, Kokkos::HostSpace>) {
+        return std::bind(&MapObjective<MemorySpace>::operator(), objective, std::placeholders::_1, std::placeholders::_2, std::placeholders::_3, map);
+    } else {
+        NLOptFunctor<MemorySpace> functor(objective, map);
+        return functor;
+    }
+}
+
+template<typename MemorySpace>
+double TrainComponent_StandardNormalDensity(std::shared_ptr<ConditionalMapBase<MemorySpace>> component, std::shared_ptr<FastGaussianReverseKLObjective<MemorySpace>> objective, TrainOptions options) {
+    if(component->outputDim != 1) {
+        throw std::runtime_error("TrainComponent_StandardNormalDensity: Component must be scalar-valued. Has output dimension " + std::to_string(component->outputDim) + ".");
+    }
+    std::shared_ptr<MapObjective<MemorySpace>> mapObjective = std::dynamic_pointer_cast<MapObjective<MemorySpace>>(objective);
+    assert(mapObjective != nullptr);
+    std::function<double(unsigned, const double*, double*)> functor = CreateNLOptObjective(mapObjective, component);
+
+    // Get the initial guess at the coefficients
+    std::vector<double> compCoeffsStd = KokkosToStd(component->Coeffs());
+
+    // Optimize the map coefficients using NLopt
+    nlopt::opt opt = SetupOptimization(component->numCoeffs, options);
+    opt.set_min_objective(functor_wrapper, reinterpret_cast<void*>(&functor));
+
+    double error = 0.;
+    
+    nlopt::result res = opt.optimize(compCoeffsStd, error);
+    Kokkos::View<const double*, Kokkos::HostSpace> compCoeffsView = VecToKokkos<double,Kokkos::HostSpace>(compCoeffsStd);
+    component->SetCoeffs(compCoeffsView);
+
+    if(options.verbose) {
+        if(res < 0) {
+            std::cerr << "WARNING: Optimization failed: " << MPART_NLOPT_FAILURE_CODES[-res] << std::endl;
+        } else {
+            std::cout << "Optimization result: " << MPART_NLOPT_SUCCESS_CODES[res] << std::endl;
+        }
+        std::cout << "Optimization error: " << error << "\n"
+                  << "Optimization evaluations: " << opt.get_numevals() << std::endl;
+    }
+    return error;
+}
+
+template<typename MemorySpace>
+double TrainMap_Triangular_StandardNormalDensity(std::shared_ptr<TriangularMap<MemorySpace>> map, std::shared_ptr<KLObjective<MemorySpace>> objective, TrainOptions options) {
+    if(options.verbose) {
+        std::cout << "Detected standard normal reference density with Triangular Map" << std::endl;
+    }
+    double total_error = 0.;
+    StridedMatrix<const double, MemorySpace> trainSamples = objective->GetTrain();
+    // Assume each component is a scalar-valued function
+    for(int i = 0; i < map->NumComponents(); i++) {
+        std::shared_ptr<ConditionalMapBase<MemorySpace>> component = map->GetComponent(i);
+        // Create a new objective for each component by slicing the KL objective
+        StridedMatrix<const double, MemorySpace> trainSlice = Kokkos::subview(trainSamples, Kokkos::make_pair(0u, component->inputDim), Kokkos::ALL());
+        std::shared_ptr<FastGaussianReverseKLObjective<MemorySpace>> componentObjective = std::make_shared<FastGaussianReverseKLObjective<MemorySpace>>(trainSlice, component->numCoeffs);
+        if(options.verbose) std::cout << "Training component " << i << "..." << std::endl;
+        total_error += TrainComponent_StandardNormalDensity(component, componentObjective, options);
+        if(options.verbose) {
+            std::cout << "Component " << i << " trained.\n" << std::endl;
+        }
+    }
+    return total_error;
+}
+
+template<typename MemorySpace>
+double mpart::TrainMap(std::shared_ptr<ConditionalMapBase<MemorySpace>> map, std::shared_ptr<MapObjective<MemorySpace>> objective, TrainOptions options) {
     if(map->Coeffs().extent(0) == 0) {
         if(options.verbose) {
             std::cout << "TrainMap: Initializing map coeffs to 1." << std::endl;
         }
         Kokkos::View<double*, Kokkos::HostSpace> coeffs ("Default coeffs", map->numCoeffs);
-	Kokkos::RangePolicy<typename MemoryToExecution<Kokkos::HostSpace>::Space> policy(0,map->numCoeffs);
-        Kokkos::parallel_for("Setting default coeff val", policy, KOKKOS_LAMBDA(const unsigned int i){
-            coeffs(i) = 1.;
-        });
-	Kokkos::View<const double*, Kokkos::HostSpace> constCoeffs = coeffs;
+        Kokkos::deep_copy(coeffs, 1.);
+	    Kokkos::View<const double*, Kokkos::HostSpace> constCoeffs = coeffs;
         map->SetCoeffs(constCoeffs);
     }
+    std::shared_ptr<TriangularMap<MemorySpace>> tri_map = std::dynamic_pointer_cast<TriangularMap<MemorySpace>>(map);
+    bool isMapTriangular = tri_map != nullptr;
+    if(isMapObjectiveStandardGaussianKL(objective) && isMapTriangular) {
+        std::shared_ptr<KLObjective<MemorySpace>> klObjective = std::dynamic_pointer_cast<KLObjective<MemorySpace>>(objective);
+        return TrainMap_Triangular_StandardNormalDensity(tri_map, klObjective, options);
+    }
+
     nlopt::opt opt = SetupOptimization(map->numCoeffs, options);
 
     // Since objective is (rightfully) separate from the map, we use std::bind to create a functor
     // from objective::operator() that keeps the map argument held.
-    std::function<double(unsigned, const double*, double*)> functor = std::bind(&MapObjective<Kokkos::HostSpace>::operator(), objective, std::placeholders::_1, std::placeholders::_2, std::placeholders::_3, map);
+    std::function<double(unsigned, const double*, double*)> functor = CreateNLOptObjective(objective, map);
     opt.set_min_objective(functor_wrapper, reinterpret_cast<void*>(&functor));
 
     // Get the initial guess at the coefficients
@@ -103,3 +214,8 @@ double mpart::TrainMap(std::shared_ptr<ConditionalMapBase<Kokkos::HostSpace>> ma
     }
     return error;
 }
+
+template double mpart::TrainMap<Kokkos::HostSpace>(std::shared_ptr<ConditionalMapBase<Kokkos::HostSpace>> map, std::shared_ptr<MapObjective<Kokkos::HostSpace>> objective, TrainOptions options);
+#if defined(MPART_ENABLE_GPU)
+template double mpart::TrainMap<DeviceSpace>(std::shared_ptr<ConditionalMapBase<DeviceSpace>> map, std::shared_ptr<MapObjective<DeviceSpace>> objective, TrainOptions options);
+#endif
\ No newline at end of file
diff --git a/src/Utilities/LinearAlgebra.cpp b/src/Utilities/LinearAlgebra.cpp
index 9543ab96..ca36eee2 100644
--- a/src/Utilities/LinearAlgebra.cpp
+++ b/src/Utilities/LinearAlgebra.cpp
@@ -86,8 +86,7 @@ template<>
 void Cholesky<Kokkos::HostSpace>::solveInPlace(Kokkos::View<double**,Kokkos::LayoutLeft,Kokkos::HostSpace> x)
 {
     auto eigX = KokkosToMat(x);
-    cholSolver_->matrixL().solveInPlace(eigX);
-    cholSolver_->matrixU().solveInPlace(eigX);
+    cholSolver_->solveInPlace(eigX);
 }
 
 template<>
diff --git a/tests/Distributions/Test_Distributions_Common.cpp b/tests/Distributions/Test_Distributions_Common.cpp
index fec92d61..66bad0b6 100644
--- a/tests/Distributions/Test_Distributions_Common.cpp
+++ b/tests/Distributions/Test_Distributions_Common.cpp
@@ -1,4 +1,5 @@
 #include "Test_Distributions_Common.h"
+using namespace Catch::Matchers;
 
 void TestStandardNormalSamples(StridedMatrix<double, Kokkos::HostSpace> samples) {
     unsigned int dim = samples.extent(0);
@@ -21,10 +22,14 @@ void TestStandardNormalSamples(StridedMatrix<double, Kokkos::HostSpace> samples)
 
     // Check that the mean is zero and the covariance is the identity matrix
     for(int i = 0; i < dim; i++) {
-        REQUIRE(mean(i) == Approx(0.0).margin(mc_margin));
+        CHECK_THAT(mean(i), WithinAbs(0.0, mc_margin));
         for(int j = 0; j < dim; j++) {
             double diag = (double) (i == j);
-            REQUIRE(covar(i, j) - mean(i)*mean(j) == Approx(diag).margin(mc_margin));
+            double cov_ij = covar(i, j) - mean(i)*mean(j);
+            if(i == j)
+                CHECK_THAT(cov_ij, WithinRel(1., mc_margin));
+            else
+                CHECK_THAT(cov_ij, WithinAbs(0., mc_margin));
         }
     }
 
@@ -49,8 +54,8 @@ void TestStandardNormalSamples(StridedMatrix<double, Kokkos::HostSpace> samples)
     double emp_two_std = 0.954499736104;
     double emp_three_std = 0.997300203937;
     for(int i = 0; i < dim; i++) {
-        REQUIRE(in_one_std[i]/(double)N_samp == Approx(emp_one_std).margin(mc_margin));
-        REQUIRE(in_two_std[i]/(double)N_samp == Approx(emp_two_std).margin(mc_margin));
-        REQUIRE(in_three_std[i]/(double)N_samp == Approx(emp_three_std).margin(mc_margin));
+        CHECK_THAT(in_one_std[i]/(double)N_samp, WithinRel(emp_one_std, mc_margin));
+        CHECK_THAT(in_two_std[i]/(double)N_samp, WithinRel(emp_two_std, mc_margin));
+        CHECK_THAT(in_three_std[i]/(double)N_samp, WithinRel(emp_three_std, mc_margin));
     }
 }
\ No newline at end of file
diff --git a/tests/Test_MapObjective.cpp b/tests/Test_MapObjective.cpp
index a89f61ac..7cb982ee 100644
--- a/tests/Test_MapObjective.cpp
+++ b/tests/Test_MapObjective.cpp
@@ -10,6 +10,7 @@
 
 using namespace mpart;
 using namespace Catch;
+using namespace Catch::Matchers;
 
 TEST_CASE( "Test KLMapObjective", "[KLMapObjective]") {
     unsigned int dim = 2;
@@ -39,7 +40,7 @@ TEST_CASE( "Test KLMapObjective", "[KLMapObjective]") {
         double inv_cov_diag = map_scale*map_scale;
         double kl_exact = -std::log(inv_cov_diag) - 1 + inv_cov_diag + map_shift*map_shift*inv_cov_diag;
         kl_exact /= 2.;
-        CHECK(kl_exact == Approx(kl_est).margin(0.5));
+        CHECK_THAT(kl_exact, WithinRel(kl_est, 0.5));
     }
 
     // Setup map for following sections
@@ -57,7 +58,7 @@ TEST_CASE( "Test KLMapObjective", "[KLMapObjective]") {
             map->Coeffs()(i) += fd_step;
             double kl_perturb_i = objective.ObjectiveImpl(reference_samples, map);
             double coeffFD_i = (kl_perturb_i - kl_est)/fd_step;
-            CHECK(coeffFD_i == Approx(coeffGrad(i)).margin(5*fd_step));
+            CHECK_THAT(coeffFD_i, WithinRel(coeffGrad(i), 10*fd_step));
             map->Coeffs()(i) -= fd_step;
         }
     }
@@ -67,9 +68,9 @@ TEST_CASE( "Test KLMapObjective", "[KLMapObjective]") {
         Kokkos::View<double*, Kokkos::HostSpace> coeffGrad ("CoeffGrad of KL Obj", map->numCoeffs);
         objective.CoeffGradImpl(reference_samples, coeffGradRef, map);
         double kl_est = objective.ObjectivePlusCoeffGradImpl(reference_samples, coeffGrad, map);
-        CHECK(kl_est_ref == Approx(kl_est).margin(1e-12));
+        CHECK_THAT(kl_est_ref, WithinRel(kl_est, 1e-12));
         for(int i = 0; i < map->numCoeffs; i++) {
-            CHECK(coeffGradRef(i) == Approx(coeffGrad(i)).margin(1e-12));
+            CHECK_THAT(coeffGradRef(i), WithinRel(coeffGrad(i), 1e-12));
         }
     }
     SECTION("MapObjectiveFunctions") {
@@ -78,19 +79,79 @@ TEST_CASE( "Test KLMapObjective", "[KLMapObjective]") {
         Kokkos::View<double*, Kokkos::HostSpace> coeffGrad ("CoeffGrad of KL Obj", map->numCoeffs);
         double kl_est_ref = objective.ObjectivePlusCoeffGradImpl(train_samples, coeffGradRef, map);
         double kl_est = objective(map->numCoeffs, map->Coeffs().data(), coeffGrad.data(), map);
-        CHECK(kl_est_ref == Approx(kl_est).margin(1e-12));
+        CHECK_THAT(kl_est_ref, WithinRel(kl_est, 1e-12));
         for(int i = 0; i < map->numCoeffs; i++) {
-            CHECK(coeffGradRef(i) == Approx(coeffGrad(i)).margin(1e-12));
+            CHECK_THAT(coeffGradRef(i), WithinRel(coeffGrad(i), 1e-12));
         }
         // TestError
         double test_error_ref = objective.ObjectiveImpl(test_samples, map);
         double test_error = objective.TestError(map);
-        CHECK(test_error_ref == Approx(test_error).margin(1e-12));
+        CHECK_THAT(test_error_ref, WithinRel(test_error, 1e-12));
         // TrainCoeffGrad
         StridedVector<double,Kokkos::HostSpace> trainCoeffGrad = objective.TrainCoeffGrad(map);
         for(int i = 0; i < map->numCoeffs; i++){
-            CHECK(coeffGradRef(i) == Approx(trainCoeffGrad(i)).margin(1e-12));
+            CHECK_THAT(coeffGradRef(i), WithinRel(trainCoeffGrad(i), 1e-12));
         }
     }
 }
 
+TEST_CASE( "Test FastGaussianReverseKLObjective", "[FastGaussianReverseKLObjective]") {
+    unsigned int dim = 1;
+    unsigned int maxOrder = 4;
+    auto map = MapFactory::CreateTriangular<Kokkos::HostSpace>(dim, dim, maxOrder);
+
+    unsigned int seed = 42;
+    unsigned int N_samples = 20000;
+    unsigned int N_testpts = N_samples/5;
+
+    std::shared_ptr<GaussianSamplerDensity<Kokkos::HostSpace>> density = std::make_shared<GaussianSamplerDensity<Kokkos::HostSpace>>(dim);
+    density->SetSeed(seed);
+    StridedMatrix<double, Kokkos::HostSpace> reference_samples = density->Sample(N_samples);
+    StridedMatrix<double, Kokkos::HostSpace> test_samples = Kokkos::subview(reference_samples, Kokkos::ALL, Kokkos::make_pair(0u,N_testpts));
+    StridedMatrix<double, Kokkos::HostSpace> train_samples = Kokkos::subview(reference_samples, Kokkos::ALL, Kokkos::make_pair(N_testpts,N_samples));
+
+    FastGaussianReverseKLObjective<Kokkos::HostSpace> objective {train_samples, test_samples, map->numCoeffs};
+    KLObjective<Kokkos::HostSpace> klObjective {train_samples, test_samples, density};
+    Kokkos::View<double*, Kokkos::HostSpace> coeffGradRef ("Reference CoeffGrad of KL Obj", map->numCoeffs);
+    Kokkos::View<double*, Kokkos::HostSpace> coeffGrad ("CoeffGrad of KL Obj", map->numCoeffs);
+    Kokkos::deep_copy(map->Coeffs(), 0.1);
+
+    SECTION("ObjectivePlusCoeffGradImpl"){
+        double kl_err_ref = klObjective.ObjectivePlusCoeffGradImpl(train_samples, coeffGradRef, map);
+        double kl_err = objective.ObjectivePlusCoeffGradImpl(train_samples, coeffGrad, map);
+        for(int i = 0; i < map->numCoeffs; i++) {
+            CHECK_THAT(coeffGradRef(i), WithinRel(coeffGrad(i), 1e-12));
+        }
+        Kokkos::deep_copy(map->Coeffs(), 0.5);
+        double kl_err_ref2 = klObjective.ObjectivePlusCoeffGradImpl(train_samples, coeffGradRef, map);
+        double kl_err2 = objective.ObjectivePlusCoeffGradImpl(train_samples, coeffGrad, map);
+        double slope = (kl_err_ref2 - kl_err_ref) / (kl_err2 - kl_err);
+        CHECK_THAT(slope, WithinRel(1., 1e-12));
+        double shift = kl_err_ref;
+        for(int i = 0; i < map->numCoeffs; i++) {
+            CHECK_THAT(coeffGradRef(i), WithinRel(coeffGrad(i), 1e-12));
+        }
+        Kokkos::deep_copy(map->Coeffs(), -0.1);
+        double kl_err_ref3 = klObjective.ObjectivePlusCoeffGradImpl(train_samples, coeffGradRef, map);
+        double kl_err3 = objective.ObjectivePlusCoeffGradImpl(train_samples, coeffGrad, map);
+        double pred_ref = (kl_err3-kl_err)*slope + shift;
+        CHECK_THAT(kl_err_ref3, WithinRel(pred_ref, 1e-12));
+        for(int i = 0; i < map->numCoeffs; i++) {
+            CHECK_THAT(coeffGradRef(i), WithinRel(coeffGrad(i), 1e-12));
+        }
+    }
+    SECTION("ObjectiveImpl") {
+        double kl_err_cgrad = objective.ObjectivePlusCoeffGradImpl(train_samples, coeffGrad, map);
+        double kl_err = objective.ObjectiveImpl(train_samples, map);
+        CHECK_THAT(kl_err_cgrad, WithinRel(kl_err, 1e-14));
+    }
+    SECTION("CoeffGradImpl") {
+        Kokkos::deep_copy(coeffGrad, 0.);
+        Kokkos::deep_copy(coeffGradRef, 0.);
+        objective.ObjectivePlusCoeffGradImpl(train_samples, coeffGradRef, map);
+        objective.CoeffGradImpl(train_samples, coeffGrad, map);
+        for(int i = 0; i < map->numCoeffs; i++) {
+            CHECK_THAT(coeffGradRef(i), WithinRel(coeffGrad(i), 1e-14));
+        }
+    }
+}
\ No newline at end of file
diff --git a/tests/Test_TrainMap.cpp b/tests/Test_TrainMap.cpp
index 57e9c142..a22b9a30 100644
--- a/tests/Test_TrainMap.cpp
+++ b/tests/Test_TrainMap.cpp
@@ -1,5 +1,6 @@
 #include <catch2/catch_all.hpp>
 
+#include "MParT/MultiIndices/MultiIndexSet.h"
 #include "MParT/MapFactory.h"
 #include "MParT/MapObjective.h"
 #include "MParT/TrainMap.h"
@@ -12,9 +13,16 @@
 using namespace mpart;
 using namespace Catch;
 
-using MemorySpace = Kokkos::HostSpace;
+#if defined(MPART_ENABLE_GPU)
+    using MemorySpace2 = mpart::DeviceSpace;
+#else
+    using MemorySpace2 = std::false_type;
+#endif
 
-TEST_CASE("Test_TrainMap", "[TrainMap]") {
+TEMPLATE_TEST_CASE("Test_TrainMap", "[TrainMap]", Kokkos::HostSpace, MemorySpace2) {
+// Skip the test if the MemorySpace is false_type, i.e., no device space is available
+if constexpr (!std::is_same_v<TestType, std::false_type>) {
+    using MemorySpace = TestType;
     unsigned int seed = 42;
     unsigned int dim = 2;
     unsigned int numPts = 5000;
@@ -22,7 +30,7 @@ TEST_CASE("Test_TrainMap", "[TrainMap]") {
     auto sampler = std::make_shared<GaussianSamplerDensity<MemorySpace>>(3);
     sampler->SetSeed(seed);
     auto samples = sampler->Sample(numPts);
-    Kokkos::View<double**, Kokkos::HostSpace> targetSamples("targetSamples", 3, numPts);
+    Kokkos::View<double**, MemorySpace> targetSamples("targetSamples", 3, numPts);
     Kokkos::RangePolicy<typename MemoryToExecution<MemorySpace>::Space> policy {0u, numPts};
     Kokkos::parallel_for("Banana", policy, KOKKOS_LAMBDA(const unsigned int i) {
         targetSamples(0,i) = samples(0,i);
@@ -31,45 +39,68 @@ TEST_CASE("Test_TrainMap", "[TrainMap]") {
     });
     unsigned int map_order = 2;
     SECTION("SquareMap") {
-        StridedMatrix<const double, Kokkos::HostSpace> testSamps = Kokkos::subview(targetSamples, Kokkos::make_pair(1u,3u), Kokkos::make_pair(0u, testPts));
-        StridedMatrix<const double, Kokkos::HostSpace> trainSamps = Kokkos::subview(targetSamples, Kokkos::make_pair(1u,3u), Kokkos::make_pair(testPts, numPts));
+        StridedMatrix<const double, MemorySpace> testSamps = Kokkos::subview(targetSamples, Kokkos::make_pair(1u,3u), Kokkos::make_pair(0u, testPts));
+        StridedMatrix<const double, MemorySpace> trainSamps = Kokkos::subview(targetSamples, Kokkos::make_pair(1u,3u), Kokkos::make_pair(testPts, numPts));
         auto obj = ObjectiveFactory::CreateGaussianKLObjective(trainSamps, testSamps);
 
         MapOptions map_options;
-        auto map = MapFactory::CreateTriangular<Kokkos::HostSpace>(dim, dim, map_order, map_options);
+        auto map = MapFactory::CreateTriangular<MemorySpace>(dim, dim, map_order, map_options);
 
         TrainOptions train_options;
         train_options.verbose = 0;
         TrainMap(map, obj, train_options);
-        auto pullback_samples = map->Evaluate(testSamps);
-        TestStandardNormalSamples(pullback_samples);
+        StridedMatrix<double, MemorySpace> pullback_samples = map->Evaluate(testSamps);
+        StridedMatrix<double, Kokkos::HostSpace> pullback_samples_h = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace(), pullback_samples);
+        TestStandardNormalSamples(pullback_samples_h);
     }
     SECTION("ComponentMap") {
-        StridedMatrix<const double, Kokkos::HostSpace> testSamps = Kokkos::subview(targetSamples, Kokkos::make_pair(1u,3u), Kokkos::pair<unsigned int, unsigned int>(0, testPts));
-        StridedMatrix<const double, Kokkos::HostSpace> trainSamps = Kokkos::subview(targetSamples, Kokkos::make_pair(1u,3u), Kokkos::pair<unsigned int, unsigned int>(testPts, numPts));
+        StridedMatrix<const double, MemorySpace> testSamps = Kokkos::subview(targetSamples, Kokkos::make_pair(1u,3u), Kokkos::pair<unsigned int, unsigned int>(0, testPts));
+        StridedMatrix<const double, MemorySpace> trainSamps = Kokkos::subview(targetSamples, Kokkos::make_pair(1u,3u), Kokkos::pair<unsigned int, unsigned int>(testPts, numPts));
         auto obj = ObjectiveFactory::CreateGaussianKLObjective(trainSamps, testSamps, 1);
 
         MapOptions map_options;
-        std::shared_ptr<ConditionalMapBase<Kokkos::HostSpace>> map = MapFactory::CreateComponent<Kokkos::HostSpace>(FixedMultiIndexSet<Kokkos::HostSpace>(dim,map_order), map_options);
+        MultiIndexSet mset = MultiIndexSet::CreateTotalOrder(dim, map_order);
+        FixedMultiIndexSet<MemorySpace> fmset = mset.Fix().ToDevice<MemorySpace>();
+        std::shared_ptr<ConditionalMapBase<MemorySpace>> map = MapFactory::CreateComponent<MemorySpace>(fmset, map_options);
 
         TrainOptions train_options;
         train_options.verbose = 0;
         TrainMap(map, obj, train_options);
-        auto pullback_samples = map->Evaluate(testSamps);
-        TestStandardNormalSamples(pullback_samples);
+        StridedMatrix<double, MemorySpace> pullback_samples = map->Evaluate(testSamps);
+        StridedMatrix<double, Kokkos::HostSpace> pullback_samples_h = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace(), pullback_samples);
+        TestStandardNormalSamples(pullback_samples_h);
     }
+    StridedMatrix<const double, MemorySpace> testSamps = Kokkos::subview(targetSamples, Kokkos::ALL(), Kokkos::pair<unsigned int, unsigned int>(0, testPts));
+    StridedMatrix<const double, MemorySpace> trainSamps = Kokkos::subview(targetSamples, Kokkos::ALL(), Kokkos::pair<unsigned int, unsigned int>(testPts, numPts));
     SECTION("RectangleMap") {
-        StridedMatrix<const double, Kokkos::HostSpace> testSamps = Kokkos::subview(targetSamples, Kokkos::ALL(), Kokkos::pair<unsigned int, unsigned int>(0, testPts));
-        StridedMatrix<const double, Kokkos::HostSpace> trainSamps = Kokkos::subview(targetSamples, Kokkos::ALL(), Kokkos::pair<unsigned int, unsigned int>(testPts, numPts));
         auto obj = ObjectiveFactory::CreateGaussianKLObjective(trainSamps, testSamps, 2);
 
         MapOptions map_options;
-        std::shared_ptr<ConditionalMapBase<Kokkos::HostSpace>> map = MapFactory::CreateTriangular<Kokkos::HostSpace>(dim+1, dim, map_order, map_options);
+        std::shared_ptr<ConditionalMapBase<MemorySpace>> map = MapFactory::CreateTriangular<MemorySpace>(dim+1, dim, map_order, map_options);
 
         TrainOptions train_options;
         train_options.verbose = 0;
         TrainMap(map, obj, train_options);
-        auto pullback_samples = map->Evaluate(testSamps);
-        TestStandardNormalSamples(pullback_samples);
+        StridedMatrix<double, MemorySpace> pullback_samples = map->Evaluate(testSamps);
+        StridedMatrix<double, Kokkos::HostSpace> pullback_samples_h = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace(), pullback_samples);
+        TestStandardNormalSamples(pullback_samples_h);
     }
+    SECTION("RectangleMap_reference") {
+        Kokkos::View<double**, Kokkos::LayoutLeft, MemorySpace> covar("reference", dim, dim);
+        Kokkos::MDRangePolicy<typename MemoryToExecution<MemorySpace>::Space, Kokkos::Rank<2>> policy({0,0}, {dim,dim});
+        Kokkos::parallel_for("Covar", policy, KOKKOS_LAMBDA(const unsigned int i, const unsigned int j) {
+            covar(i,j) = double(i == j);
+        });
+        std::shared_ptr<DensityBase<MemorySpace>> dens = std::make_shared<GaussianSamplerDensity<MemorySpace>>(covar); // Ensures that the KLObjective isn't shortcut-ed
+        std::shared_ptr<MapObjective<MemorySpace>> obj = std::make_shared<KLObjective<MemorySpace>>(trainSamps, testSamps, dens);
+        MapOptions map_options;
+        std::shared_ptr<ConditionalMapBase<MemorySpace>> map = MapFactory::CreateTriangular<MemorySpace>(dim+1, dim, map_order, map_options);
+        TrainOptions train_options;
+        train_options.verbose = 0;
+        TrainMap(map, obj, train_options);
+        StridedMatrix<double, MemorySpace> pullback_samples = map->Evaluate(testSamps);
+        StridedMatrix<double, Kokkos::HostSpace> pullback_samples_h = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace(), pullback_samples);
+        TestStandardNormalSamples(pullback_samples_h);
+    }
+} // End possible skip of GPU test
 }
\ No newline at end of file