reparameterize.hpp

Go to the documentation of this file.

 
#pragma once
 
#include "pulsatrix/device_backend.hpp"
#include "pulsatrix/tensor.hpp"
 
namespace pulsatrix {
 
struct ReparamGrad {
    Tensor grad_mu;
    Tensor grad_log_sigma;
};
 
class Reparameterize {
public:
    explicit Reparameterize(DeviceBackend* backend);
 
    [[nodiscard]] Tensor forward(const Tensor& mu, const Tensor& log_sigma, const Tensor& epsilon);
 
    [[nodiscard]] ReparamGrad backward(const Tensor& grad_z) const;
 
private:
    DeviceBackend* backend_;
    Tensor last_log_sigma_;
    Tensor last_epsilon_;
    bool has_forwarded_ = false;
};
 
}  // namespace pulsatrix