torch_distribution.h

Go to the documentation of this file.

#ifndef TORCH_DISTRIBUTION_H
#define TORCH_DISTRIBUTION_H
 
#include "cuberl/base/cubeai_config.h"
 
#ifdef USE_PYTORCH
 
#include "cuberl/base/cuberl_types.h"
#include "torch/torch.h"
 
#include <vector>
 
namespace cuberl {
namespace maths {
namespace stats {
 
class TorchDistributionBase
{
 
public:
 
    virtual ~TorchDistributionBase() = default;
 
    virtual torch_tensor_t entropy() const = 0;
 
    virtual torch_tensor_t log_prob(torch_tensor_t value) = 0;
 
    virtual torch_tensor_t sample(c10::ArrayRef<int64_t> sample_shape) = 0;
 
protected:
 
    TorchDistributionBase() = default;
 
    std::vector<int64_t> batch_shape_;
 
    std::vector<int64_t> event_shape_;
 
    std::vector<int64_t> extended_shape(c10::ArrayRef<int64_t> sample_shape);
 
};
 
}
}
}
#endif
#endif // TORCH_DISTRIBUTION_H