detailed_balance_loss.hpp

Go to the documentation of this file.

 
#pragma once
 
namespace pulsatrix {
 
class DetailedBalanceLoss {
public:
    DetailedBalanceLoss() = default;
 
    [[nodiscard]] float forward(float log_flow_s, float log_pf, float log_flow_s_next, float log_pb);
 
    [[nodiscard]] float delta() const;
 
    [[nodiscard]] float grad_log_flow_s() const;
 
    [[nodiscard]] float grad_log_flow_s_next() const;
 
    [[nodiscard]] float grad_weight_for_log_pf() const;
 
private:
    float last_delta_ = 0.0f;
    bool has_forwarded_ = false;
 
    void require_forwarded() const;
};
 
}  // namespace pulsatrix