divop.h

Go to the documentation of this file.

#pragma once

#include "compute/operation.h"
#include "tensor/tensor.h"
#include "compute/div/div_internal.h"

namespace magmadnn {

namespace internal {
enum div_op_t {
    TENSOR_DIV_TENSOR,
    SCALAR_DIV_TENSOR,
    TENSOR_DIV_SCALAR,
    VEC_DIV_SCALAR,
    SCALAR_DIV_VEC
};
}   // namespace internal

namespace op {


template <typename T>
class DivOp : public Operation<T> {
public:
    DivOp(Operation<T> *a, Operation<T> *b, bool copy, bool needs_grad);

    
    std::string to_string() { return "( " + a->to_string() + " / " + b->to_string() + " )"; }
protected:
    Tensor<T> *_eval(bool recompute=true);
    Tensor<T> *_grad(Operation<T> *consumer, Operation<T> *var, Tensor<T> *grad);

    Operation<T> *a, *b;
    Tensor<T> *a_tensor, *b_tensor;

    internal::div_op_t op_type;

    bool copy;
};

template <typename T>
DivOp<T>* div(Operation<T> *a, Operation<T> *b, bool copy, bool needs_grad);

} // namespace op
} // namespace magmadnn