Neural Network Modules

namespace nn

Typedefs

using LSTMState = std::pair<Tensor, Tensor>

Functions

Tensor sinusoidal_position_embedding(int seq_len, int d_model, DType dtype = DType::Float32, Device device = Device::CPU)

class ReLU : public axiom::nn::Module

#include <activation.hpp>

Public Functions

ReLU() = default

virtual Tensor forward(const Tensor &input) const override

inline Tensor operator()(const Tensor &input) const

class GELU : public axiom::nn::Module

#include <activation.hpp>

Public Functions

GELU() = default

virtual Tensor forward(const Tensor &input) const override

inline Tensor operator()(const Tensor &input) const

class SiLU : public axiom::nn::Module

#include <activation.hpp>

Public Functions

SiLU() = default

virtual Tensor forward(const Tensor &input) const override

inline Tensor operator()(const Tensor &input) const

class Sigmoid : public axiom::nn::Module

#include <activation.hpp>

Public Functions

Sigmoid() = default

virtual Tensor forward(const Tensor &input) const override

inline Tensor operator()(const Tensor &input) const

class Tanh : public axiom::nn::Module

#include <activation.hpp>

Public Functions

Tanh() = default

virtual Tensor forward(const Tensor &input) const override

inline Tensor operator()(const Tensor &input) const

class LeakyReLU : public axiom::nn::Module

#include <activation.hpp>

Public Functions

explicit LeakyReLU(float negative_slope = 0.01f)

virtual Tensor forward(const Tensor &input) const override

inline Tensor operator()(const Tensor &input) const

class Dropout : public axiom::nn::Module

#include <activation.hpp>

Public Functions

explicit Dropout(float p = 0.5f)

virtual Tensor forward(const Tensor &input) const override

inline Tensor operator()(const Tensor &input) const

class Flatten : public axiom::nn::Module

#include <activation.hpp>

Public Functions

explicit Flatten(int start_dim = 1, int end_dim = -1)

virtual Tensor forward(const Tensor &input) const override

inline Tensor operator()(const Tensor &input) const

class Upsample : public axiom::nn::Module

#include <activation.hpp>

Public Functions

explicit Upsample(std::vector<size_t> size, InterpolateMode mode = InterpolateMode::Nearest, bool align_corners = false)

explicit Upsample(std::vector<float> scale_factor, InterpolateMode mode = InterpolateMode::Nearest, bool align_corners = false)

virtual Tensor forward(const Tensor &input) const override

inline Tensor operator()(const Tensor &input) const

class MultiHeadAttention : public axiom::nn::Module

#include <attention.hpp>

Public Functions

explicit MultiHeadAttention(int num_heads, bool bias = true)

Tensor forward(const Tensor &query, const Tensor &key, const Tensor &value, const Tensor &mask = Tensor(), bool is_causal = false) const

inline const Linear &q_proj() const

inline const Linear &k_proj() const

inline const Linear &v_proj() const

inline const Linear &out_proj() const

inline int num_heads() const

Tensor forward(const Tensor &input) const

class ModuleList : public axiom::nn::Module

#include <container.hpp>

Public Functions

ModuleList() = default

template<typename T, typename ...Args>
inline T &emplace_back(Args&&... args)

Module &push_back(std::unique_ptr<Module> module)

Module &operator[](size_t index)

const Module &operator[](size_t index) const

size_t size() const

inline auto begin()

inline auto end()

inline auto begin() const

inline auto end() const

template<typename T>
inline TypedModuleRangeImpl<T, true> each() const

template<typename T>
inline TypedModuleRangeImpl<T, false> each()

template<typename T, bool Const>
struct TypedModuleRangeImpl

#include <container.hpp>

Public Types

using Vec = std::vector<std::unique_ptr<Module>>

using VecRef = std::conditional_t<Const, const Vec&, Vec&>

using Ref = std::conditional_t<Const, const T&, T&>

using It = std::conditional_t<Const, typename Vec::const_iterator, typename Vec::iterator>

Public Functions

inline Iterator begin() const

inline Iterator end() const

Public Members

VecRef modules

struct Iterator

#include <container.hpp>

Public Functions

inline Ref operator*() const

inline Iterator &operator++()

inline bool operator!=(const Iterator &o) const

Public Members

It it

class ModuleDict : public axiom::nn::Module

#include <container.hpp>

Public Functions

ModuleDict() = default

template<typename T, typename ...Args>
inline T &emplace(const std::string &key, Args&&... args)

Module &insert(const std::string &key, std::unique_ptr<Module> module)

Module &operator[](const std::string &key)

const Module &operator[](const std::string &key) const

bool contains(const std::string &key) const

std::vector<std::string> keys() const

size_t size() const

template<typename T>
inline T &get(const std::string &key)

template<typename T>
inline const T &get(const std::string &key) const

inline auto begin()

inline auto end()

inline auto begin() const

inline auto end() const

class ParameterDict : public axiom::nn::Module

#include <container.hpp>

Public Functions

ParameterDict() = default

void insert(const std::string &key, const Tensor &param)

Tensor &operator[](const std::string &key)

const Tensor &operator[](const std::string &key) const

bool contains(const std::string &key) const

std::vector<std::string> keys() const

size_t size() const

class Sequential : public axiom::nn::Module

#include <container.hpp>

Public Functions

Sequential() = default

template<typename T, typename ...Args>
inline T &emplace_back(Args&&... args)

virtual Tensor forward(const Tensor &input) const override

inline Tensor operator()(const Tensor &input) const

inline size_t size() const

struct Conv1dConfig

#include <conv.hpp>

Public Members

int stride = 1

int padding = 0

int dilation = 1

int groups = 1

bool bias = true

struct Conv2dConfig

#include <conv.hpp>

Public Members

std::array<int, 2> stride = {1, 1}

std::array<int, 2> padding = {0, 0}

std::array<int, 2> dilation = {1, 1}

int groups = 1

bool bias = true

class Conv1d : public axiom::nn::Module

#include <conv.hpp>

Public Functions

Conv1d(int stride = 1, int padding = 0, int dilation = 1, int groups = 1, bool bias = true)

explicit Conv1d(const Conv1dConfig &config)

virtual Tensor forward(const Tensor &input) const override

inline Tensor operator()(const Tensor &input) const

inline const Tensor &weight() const

inline const Tensor &bias() const

inline int padding() const

inline int groups() const

class Conv2d : public axiom::nn::Module

#include <conv.hpp>

Public Functions

Conv2d(std::array<int, 2> stride = {1, 1}, std::array<int, 2> padding = {0, 0}, std::array<int, 2> dilation = {1, 1}, int groups = 1, bool bias = true)

explicit Conv2d(const Conv2dConfig &config)

virtual Tensor forward(const Tensor &input) const override

inline Tensor operator()(const Tensor &input) const

struct ConvTranspose1dConfig

#include <conv.hpp>

Public Members

int stride = 1

int padding = 0

int output_padding = 0

int dilation = 1

int groups = 1

bool bias = true

struct ConvTranspose2dConfig

#include <conv.hpp>

Public Members

std::array<int, 2> stride = {1, 1}

std::array<int, 2> padding = {0, 0}

std::array<int, 2> output_padding = {0, 0}

std::array<int, 2> dilation = {1, 1}

int groups = 1

bool bias = true

class ConvTranspose1d : public axiom::nn::Module

#include <conv.hpp>

Public Functions

ConvTranspose1d(int stride = 1, int padding = 0, int output_padding = 0, int dilation = 1, int groups = 1, bool bias = true)

explicit ConvTranspose1d(const ConvTranspose1dConfig &config)

virtual Tensor forward(const Tensor &input) const override

inline Tensor operator()(const Tensor &input) const

class ConvTranspose2d : public axiom::nn::Module

#include <conv.hpp>

Public Functions

ConvTranspose2d(std::array<int, 2> stride = {1, 1}, std::array<int, 2> padding = {0, 0}, std::array<int, 2> output_padding = {0, 0}, std::array<int, 2> dilation = {1, 1}, int groups = 1, bool bias = true)

explicit ConvTranspose2d(const ConvTranspose2dConfig &config)

virtual Tensor forward(const Tensor &input) const override

inline Tensor operator()(const Tensor &input) const

class Embedding : public axiom::nn::Module

#include <embedding.hpp>

Public Functions

Embedding()

virtual Tensor forward(const Tensor &indices) const override

inline Tensor operator()(const Tensor &indices) const

class Linear : public axiom::nn::Module

#include <linear.hpp>

Public Functions

explicit Linear(bool bias = true)

virtual Tensor forward(const Tensor &input) const override

inline Tensor operator()(const Tensor &input) const

inline const Tensor &weight() const

inline const Tensor &bias() const

inline bool has_bias() const

class Module

#include <module.hpp>

Subclassed by axiom::nn::AdaptiveAvgPool1d, axiom::nn::AdaptiveAvgPool2d, axiom::nn::AdaptiveMaxPool1d, axiom::nn::AdaptiveMaxPool2d, axiom::nn::AvgPool1d, axiom::nn::AvgPool2d, axiom::nn::BatchNorm1d, axiom::nn::BatchNorm2d, axiom::nn::Conv1d, axiom::nn::Conv2d, axiom::nn::ConvTranspose1d, axiom::nn::ConvTranspose2d, axiom::nn::Dropout, axiom::nn::Embedding, axiom::nn::Flatten, axiom::nn::GELU, axiom::nn::GroupNorm, axiom::nn::InstanceNorm1d, axiom::nn::InstanceNorm2d, axiom::nn::LSTM, axiom::nn::LSTMCell, axiom::nn::LayerNorm, axiom::nn::LeakyReLU, axiom::nn::Linear, axiom::nn::MaxPool1d, axiom::nn::MaxPool2d, axiom::nn::ModuleDict, axiom::nn::ModuleList, axiom::nn::MultiHeadAttention, axiom::nn::ParameterDict, axiom::nn::RMSNorm, axiom::nn::ReLU, axiom::nn::Sequential, axiom::nn::SiLU, axiom::nn::Sigmoid, axiom::nn::Tanh, axiom::nn::Upsample

Public Functions

Module() = default

virtual ~Module() = default

Module(const Module&) = delete

Module &operator=(const Module&) = delete

Module(Module&&) = delete

Module &operator=(Module&&) = delete

virtual Tensor forward(const Tensor &input) const

virtual Module &to(Device device)

virtual Module &to(DType dtype)

virtual void load_state_dict(const std::map<std::string, Tensor> &state_dict, const std::string &prefix = "", bool strict = true)

std::map<std::string, Tensor> state_dict(const std::string &prefix = "") const

std::vector<std::pair<std::string, Tensor*>> named_parameters(const std::string &prefix = "") const

std::vector<Tensor*> parameters() const

class LayerNorm : public axiom::nn::Module

#include <normalization.hpp>

Public Functions

explicit LayerNorm(float eps = 1e-5f)

virtual Tensor forward(const Tensor &input) const override

inline Tensor operator()(const Tensor &input) const

inline const Tensor &weight() const

inline const Tensor &bias() const

inline float eps() const

class RMSNorm : public axiom::nn::Module

#include <normalization.hpp>

Public Functions

explicit RMSNorm(float eps = 1e-5f)

virtual Tensor forward(const Tensor &input) const override

inline Tensor operator()(const Tensor &input) const

class BatchNorm1d : public axiom::nn::Module

#include <normalization.hpp>

Public Functions

explicit BatchNorm1d(float eps = 1e-5f)

virtual Tensor forward(const Tensor &input) const override

inline Tensor operator()(const Tensor &input) const

inline const Tensor &weight() const

inline const Tensor &bias() const

inline const Tensor &running_mean() const

inline const Tensor &running_var() const

inline float eps() const

class BatchNorm2d : public axiom::nn::Module

#include <normalization.hpp>

Public Functions

explicit BatchNorm2d(float eps = 1e-5f)

virtual Tensor forward(const Tensor &input) const override

inline Tensor operator()(const Tensor &input) const

class GroupNorm : public axiom::nn::Module

#include <normalization.hpp>

Public Functions

explicit GroupNorm(int num_groups, float eps = 1e-5f)

virtual Tensor forward(const Tensor &input) const override

inline Tensor operator()(const Tensor &input) const

class InstanceNorm1d : public axiom::nn::Module

#include <normalization.hpp>

Public Functions

explicit InstanceNorm1d(float eps = 1e-5f, bool affine = true)

virtual Tensor forward(const Tensor &input) const override

inline Tensor operator()(const Tensor &input) const

class InstanceNorm2d : public axiom::nn::Module

#include <normalization.hpp>

Public Functions

explicit InstanceNorm2d(float eps = 1e-5f, bool affine = true)

virtual Tensor forward(const Tensor &input) const override

inline Tensor operator()(const Tensor &input) const

class MaxPool1d : public axiom::nn::Module

#include <pooling.hpp>

Public Functions

explicit MaxPool1d(int kernel_size, int stride = 0, int padding = 0)

virtual Tensor forward(const Tensor &input) const override

inline Tensor operator()(const Tensor &input) const

class MaxPool2d : public axiom::nn::Module

#include <pooling.hpp>

Public Functions

explicit MaxPool2d(std::vector<int> kernel_size, std::vector<int> stride = {}, std::vector<int> padding = {})

virtual Tensor forward(const Tensor &input) const override

inline Tensor operator()(const Tensor &input) const

class AvgPool1d : public axiom::nn::Module

#include <pooling.hpp>

Public Functions

explicit AvgPool1d(int kernel_size, int stride = 0, int padding = 0, bool count_include_pad = true)

virtual Tensor forward(const Tensor &input) const override

inline Tensor operator()(const Tensor &input) const

class AvgPool2d : public axiom::nn::Module

#include <pooling.hpp>

Public Functions

explicit AvgPool2d(std::vector<int> kernel_size, std::vector<int> stride = {}, std::vector<int> padding = {}, bool count_include_pad = true)

virtual Tensor forward(const Tensor &input) const override

inline Tensor operator()(const Tensor &input) const

class AdaptiveAvgPool1d : public axiom::nn::Module

#include <pooling.hpp>

Public Functions

explicit AdaptiveAvgPool1d(int output_size)

virtual Tensor forward(const Tensor &input) const override

inline Tensor operator()(const Tensor &input) const

class AdaptiveAvgPool2d : public axiom::nn::Module

#include <pooling.hpp>

Public Functions

explicit AdaptiveAvgPool2d(std::vector<int> output_size)

virtual Tensor forward(const Tensor &input) const override

inline Tensor operator()(const Tensor &input) const

class AdaptiveMaxPool1d : public axiom::nn::Module

#include <pooling.hpp>

Public Functions

explicit AdaptiveMaxPool1d(int output_size)

virtual Tensor forward(const Tensor &input) const override

inline Tensor operator()(const Tensor &input) const

class AdaptiveMaxPool2d : public axiom::nn::Module

#include <pooling.hpp>

Public Functions

explicit AdaptiveMaxPool2d(std::vector<int> output_size)

virtual Tensor forward(const Tensor &input) const override

inline Tensor operator()(const Tensor &input) const

class LSTMCell : public axiom::nn::Module

#include <rnn.hpp>

Public Functions

LSTMCell()

LSTMState forward(const Tensor &input, const LSTMState &state) const

class LSTM : public axiom::nn::Module

#include <rnn.hpp>

Public Functions

explicit LSTM(int num_layers = 1)

Tensor forward(const Tensor &input, std::vector<LSTMState> &states) const

Tensor step(const Tensor &input, std::vector<LSTMState> &states) const

inline int num_layers() const