torch_bernoulli_dist.h

Go to the documentation of this file.

#ifndef TORCH_BERNOULLI_DIST_H
#define TORCH_BERNOULLI_DIST_H
 
 
#include "cuberl/base/cubeai_config.h"
 
#ifdef USE_PYTORCH
 
#include "cuberl/base/cuberl_types.h"
#include "cuberl/maths/statistics/distributions/torch_distribution.h"
 
#include <torch/torch.h>
 
namespace cuberl {
namespace maths {
namespace stats {
 
class TorchBernoulliDist final : public TorchDistributionBase
{
public:
 
    TorchBernoulliDist() = default;
 
    TorchBernoulliDist(torch_tensor_t probs, bool do_build_from_logits=false);
 
    virtual ~TorchBernoulliDist() = default;
 
    virtual torch_tensor_t entropy() const override;
 
    virtual torch_tensor_t log_prob(torch_tensor_t value) override;
 
    virtual torch_tensor_t sample(c10::ArrayRef<int64_t> sample_shape = {})override;
 
 
    void build_from_logits(torch_tensor_t logits);
 
    void build_from_probabilities(torch_tensor_t probs);
 
    torch_tensor_t get_logits()const { return logits_; }
 
    torch_tensor_t get_probs()const { return probs_; }
 
private:
  torch_tensor_t probs_;
  torch_tensor_t logits_;
  torch_tensor_t param_;
  int num_events_;
 
 
};
 
}
}
}
 
#endif // TORCH_BERNOULLI_DIST_H
#endif