torch_categorical.h

Go to the documentation of this file.

#ifndef TORCH_CATEGORICAL_H
#define TORCH_CATEGORICAL_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::stats
{
 
    class TorchCategorical final : public TorchDistributionBase
    {
 
 
    public:
 
        TorchCategorical() = default;
 
        TorchCategorical(torch_tensor_t probs, bool do_build_from_logits=false);
 
        ~TorchCategorical() override = 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
#endif