deep_biaffine_attn_decoder.h

Go to the documentation of this file.

// Copyright 2021 CEA LIST
// SPDX-FileCopyrightText: 2022 CEA LIST <gael.de-chalendar@cea.fr>
//
// SPDX-License-Identifier: MIT
 
#ifndef DEEPLIMA_SRC_INFERENCE_EIGEN_DEEP_BIAFFINE_ATTN_DECODER_H
#define DEEPLIMA_SRC_INFERENCE_EIGEN_DEEP_BIAFFINE_ATTN_DECODER_H
 
#include <eigen3/Eigen/Dense>
#include "op_base.h"
#include "deeplima/graph_dp/impl/arborescence.h"
 
namespace deeplima
{
namespace eigen_impl
{
 
template<class M=Eigen::MatrixXf, class V=Eigen::VectorXf>
struct params_deep_biaffine_attn_decoder_t : public param_base_t
{
  M m_weight_head;
  V m_bias_head;
  M m_weight_dep;
  V m_bias_dep;
 
  M m_U1;
  V m_u2;
};
 
template<class M, class V, class T>
class Op_DeepBiaffineAttnDecoder : public Op_Base
{
protected:
  struct workbench_t : public Op_Base::workbench_t
  {
    workbench_t()
    {
    }
    virtual ~workbench_t() {}
 
    M output;
  };
 
public:
  typedef V Vector;
  typedef params_deep_biaffine_attn_decoder_t<M, V> params_t;
 
  virtual std::shared_ptr<Op_Base::workbench_t> create_workbench([[maybe_unused]] uint32_t input_size,
                                        [[maybe_unused]] const std::shared_ptr<param_base_t> params,
                                        bool /*precomputed_input=false*/) const override
  {
    assert(input_size > 0);
    assert(nullptr != params);
    // TODO should it be used
    // const params_deep_biaffine_attn_decoder_t<M, V>& layer = *static_cast<const params_t*>(params);
 
    return std::make_shared<workbench_t>();
  }
 
  virtual size_t execute([[maybe_unused]] std::shared_ptr<Op_Base::workbench_t> pwb,
                         const M& input_matrix,
                         const std::shared_ptr<param_base_t> params,
                         const size_t input_begin,
                         const size_t input_end,
                         std::vector<uint32_t>& output)
  {
    // std::cerr << "Op_DeepBiaffineAttnDecoder::execute input begin=" << input_begin << ", end=" << input_end
    //           << ", output size=" << output.size()
    //           << std::endl;
    assert(nullptr != pwb);
    assert(nullptr != params);
    auto layer = std::dynamic_pointer_cast<const params_t>(params);
    // TODO check if layer is empty, signaling that the cast failed
    // TODO should it be used?
    // auto wb = std::dynamic_pointer_cast<workbench_t>(pwb);
 
    const M input = input_matrix.block(0, 0, input_matrix.rows(), input_end - input_begin);
 
    M arc_head = ((layer->m_weight_head * input).colwise() + layer->m_bias_head).transpose();
    elu_inplace(arc_head);
    M arc_dep = (layer->m_weight_dep * input).colwise() + layer->m_bias_dep;
    elu_inplace(arc_dep);
 
    // Head bias: (arc_head * m_u2) depends only on the head, so it is the same
    // for every dependent column. Compute it once and broadcast, instead of
    // recomputing the identical matrix-vector product for every column.
    const V head_bias = arc_head * layer->m_u2; // [n_head]
 
    M logits = (arc_head * layer->m_U1) * arc_dep;
    logits.colwise() += head_bias;
 
    for (Eigen::Index i = 0; i < logits.rows(); ++i)
    {
      Eigen::Index idx = 0;
      /*typename M::Scalar v = */ logits.row(i).maxCoeff(&idx);
      assert(idx >= 0);
      assert(idx < std::numeric_limits<uint32_t>::max());
      output[input_begin + i] = (uint32_t) idx;
      // std::cerr << i << "\t" << idx << std::endl;
    }
    // std::cerr << std::endl;
 
    // TODO comment out that and output
    arborescence<M, uint32_t, typename M::Scalar>(logits, output, input_begin);
 
    return 0;
  }
 
protected:
  inline void elu_inplace(M& m)
  {
    for (Eigen::Index r = 0; r < m.rows(); ++r)
    {
      for (Eigen::Index c = 0; c < m.cols(); ++c)
      {
        if (m(r,c) < 0)
        {
          m(r,c) = expf(m(r,c)) - 1;
        }
      }
    }
  }
};
 
} // namespace eigen_impl
} // namespace deeplima
 
#endif