matmulop.h

Go to the documentation of this file.

#pragma once
#include <vector>
#include "compute/operation.h"
#include "compute/variable.h"
#include "compute/dot/dotop.h"
#include "compute/transpose/transposeop.h"
#include "tensor/tensor.h"
#include "math/matmul.h"
#include "math/dot.h"
#include "gemm_internal.h"

namespace magmadnn {
namespace op {

template <typename T>
class MatmulOp : public Operation<T> {
public:
    MatmulOp(T alpha, Operation<T>* a, Operation<T>* b, T beta, Operation<T> *c, bool copy=true, bool needs_grad=true);

    
    std::string to_string() { return "(" + a->to_string() + " x " + 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;
    Operation<T> *b;
    Operation<T> *c;

    Tensor<T> *a_tensor;
    Tensor<T> *b_tensor;
    Tensor<T> *c_tensor;

    T alpha;
    T beta;
    bool copy;
};

template <typename T>
MatmulOp<T>* matmul(Operation<T> *a, Operation<T> *b, bool needs_grad=true);

template <typename T>
MatmulOp<T>* matmul(T alpha, Operation<T> *a, Operation<T> *b, T beta, Tensor<T> *c, bool copy=true, bool needs_grad=true);

} // namespace op
} // namespace magmadnn