operator.hpp

#pragma once
#include <dlprim/definitions.hpp>
#include <dlprim/tensor.hpp>
#include <dlprim/context.hpp>

namespace dlprim {
    namespace json { class value; }

    class SharedResource;


    class Operator {
    public:
        Operator(Context const &ctx) : 
            ctx_(ctx),
            mode_(CalculationsMode::predict)
        {
        }

        SharedResource &shared_resource()
        {
            DLPRIM_CHECK(shared_resource_);
            return *shared_resource_;
        }

        void shared_resource(std::shared_ptr<SharedResource> r)
        {
            shared_resource_ = r;
        }

        virtual ~Operator() 
        {
        }
        
        virtual char const *operator_type() const = 0;
        
        virtual void mode(CalculationsMode mode)
        {
            mode_ = mode;
        }

        virtual CalculationsMode mode()
        {
            return mode_;
        }

        Operator(Operator const &) = delete;
        void operator=(Operator const &) = delete;
        Operator(Operator &&) = delete;
        void operator=(Operator &&) = delete;


        virtual bool alias_generator() 
        {
            return false;
        }

        virtual void initialize_params(std::vector<Tensor> &/*parameters*/,ExecutionContext const &/*e*/)
        {
        }

        virtual void setup(std::vector<TensorSpecs> const &in,
                           std::vector<TensorSpecs> &out,
                           std::vector<TensorSpecs> &parameters,
                           size_t &workspace) = 0;

        virtual void reshape(std::vector<Shape> const &in,
                             std::vector<Shape> &out,
                             size_t &workspace) = 0;

                virtual void forward(std::vector<Tensor> &input,
                             std::vector<Tensor> &output,
                             std::vector<Tensor> &parameters,
                             Tensor &workspace,
                             ExecutionContext const &ctx) = 0;
                
        virtual void backward(std::vector<TensorAndGradient> & /*input*/,
                              std::vector<TensorAndGradient> & /*output*/,
                              std::vector<TensorAndGradient> & /*parameters*/,
                              Tensor &/*workspace*/,
                              ExecutionContext const &/*ctx*/)
        {
            throw NotImplementedError("backpropogation is not implemented for " + std::string(operator_type()));
        }

    protected:
        Context ctx_; 
        CalculationsMode mode_; 
        std::shared_ptr<SharedResource> shared_resource_;
    };
   
    std::unique_ptr<Operator> create_by_name(Context &ctx,
                                             std::string const &name,
                                             json::value const &parameters);


} // dlprim