ct2rs 0.9.21

Rust bindings for OpenNMT/CTranslate2
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64

#pragma once

#include "activation.h"
#include "op.h"

namespace ctranslate2 {
  namespace ops {

    void apply_bias_and_activation(StorageView& x,
                                   const StorageView* bias,
                                   const ActivationType* activation_type,
                                   const StorageView* residual = nullptr,
                                   const dim_t axis = -1);

    class Gemm : public Op {
    public:
      Gemm(float alpha = 1,
           float beta = 1,
           bool trans_a = false,
           bool trans_b = false,
           bool a_is_packed = false,
           bool b_is_packed = false,
           const ActivationType* activation_type = nullptr);

      void operator()(const StorageView& a,
                      const StorageView& b,
                      StorageView& c,
                      const StorageView* a_shift_compensation = nullptr,
                      const StorageView* bias = nullptr,
                      const StorageView* residual = nullptr) const;

      // Return the packed representation of b, if implemented by the GEMM backend.
      static StorageView pack_b_input(const StorageView& b,
                                      const bool transpose,
                                      const dim_t k,
                                      const dim_t n,
                                      const float alpha);

      // Return the compensation term when s8s8s32 is implemented with u8s8s32.
      static StorageView compensate_u8_input(const StorageView& b,
                                             const bool transpose,
                                             const dim_t k,
                                             const dim_t n,
                                             const float alpha);
    protected:
      const ActivationType* _activation_type;

    private:
      float _alpha;
      float _beta;
      bool _trans_a;
      bool _trans_b;
      bool _a_is_packed;
      bool _b_is_packed;

      template <Device D, typename In, typename Out>
      void compute(const StorageView& a,
                   const StorageView& b,
                   StorageView& c,
                   const StorageView* a_shift_compensation) const;
    };

  }
}