Adding upstream version 3.8.1.upstream/3.8.1upstream

Signed-off-by: Daniel Baumann <daniel.baumann@progress-linux.org>

1 files changed, 1361 insertions, 0 deletions

diff --git a/src/lua/lua_kann.c b/src/lua/lua_kann.c
new file mode 100644
index 0000000..e42fbfb
--- /dev/null
+++ b/src/lua/lua_kann.c
@@ -0,0 +1,1361 @@
+/*
+ * Copyright 2023 Vsevolod Stakhov
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "lua_common.h"
+#include "lua_tensor.h"
+#include "contrib/kann/kann.h"
+
+/***
+ * @module rspamd_kann
+ * `rspamd_kann` is a Lua interface to kann library
+ */
+
+#define KANN_NODE_CLASS "rspamd{kann_node}"
+#define KANN_NETWORK_CLASS "rspamd{kann}"
+
+/* Simple macros to define behaviour */
+#define KANN_LAYER_DEF(name) static int lua_kann_layer_##name(lua_State *L)
+#define KANN_LAYER_INTERFACE(name)   \
+	{                                \
+		#name, lua_kann_layer_##name \
+	}
+
+#define KANN_TRANSFORM_DEF(name) static int lua_kann_transform_##name(lua_State *L)
+#define KANN_TRANSFORM_INTERFACE(name)   \
+	{                                    \
+		#name, lua_kann_transform_##name \
+	}
+
+#define KANN_LOSS_DEF(name) static int lua_kann_loss_##name(lua_State *L)
+#define KANN_LOSS_INTERFACE(name)   \
+	{                               \
+		#name, lua_kann_loss_##name \
+	}
+
+#define KANN_NEW_DEF(name) static int lua_kann_new_##name(lua_State *L)
+#define KANN_NEW_INTERFACE(name)   \
+	{                              \
+		#name, lua_kann_new_##name \
+	}
+
+
+/*
+ * Forwarded declarations
+ */
+static kad_node_t *lua_check_kann_node(lua_State *L, int pos);
+
+/* Layers */
+KANN_LAYER_DEF(input);
+KANN_LAYER_DEF(dense);
+KANN_LAYER_DEF(layernorm);
+KANN_LAYER_DEF(rnn);
+KANN_LAYER_DEF(lstm);
+KANN_LAYER_DEF(gru);
+KANN_LAYER_DEF(conv2d);
+KANN_LAYER_DEF(conv1d);
+KANN_LAYER_DEF(cost);
+
+static luaL_reg rspamd_kann_layers_f[] = {
+	KANN_LAYER_INTERFACE(input),
+	KANN_LAYER_INTERFACE(dense),
+	KANN_LAYER_INTERFACE(layernorm),
+	KANN_LAYER_INTERFACE(rnn),
+	KANN_LAYER_INTERFACE(lstm),
+	KANN_LAYER_INTERFACE(gru),
+	KANN_LAYER_INTERFACE(conv2d),
+	KANN_LAYER_INTERFACE(conv1d),
+	KANN_LAYER_INTERFACE(cost),
+	{NULL, NULL},
+};
+
+/* Transition and composition functions */
+
+/* General transform */
+KANN_TRANSFORM_DEF(add);
+KANN_TRANSFORM_DEF(sub);
+KANN_TRANSFORM_DEF(mul);
+KANN_TRANSFORM_DEF(cmul);
+KANN_TRANSFORM_DEF(matmul);
+
+KANN_TRANSFORM_DEF(square);
+KANN_TRANSFORM_DEF(sigm);
+KANN_TRANSFORM_DEF(tanh);
+KANN_TRANSFORM_DEF(relu);
+KANN_TRANSFORM_DEF(softmax);
+KANN_TRANSFORM_DEF(1minus);
+KANN_TRANSFORM_DEF(exp);
+KANN_TRANSFORM_DEF(log);
+KANN_TRANSFORM_DEF(sin);
+static luaL_reg rspamd_kann_transform_f[] = {
+	KANN_TRANSFORM_INTERFACE(add),
+	KANN_TRANSFORM_INTERFACE(sub),
+	KANN_TRANSFORM_INTERFACE(mul),
+	KANN_TRANSFORM_INTERFACE(cmul),
+	KANN_TRANSFORM_INTERFACE(matmul),
+
+	KANN_TRANSFORM_INTERFACE(square),
+	KANN_TRANSFORM_INTERFACE(sigm),
+	KANN_TRANSFORM_INTERFACE(tanh),
+	KANN_TRANSFORM_INTERFACE(relu),
+	KANN_TRANSFORM_INTERFACE(softmax),
+	KANN_TRANSFORM_INTERFACE(1minus),
+	KANN_TRANSFORM_INTERFACE(exp),
+	KANN_TRANSFORM_INTERFACE(log),
+	KANN_TRANSFORM_INTERFACE(sin),
+	{NULL, NULL},
+};
+
+/* Loss functions */
+KANN_LOSS_DEF(mse);
+KANN_LOSS_DEF(ce_multi);
+KANN_LOSS_DEF(ce_bin);
+KANN_LOSS_DEF(ce_bin_neg);
+KANN_LOSS_DEF(ce_multi_weighted);
+static luaL_reg rspamd_kann_loss_f[] = {
+	KANN_LOSS_INTERFACE(mse),
+	KANN_LOSS_INTERFACE(ce_multi),
+	KANN_LOSS_INTERFACE(ce_bin),
+	KANN_LOSS_INTERFACE(ce_bin_neg),
+	KANN_LOSS_INTERFACE(ce_multi_weighted),
+	{NULL, NULL},
+};
+
+/* Creation functions */
+KANN_NEW_DEF(leaf);
+KANN_NEW_DEF(scalar);
+KANN_NEW_DEF(weight);
+KANN_NEW_DEF(bias);
+KANN_NEW_DEF(weight_conv2d);
+KANN_NEW_DEF(weight_conv1d);
+KANN_NEW_DEF(kann);
+
+static luaL_reg rspamd_kann_new_f[] = {
+	KANN_NEW_INTERFACE(leaf),
+	KANN_NEW_INTERFACE(scalar),
+	KANN_NEW_INTERFACE(weight),
+	KANN_NEW_INTERFACE(bias),
+	KANN_NEW_INTERFACE(weight_conv2d),
+	KANN_NEW_INTERFACE(weight_conv1d),
+	KANN_NEW_INTERFACE(kann),
+	{NULL, NULL},
+};
+
+LUA_FUNCTION_DEF(kann, load);
+LUA_FUNCTION_DEF(kann, destroy);
+LUA_FUNCTION_DEF(kann, save);
+LUA_FUNCTION_DEF(kann, train1);
+LUA_FUNCTION_DEF(kann, apply1);
+
+static luaL_reg rspamd_kann_m[] = {
+	LUA_INTERFACE_DEF(kann, save),
+	LUA_INTERFACE_DEF(kann, train1),
+	LUA_INTERFACE_DEF(kann, apply1),
+	{"__gc", lua_kann_destroy},
+	{NULL, NULL},
+};
+
+static int
+rspamd_kann_table_to_flags(lua_State *L, int table_pos)
+{
+	int result = 0;
+
+	lua_pushvalue(L, table_pos);
+
+	for (lua_pushnil(L); lua_next(L, -2); lua_pop(L, 1)) {
+		int fl = lua_tointeger(L, -1);
+
+		result |= fl;
+	}
+
+	lua_pop(L, 1);
+
+	return result;
+}
+
+static gint
+lua_load_kann(lua_State *L)
+{
+	lua_newtable(L);
+
+	/* Flags */
+	lua_pushstring(L, "flag");
+	lua_newtable(L);
+	lua_pushinteger(L, KANN_F_IN);
+	lua_setfield(L, -2, "in");
+	lua_pushinteger(L, KANN_F_COST);
+	lua_setfield(L, -2, "cost");
+	lua_pushinteger(L, KANN_F_OUT);
+	lua_setfield(L, -2, "out");
+	lua_pushinteger(L, KANN_F_TRUTH);
+	lua_setfield(L, -2, "truth");
+	lua_settable(L, -3);
+
+	/* Cost type */
+	lua_pushstring(L, "cost");
+	lua_newtable(L);
+	/* binary cross-entropy cost, used with sigmoid */
+	lua_pushinteger(L, KANN_C_CEB);
+	lua_setfield(L, -2, "ceb");
+	/* multi-class cross-entropy cost, used with softmax */
+	lua_pushinteger(L, KANN_C_CEM);
+	lua_setfield(L, -2, "cem");
+	/* binary cross-entropy-like cost, used with tanh */
+	lua_pushinteger(L, KANN_C_CEB_NEG);
+	lua_setfield(L, -2, "ceb_neg");
+	lua_pushinteger(L, KANN_C_MSE);
+	lua_setfield(L, -2, "mse");
+	lua_settable(L, -3);
+
+	/* RNN flag */
+	lua_pushstring(L, "rnn");
+	lua_newtable(L);
+	/* apply layer normalization */
+	lua_pushinteger(L, KANN_RNN_NORM);
+	lua_setfield(L, -2, "norm");
+	/* take the initial hidden values as variables */
+	lua_pushinteger(L, KANN_RNN_VAR_H0);
+	lua_setfield(L, -2, "var_h0");
+	lua_settable(L, -3);
+
+	/* Layers */
+	lua_pushstring(L, "layer");
+	lua_newtable(L);
+	luaL_register(L, NULL, rspamd_kann_layers_f);
+	lua_settable(L, -3);
+
+	/* Transforms */
+	lua_pushstring(L, "transform");
+	lua_newtable(L);
+	luaL_register(L, NULL, rspamd_kann_transform_f);
+	lua_settable(L, -3);
+
+	/* Cost */
+	lua_pushstring(L, "loss");
+	lua_newtable(L);
+	luaL_register(L, NULL, rspamd_kann_loss_f);
+	lua_settable(L, -3);
+
+	/* Create functions */
+	lua_pushstring(L, "new");
+	lua_newtable(L);
+	luaL_register(L, NULL, rspamd_kann_new_f);
+	lua_settable(L, -3);
+
+	/* Load ann from memory or file */
+	lua_pushstring(L, "load");
+	lua_pushcfunction(L, lua_kann_load);
+	lua_settable(L, -3);
+
+	return 1;
+}
+
+static kad_node_t *
+lua_check_kann_node(lua_State *L, int pos)
+{
+	void *ud = rspamd_lua_check_udata(L, pos, KANN_NODE_CLASS);
+	luaL_argcheck(L, ud != NULL, pos, "'kann_node' expected");
+	return ud ? *((kad_node_t **) ud) : NULL;
+}
+
+static kann_t *
+lua_check_kann(lua_State *L, int pos)
+{
+	void *ud = rspamd_lua_check_udata(L, pos, KANN_NETWORK_CLASS);
+	luaL_argcheck(L, ud != NULL, pos, "'kann' expected");
+	return ud ? *((kann_t **) ud) : NULL;
+}
+
+void luaopen_kann(lua_State *L)
+{
+	/* Metatables */
+	rspamd_lua_new_class(L, KANN_NODE_CLASS, NULL); /* TODO: add methods */
+	lua_pop(L, 1);                                  /* No need in metatable... */
+	rspamd_lua_new_class(L, KANN_NETWORK_CLASS, rspamd_kann_m);
+	lua_pop(L, 1); /* No need in metatable... */
+	rspamd_lua_add_preload(L, "rspamd_kann", lua_load_kann);
+	lua_settop(L, 0);
+}
+
+/* Layers implementation */
+#define PUSH_KAD_NODE(n)                               \
+	do {                                               \
+		kad_node_t **pt;                               \
+		pt = lua_newuserdata(L, sizeof(kad_node_t *)); \
+		*pt = (n);                                     \
+		rspamd_lua_setclass(L, KANN_NODE_CLASS, -1);   \
+	} while (0)
+
+#define PUSH_KAN_NETWORK(n)                             \
+	do {                                                \
+		kann_t **pn;                                    \
+		pn = lua_newuserdata(L, sizeof(kann_t *));      \
+		*pn = (n);                                      \
+		rspamd_lua_setclass(L, KANN_NETWORK_CLASS, -1); \
+	} while (0)
+
+#define PROCESS_KAD_FLAGS(n, pos)                                                            \
+	do {                                                                                     \
+		int fl = 0;                                                                          \
+		if (lua_type(L, (pos)) == LUA_TTABLE) { fl = rspamd_kann_table_to_flags(L, (pos)); } \
+		else if (lua_type(L, (pos)) == LUA_TNUMBER) {                                        \
+			fl = lua_tointeger(L, (pos));                                                    \
+		}                                                                                    \
+		(n)->ext_flag |= fl;                                                                 \
+	} while (0)
+
+/***
+ * @function kann.layer.input(ninputs[, flags])
+ * Creates an input layer for ANN
+ * @param {int} ninputs number of inputs
+ * @param {table|int} flags optional flags
+ * @return {kann_node} kann node object (should be used to combine ANN)
+*/
+static int
+lua_kann_layer_input(lua_State *L)
+{
+	gint nnodes = luaL_checkinteger(L, 1);
+
+	if (nnodes > 0) {
+		kad_node_t *t;
+
+		t = kann_layer_input(nnodes);
+
+		PROCESS_KAD_FLAGS(t, 2);
+		PUSH_KAD_NODE(t);
+	}
+	else {
+		return luaL_error(L, "invalid arguments, nnodes required");
+	}
+
+	return 1;
+}
+
+/***
+ * @function kann.layer.dense(in, ninputs[, flags])
+ * Creates a dense layer (e.g. for hidden layer)
+ * @param {kann_node} in kann node
+ * @param {int} ninputs number of dense nodes
+ * @param {table|int} flags optional flags
+ * @return {kann_node} kann node object (should be used to combine ANN)
+*/
+static int
+lua_kann_layer_dense(lua_State *L)
+{
+	kad_node_t *in = lua_check_kann_node(L, 1);
+	gint nnodes = luaL_checkinteger(L, 2);
+
+	if (in != NULL && nnodes > 0) {
+		kad_node_t *t;
+
+		t = kann_layer_dense(in, nnodes);
+
+		PROCESS_KAD_FLAGS(t, 3);
+		PUSH_KAD_NODE(t);
+	}
+	else {
+		return luaL_error(L, "invalid arguments, input + nnodes required");
+	}
+
+	return 1;
+}
+
+/***
+ * @function kann.layer.dropout(in, ratio[, flags])
+ * Creates a dropout layer
+ * @param {kann_node} in kann node
+ * @param {float} ratio drop ratio
+ * @param {table|int} flags optional flags
+ * @return {kann_node} kann node object (should be used to combine ANN)
+*/
+static int
+lua_kann_layer_layerdropout(lua_State *L)
+{
+	kad_node_t *in = lua_check_kann_node(L, 1);
+	double r = luaL_checknumber(L, 2);
+
+	if (in != NULL) {
+		kad_node_t *t;
+
+		t = kann_layer_dropout(in, r);
+
+		PROCESS_KAD_FLAGS(t, 3);
+		PUSH_KAD_NODE(t);
+	}
+	else {
+		return luaL_error(L, "invalid arguments, input + rate required");
+	}
+
+	return 1;
+}
+
+/***
+ * @function kann.layer.dropout(in [, flags])
+ * Creates a normalisation layer
+ * @param {kann_node} in kann node
+ * @param {table|int} flags optional flags
+ * @return {kann_node} kann node object (should be used to combine ANN)
+*/
+static int
+lua_kann_layer_layernorm(lua_State *L)
+{
+	kad_node_t *in = lua_check_kann_node(L, 1);
+
+	if (in != NULL) {
+		kad_node_t *t;
+
+		t = kann_layer_layernorm(in);
+
+		PROCESS_KAD_FLAGS(t, 2);
+		PUSH_KAD_NODE(t);
+	}
+	else {
+		return luaL_error(L, "invalid arguments, input required");
+	}
+
+	return 1;
+}
+
+/***
+ * @function kann.layer.rnn(in, nnodes[, rnn_flags, [, flags]])
+ * Creates a recursive NN layer
+ * @param {kann_node} in kann node
+ * @param {int} nnodes number of cells
+ * @param {int} rnnflags rnn flags
+ * @param {table|int} flags optional flags
+ * @return {kann_node} kann node object (should be used to combine ANN)
+*/
+static int
+lua_kann_layer_rnn(lua_State *L)
+{
+	kad_node_t *in = lua_check_kann_node(L, 1);
+	gint nnodes = luaL_checkinteger(L, 2);
+	gint rnnflags = 0;
+
+	if (in != NULL && nnodes > 0) {
+		kad_node_t *t;
+
+		if (lua_type(L, 3) == LUA_TNUMBER) {
+			rnnflags = lua_tointeger(L, 3);
+		}
+
+		t = kann_layer_rnn(in, nnodes, rnnflags);
+
+		PROCESS_KAD_FLAGS(t, 4);
+		PUSH_KAD_NODE(t);
+	}
+	else {
+		return luaL_error(L, "invalid arguments, input + nnodes required");
+	}
+
+	return 1;
+}
+
+/***
+ * @function kann.layer.lstm(in, nnodes[, rnn_flags, [, flags]])
+ * Creates a recursive NN layer using LSTM cells
+ * @param {kann_node} in kann node
+ * @param {int} nnodes number of cells
+ * @param {int} rnnflags rnn flags
+ * @param {table|int} flags optional flags
+ * @return {kann_node} kann node object (should be used to combine ANN)
+*/
+static int
+lua_kann_layer_lstm(lua_State *L)
+{
+	kad_node_t *in = lua_check_kann_node(L, 1);
+	gint nnodes = luaL_checkinteger(L, 2);
+	gint rnnflags = 0;
+
+	if (in != NULL && nnodes > 0) {
+		kad_node_t *t;
+
+		if (lua_type(L, 3) == LUA_TNUMBER) {
+			rnnflags = lua_tointeger(L, 3);
+		}
+
+		t = kann_layer_lstm(in, nnodes, rnnflags);
+
+		PROCESS_KAD_FLAGS(t, 4);
+		PUSH_KAD_NODE(t);
+	}
+	else {
+		return luaL_error(L, "invalid arguments, input + nnodes required");
+	}
+
+	return 1;
+}
+
+/***
+ * @function kann.layer.rnn(in, nnodes[, rnn_flags, [, flags]])
+ * Creates a recursive NN layer using GRU cells
+ * @param {kann_node} in kann node
+ * @param {int} nnodes number of cells
+ * @param {int} rnnflags rnn flags
+ * @param {table|int} flags optional flags
+ * @return {kann_node} kann node object (should be used to combine ANN)
+*/
+static int
+lua_kann_layer_gru(lua_State *L)
+{
+	kad_node_t *in = lua_check_kann_node(L, 1);
+	gint nnodes = luaL_checkinteger(L, 2);
+	gint rnnflags = 0;
+
+	if (in != NULL && nnodes > 0) {
+		kad_node_t *t;
+
+		if (lua_type(L, 3) == LUA_TNUMBER) {
+			rnnflags = lua_tointeger(L, 3);
+		}
+
+		t = kann_layer_gru(in, nnodes, rnnflags);
+
+		PROCESS_KAD_FLAGS(t, 4);
+		PUSH_KAD_NODE(t);
+	}
+	else {
+		return luaL_error(L, "invalid arguments, input + nnodes required");
+	}
+
+	return 1;
+}
+
+/***
+ * @function kann.layer.conv2d(in, n_flt, k_rows, k_cols, stride_rows, stride_cols, pad_rows, pad_columns[, flags])
+ * Creates a 2D convolution layer
+ * @param {kann_node} in kann node
+ * @param {int} n_flt number of filters
+ * @param {int} k_rows kernel rows
+ * @param {int} k_cols kernel columns
+ * @param {int} stride_rows stride rows
+ * @param {int} stride_cols stride columns
+ * @param {int} pad_rows padding rows
+ * @param {int} pad_columns padding columns
+ * @param {table|int} flags optional flags
+ * @return {kann_node} kann node object (should be used to combine ANN)
+*/
+static int
+lua_kann_layer_conv2d(lua_State *L)
+{
+	kad_node_t *in = lua_check_kann_node(L, 1);
+	int n_flt = luaL_checkinteger(L, 2);
+	int k_rows = luaL_checkinteger(L, 3);
+	int k_cols = luaL_checkinteger(L, 4);
+	int stride_r = luaL_checkinteger(L, 5);
+	int stride_c = luaL_checkinteger(L, 6);
+	int pad_r = luaL_checkinteger(L, 7);
+	int pad_c = luaL_checkinteger(L, 8);
+
+	if (in != NULL) {
+		kad_node_t *t;
+		t = kann_layer_conv2d(in, n_flt, k_rows, k_cols, stride_r, stride_c,
+							  pad_r, pad_c);
+
+		PROCESS_KAD_FLAGS(t, 9);
+		PUSH_KAD_NODE(t);
+	}
+	else {
+		return luaL_error(L, "invalid arguments, input, nflt, kx, ky, stridex, stridey, padx, pady are required");
+	}
+
+	return 1;
+}
+
+/***
+ * @function kann.layer.conv1d(in, n_flt, kern_size, stride_size, pad_size[, flags])
+ * Creates 1D convolution layer
+ * @param {kann_node} in kann node
+ * @param {int} n_flt number of filters
+ * @param {int} kern_size kernel rows
+ * @param {int} stride_size stride rows
+ * @param {int} pad_size padding rows
+ * @param {table|int} flags optional flags
+ * @return {kann_node} kann node object (should be used to combine ANN)
+*/
+static int
+lua_kann_layer_conv1d(lua_State *L)
+{
+	kad_node_t *in = lua_check_kann_node(L, 1);
+	int n_flt = luaL_checkinteger(L, 2);
+	int k_size = luaL_checkinteger(L, 3);
+	int stride = luaL_checkinteger(L, 4);
+	int pad = luaL_checkinteger(L, 5);
+
+	if (in != NULL) {
+		kad_node_t *t;
+		t = kann_layer_conv1d(in, n_flt, k_size, stride, pad);
+
+		PROCESS_KAD_FLAGS(t, 6);
+		PUSH_KAD_NODE(t);
+	}
+	else {
+		return luaL_error(L, "invalid arguments, input, nflt, k, stride, pad required");
+	}
+
+	return 1;
+}
+
+/***
+ * @function kann.layer.cost(in, nout, cost_type[, flags])
+ * Creates 1D convolution layer
+ * @param {kann_node} in kann node
+ * @param {int} nout number of outputs
+ * @param {int} cost_type see kann.cost table
+ * @param {table|int} flags optional flags
+ * @return {kann_node} kann node object (should be used to combine ANN)
+*/
+static int
+lua_kann_layer_cost(lua_State *L)
+{
+	kad_node_t *in = lua_check_kann_node(L, 1);
+	int nout = luaL_checkinteger(L, 2);
+	int cost_type = luaL_checkinteger(L, 3);
+
+	if (in != NULL && nout > 0) {
+		kad_node_t *t;
+		t = kann_layer_cost(in, nout, cost_type);
+
+		PROCESS_KAD_FLAGS(t, 4);
+		PUSH_KAD_NODE(t);
+	}
+	else {
+		return luaL_error(L, "invalid arguments, input, nout and cost_type are required");
+	}
+
+	return 1;
+}
+
+/* Generic helpers */
+static int
+lua_kann_call_unary_function(lua_State *L, const char *name,
+							 kad_node_t *(*func)(kad_node_t *) )
+{
+	kad_node_t *in = lua_check_kann_node(L, 1);
+
+	if (in != NULL) {
+		kad_node_t *t;
+		t = func(in);
+
+		PUSH_KAD_NODE(t);
+	}
+	else {
+		return luaL_error(L, "invalid arguments for %s, input required", name);
+	}
+
+	return 1;
+}
+static int
+lua_kann_call_binary_function(lua_State *L, const char *name,
+							  kad_node_t *(*func)(kad_node_t *, kad_node_t *) )
+{
+	kad_node_t *x = lua_check_kann_node(L, 1);
+	kad_node_t *y = lua_check_kann_node(L, 2);
+
+	if (x != NULL && y != NULL) {
+		kad_node_t *t;
+		t = func(x, y);
+
+		PUSH_KAD_NODE(t);
+	}
+	else {
+		return luaL_error(L, "invalid arguments for %s, 2 inputs required", name);
+	}
+
+	return 1;
+}
+
+#define LUA_UNARY_TRANSFORM_FUNC_IMPL(name)                        \
+	static int lua_kann_transform_##name(lua_State *L)             \
+	{                                                              \
+		return lua_kann_call_unary_function(L, #name, kad_##name); \
+	}
+
+#define LUA_BINARY_TRANSFORM_FUNC_IMPL(name)                        \
+	static int lua_kann_transform_##name(lua_State *L)              \
+	{                                                               \
+		return lua_kann_call_binary_function(L, #name, kad_##name); \
+	}
+
+#define LUA_LOSS_FUNC_IMPL(name)                                    \
+	static int lua_kann_loss_##name(lua_State *L)                   \
+	{                                                               \
+		return lua_kann_call_binary_function(L, #name, kad_##name); \
+	}
+
+/* Transform functions registered via macro helpers */
+LUA_BINARY_TRANSFORM_FUNC_IMPL(add)
+LUA_BINARY_TRANSFORM_FUNC_IMPL(sub)
+LUA_BINARY_TRANSFORM_FUNC_IMPL(mul)
+LUA_BINARY_TRANSFORM_FUNC_IMPL(cmul)
+LUA_BINARY_TRANSFORM_FUNC_IMPL(matmul)
+
+LUA_UNARY_TRANSFORM_FUNC_IMPL(square)
+LUA_UNARY_TRANSFORM_FUNC_IMPL(sigm)
+LUA_UNARY_TRANSFORM_FUNC_IMPL(tanh)
+LUA_UNARY_TRANSFORM_FUNC_IMPL(relu)
+LUA_UNARY_TRANSFORM_FUNC_IMPL(softmax)
+LUA_UNARY_TRANSFORM_FUNC_IMPL(1minus)
+LUA_UNARY_TRANSFORM_FUNC_IMPL(exp)
+LUA_UNARY_TRANSFORM_FUNC_IMPL(log)
+LUA_UNARY_TRANSFORM_FUNC_IMPL(sin)
+
+/* Generic cost functions */
+LUA_LOSS_FUNC_IMPL(mse)
+LUA_LOSS_FUNC_IMPL(ce_multi)
+LUA_LOSS_FUNC_IMPL(ce_bin)
+LUA_LOSS_FUNC_IMPL(ce_bin_neg)
+
+/* The only case of ternary weight function */
+static int
+lua_kann_loss_ce_multi_weighted(lua_State *L)
+{
+	kad_node_t *pred = lua_check_kann_node(L, 1);
+	kad_node_t *truth = lua_check_kann_node(L, 2);
+	kad_node_t *weight = lua_check_kann_node(L, 3);
+
+	if (pred != NULL && truth != NULL && weight != NULL) {
+		kad_node_t *t;
+		t = kad_ce_multi_weighted(pred, truth, weight);
+
+		PUSH_KAD_NODE(t);
+	}
+	else {
+		return luaL_error(L, "invalid arguments for ce_multi_weighted, 3 inputs required");
+	}
+
+	return 1;
+}
+
+/* Creation functions */
+static int
+lua_kann_new_scalar(lua_State *L)
+{
+	gint flag = luaL_checkinteger(L, 1);
+	double x = luaL_checknumber(L, 2);
+	kad_node_t *t;
+
+	t = kann_new_scalar(flag, x);
+
+	PROCESS_KAD_FLAGS(t, 3);
+	PUSH_KAD_NODE(t);
+
+	return 1;
+}
+
+static int
+lua_kann_new_weight(lua_State *L)
+{
+	gint nrow = luaL_checkinteger(L, 1);
+	gint ncol = luaL_checkinteger(L, 2);
+	kad_node_t *t;
+
+	t = kann_new_weight(nrow, ncol);
+
+	PROCESS_KAD_FLAGS(t, 3);
+	PUSH_KAD_NODE(t);
+
+	return 1;
+}
+
+static int
+lua_kann_new_bias(lua_State *L)
+{
+	gint n = luaL_checkinteger(L, 1);
+	kad_node_t *t;
+
+	t = kann_new_bias(n);
+
+	PROCESS_KAD_FLAGS(t, 2);
+	PUSH_KAD_NODE(t);
+
+	return 1;
+}
+
+static int
+lua_kann_new_weight_conv2d(lua_State *L)
+{
+	gint nout = luaL_checkinteger(L, 1);
+	gint nin = luaL_checkinteger(L, 2);
+	gint krow = luaL_checkinteger(L, 3);
+	gint kcol = luaL_checkinteger(L, 4);
+	kad_node_t *t;
+
+	t = kann_new_weight_conv2d(nout, nin, krow, kcol);
+
+	PROCESS_KAD_FLAGS(t, 5);
+	PUSH_KAD_NODE(t);
+
+	return 1;
+}
+
+static int
+lua_kann_new_weight_conv1d(lua_State *L)
+{
+	gint nout = luaL_checkinteger(L, 1);
+	gint nin = luaL_checkinteger(L, 2);
+	gint klen = luaL_checkinteger(L, 3);
+	kad_node_t *t;
+
+	t = kann_new_weight_conv1d(nout, nin, klen);
+
+	PROCESS_KAD_FLAGS(t, 4);
+	PUSH_KAD_NODE(t);
+
+	return 1;
+}
+
+static int
+lua_kann_new_leaf(lua_State *L)
+{
+	int dim = luaL_checkinteger(L, 1), i, *ar;
+	kad_node_t *t;
+
+	if (dim >= 1 && dim < KAD_MAX_DIM && lua_istable(L, 2)) {
+		ar = g_new0(int, KAD_MAX_DIM);
+
+		for (i = 0; i < dim; i++) {
+			lua_rawgeti(L, 2, i + 1);
+			ar[i] = lua_tointeger(L, -1);
+			lua_pop(L, 1);
+		}
+
+		t = kann_new_leaf_array(NULL, NULL, 0, 0.0, dim, ar);
+
+		PROCESS_KAD_FLAGS(t, 3);
+		PUSH_KAD_NODE(t);
+
+		g_free(ar);
+	}
+	else {
+		return luaL_error(L, "invalid arguments for new.leaf, "
+							 "dim and vector of elements are required");
+	}
+
+	return 1;
+}
+
+static int
+lua_kann_new_kann(lua_State *L)
+{
+	kad_node_t *cost = lua_check_kann_node(L, 1);
+	kann_t *k;
+
+	if (cost) {
+		k = kann_new(cost, 0);
+
+		PUSH_KAN_NETWORK(k);
+	}
+	else {
+		return luaL_error(L, "invalid arguments for new.kann, "
+							 "cost node is required");
+	}
+
+	return 1;
+}
+
+static int
+lua_kann_destroy(lua_State *L)
+{
+	kann_t *k = lua_check_kann(L, 1);
+
+	kann_delete(k);
+
+	return 0;
+}
+
+static int
+lua_kann_save(lua_State *L)
+{
+	kann_t *k = lua_check_kann(L, 1);
+
+	if (k) {
+		if (lua_istable(L, 2)) {
+			lua_getfield(L, 2, "filename");
+
+			if (lua_isstring(L, -1)) {
+				const gchar *fname = lua_tostring(L, -1);
+				FILE *f;
+
+				f = fopen(fname, "w");
+
+				if (!f) {
+					lua_pop(L, 1);
+
+					return luaL_error(L, "cannot open %s for writing: %s",
+									  fname, strerror(errno));
+				}
+
+				kann_save_fp(f, k);
+				fclose(f);
+
+				lua_pushboolean(L, true);
+			}
+			else {
+				lua_pop(L, 1);
+
+				return luaL_error(L, "invalid arguments: missing filename");
+			}
+
+			lua_pop(L, 1);
+		}
+		else {
+			/* Save to Rspamd text */
+#ifndef HAVE_OPENMEMSTREAM
+			return luaL_error(L, "no support of saving to memory on your system");
+#endif
+			FILE *f;
+			char *buf = NULL;
+			size_t buflen;
+			struct rspamd_lua_text *t;
+
+			f = open_memstream(&buf, &buflen);
+			g_assert(f != NULL);
+
+			kann_save_fp(f, k);
+			fclose(f);
+
+			t = lua_newuserdata(L, sizeof(*t));
+			rspamd_lua_setclass(L, "rspamd{text}", -1);
+			t->flags = RSPAMD_TEXT_FLAG_OWN;
+			t->start = (const gchar *) buf;
+			t->len = buflen;
+		}
+	}
+	else {
+		return luaL_error(L, "invalid arguments");
+	}
+
+	return 1;
+}
+
+static int
+lua_kann_load(lua_State *L)
+{
+	kann_t *k;
+	FILE *f = NULL;
+
+	if (lua_istable(L, 1)) {
+		lua_getfield(L, 2, "filename");
+
+		if (lua_isstring(L, -1)) {
+			const gchar *fname = lua_tostring(L, -1);
+
+			f = fopen(fname, "rb");
+		}
+		else {
+			lua_pop(L, 1);
+
+			return luaL_error(L, "invalid arguments: missing filename");
+		}
+
+		lua_pop(L, 1);
+	}
+	else if (lua_isstring(L, 1)) {
+		gsize dlen;
+		const gchar *data;
+
+		data = lua_tolstring(L, 1, &dlen);
+
+#ifndef HAVE_FMEMOPEN
+		return luaL_error(L, "no support of loading from memory on your system");
+#endif
+		f = fmemopen((void *) data, dlen, "rb");
+	}
+	else if (lua_isuserdata(L, 1)) {
+		struct rspamd_lua_text *t;
+
+		t = lua_check_text(L, 1);
+
+		if (!t) {
+			return luaL_error(L, "invalid arguments");
+		}
+
+#ifndef HAVE_FMEMOPEN
+		return luaL_error(L, "no support of loading from memory on your system");
+#endif
+		f = fmemopen((void *) t->start, t->len, "rb");
+	}
+
+	if (f == NULL) {
+		return luaL_error(L, "invalid arguments or cannot open file");
+	}
+
+	k = kann_load_fp(f);
+	fclose(f);
+
+	if (k == NULL) {
+		lua_pushnil(L);
+	}
+	else {
+		PUSH_KAN_NETWORK(k);
+	}
+
+	return 1;
+}
+
+struct rspamd_kann_train_cbdata {
+	lua_State *L;
+	kann_t *k;
+	gint cbref;
+};
+
+static void
+lua_kann_train_cb(int iter, float train_cost, float val_cost, void *ud)
+{
+	struct rspamd_kann_train_cbdata *cbd = (struct rspamd_kann_train_cbdata *) ud;
+
+	if (cbd->cbref != -1) {
+		gint err_idx;
+		lua_State *L = cbd->L;
+
+		lua_pushcfunction(L, &rspamd_lua_traceback);
+		err_idx = lua_gettop(L);
+
+		lua_rawgeti(L, LUA_REGISTRYINDEX, cbd->cbref);
+		lua_pushinteger(L, iter);
+		lua_pushnumber(L, train_cost);
+		lua_pushnumber(L, val_cost);
+
+		if (lua_pcall(L, 3, 0, err_idx) != 0) {
+			msg_err("cannot run lua train callback: %s",
+					lua_tostring(L, -1));
+		}
+
+		lua_settop(L, err_idx - 1);
+	}
+}
+
+#define FREE_VEC(a, n)                                \
+	do {                                              \
+		for (int i = 0; i < (n); i++) g_free((a)[i]); \
+		g_free(a);                                    \
+	} while (0)
+
+static int
+lua_kann_train1(lua_State *L)
+{
+	kann_t *k = lua_check_kann(L, 1);
+	struct rspamd_lua_tensor *pca = NULL;
+
+	/* Default train params */
+	double lr = 0.001;
+	gint64 mini_size = 64;
+	gint64 max_epoch = 25;
+	gint64 max_drop_streak = 10;
+	double frac_val = 0.1;
+	gint cbref = -1;
+
+	if (k && lua_istable(L, 2) && lua_istable(L, 3)) {
+		int n = rspamd_lua_table_size(L, 2);
+		int n_in = kann_dim_in(k);
+		int n_out = kann_dim_out(k);
+
+		if (n_in <= 0) {
+			return luaL_error(L, "invalid inputs count: %d", n_in);
+		}
+
+		if (n_out <= 0) {
+			return luaL_error(L, "invalid outputs count: %d", n_out);
+		}
+
+		if (n != rspamd_lua_table_size(L, 3) || n == 0) {
+			return luaL_error(L, "invalid dimensions: outputs size must be "
+								 "equal to inputs and non zero");
+		}
+
+		if (lua_istable(L, 4)) {
+			GError *err = NULL;
+
+			if (!rspamd_lua_parse_table_arguments(L, 4, &err,
+												  RSPAMD_LUA_PARSE_ARGUMENTS_IGNORE_MISSING,
+												  "lr=N;mini_size=I;max_epoch=I;max_drop_streak=I;frac_val=N;cb=F;pca=u{tensor}",
+												  &lr, &mini_size, &max_epoch, &max_drop_streak, &frac_val, &cbref, &pca)) {
+				n = luaL_error(L, "invalid params: %s",
+							   err ? err->message : "unknown error");
+				g_error_free(err);
+
+				return n;
+			}
+		}
+
+		if (pca) {
+			/* Check pca matrix validity */
+			if (pca->ndims != 2) {
+				return luaL_error(L, "invalid pca tensor: matrix expected, got a row");
+			}
+
+			if (pca->dim[0] != n_in) {
+				return luaL_error(L, "invalid pca tensor: "
+									 "matrix must have %d rows and it has %d rows instead",
+								  n_in, pca->dim[0]);
+			}
+		}
+
+		float **x, **y, *tmp_row = NULL;
+
+		/* Fill vectors row by row */
+		x = (float **) g_malloc0(sizeof(float *) * n);
+		y = (float **) g_malloc0(sizeof(float *) * n);
+
+		if (pca) {
+			tmp_row = g_malloc(sizeof(float) * pca->dim[1]);
+		}
+
+		for (int s = 0; s < n; s++) {
+			/* Inputs */
+			lua_rawgeti(L, 2, s + 1);
+			x[s] = (float *) g_malloc(sizeof(float) * n_in);
+
+			if (pca == NULL) {
+				if (rspamd_lua_table_size(L, -1) != n_in) {
+					FREE_VEC(x, n);
+					FREE_VEC(y, n);
+
+					n = luaL_error(L, "invalid params at pos %d: "
+									  "bad input dimension %d; %d expected",
+								   s + 1,
+								   (int) rspamd_lua_table_size(L, -1),
+								   n_in);
+					lua_pop(L, 1);
+
+					return n;
+				}
+
+				for (int i = 0; i < n_in; i++) {
+					lua_rawgeti(L, -1, i + 1);
+					x[s][i] = lua_tonumber(L, -1);
+
+					lua_pop(L, 1);
+				}
+			}
+			else {
+				if (rspamd_lua_table_size(L, -1) != pca->dim[1]) {
+					FREE_VEC(x, n);
+					FREE_VEC(y, n);
+					g_free(tmp_row);
+
+					n = luaL_error(L, "(pca on) invalid params at pos %d: "
+									  "bad input dimension %d; %d expected",
+								   s + 1,
+								   (int) rspamd_lua_table_size(L, -1),
+								   pca->dim[1]);
+					lua_pop(L, 1);
+
+					return n;
+				}
+
+
+				for (int i = 0; i < pca->dim[1]; i++) {
+					lua_rawgeti(L, -1, i + 1);
+					tmp_row[i] = lua_tonumber(L, -1);
+
+					lua_pop(L, 1);
+				}
+
+				kad_sgemm_simple(0, 1, 1, n_in,
+								 pca->dim[1], tmp_row, pca->data,
+								 x[s]);
+			}
+
+			lua_pop(L, 1);
+
+			/* Outputs */
+			y[s] = (float *) g_malloc(sizeof(float) * n_out);
+			lua_rawgeti(L, 3, s + 1);
+
+			if (rspamd_lua_table_size(L, -1) != n_out) {
+				FREE_VEC(x, n);
+				FREE_VEC(y, n);
+				g_free(tmp_row);
+
+				n = luaL_error(L, "invalid params at pos %d: "
+								  "bad output dimension %d; "
+								  "%d expected",
+							   s + 1,
+							   (int) rspamd_lua_table_size(L, -1),
+							   n_out);
+				lua_pop(L, 1);
+
+				return n;
+			}
+
+			for (int i = 0; i < n_out; i++) {
+				lua_rawgeti(L, -1, i + 1);
+				y[s][i] = lua_tonumber(L, -1);
+
+				lua_pop(L, 1);
+			}
+
+			lua_pop(L, 1);
+		}
+
+		struct rspamd_kann_train_cbdata cbd;
+
+		cbd.cbref = cbref;
+		cbd.k = k;
+		cbd.L = L;
+
+		int niters = kann_train_fnn1(k, lr,
+									 mini_size, max_epoch, max_drop_streak,
+									 frac_val, n, x, y, lua_kann_train_cb, &cbd);
+
+		lua_pushinteger(L, niters);
+
+		FREE_VEC(x, n);
+		FREE_VEC(y, n);
+		g_free(tmp_row);
+	}
+	else {
+		return luaL_error(L, "invalid arguments: kann, inputs, outputs and"
+							 " optional params are expected");
+	}
+
+	return 1;
+}
+
+static int
+lua_kann_apply1(lua_State *L)
+{
+	kann_t *k = lua_check_kann(L, 1);
+	struct rspamd_lua_tensor *pca = NULL;
+
+	if (k) {
+		if (lua_istable(L, 2)) {
+			gsize vec_len = rspamd_lua_table_size(L, 2);
+			float *vec = (float *) g_malloc(sizeof(float) * vec_len),
+				  *pca_out = NULL;
+			int i_out;
+			int n_in = kann_dim_in(k);
+
+			if (n_in <= 0) {
+				g_free(vec);
+				return luaL_error(L, "invalid inputs count: %d", n_in);
+			}
+
+			if (lua_isuserdata(L, 3)) {
+				pca = lua_check_tensor(L, 3);
+
+				if (pca) {
+					if (pca->ndims != 2) {
+						g_free(vec);
+						return luaL_error(L, "invalid pca tensor: matrix expected, got a row");
+					}
+
+					if (pca->dim[0] != n_in) {
+						g_free(vec);
+						return luaL_error(L, "invalid pca tensor: "
+											 "matrix must have %d rows and it has %d rows instead",
+										  n_in, pca->dim[0]);
+					}
+				}
+				else {
+					g_free(vec);
+					return luaL_error(L, "invalid params: pca matrix expected");
+				}
+			}
+			else {
+				if (n_in != vec_len) {
+					g_free(vec);
+					return luaL_error(L, "invalid params: bad input dimension %d; %d expected",
+									  (int) vec_len, n_in);
+				}
+			}
+
+			for (gsize i = 0; i < vec_len; i++) {
+				lua_rawgeti(L, 2, i + 1);
+				vec[i] = lua_tonumber(L, -1);
+				lua_pop(L, 1);
+			}
+
+			i_out = kann_find(k, KANN_F_OUT, 0);
+
+			if (i_out <= 0) {
+				g_free(vec);
+				return luaL_error(L, "invalid ANN: output layer is missing or is "
+									 "at the input pos");
+			}
+
+			kann_set_batch_size(k, 1);
+			if (pca) {
+				pca_out = g_malloc(sizeof(float) * n_in);
+
+				kad_sgemm_simple(0, 1, 1, n_in,
+								 vec_len, vec, pca->data,
+								 pca_out);
+
+				kann_feed_bind(k, KANN_F_IN, 0, &pca_out);
+			}
+			else {
+				kann_feed_bind(k, KANN_F_IN, 0, &vec);
+			}
+
+			kad_eval_at(k->n, k->v, i_out);
+
+			gsize outlen = kad_len(k->v[i_out]);
+			lua_createtable(L, outlen, 0);
+
+			for (gsize i = 0; i < outlen; i++) {
+				lua_pushnumber(L, k->v[i_out]->x[i]);
+				lua_rawseti(L, -2, i + 1);
+			}
+
+			g_free(vec);
+			g_free(pca_out);
+		}
+		else if (lua_isuserdata(L, 2)) {
+			struct rspamd_lua_tensor *t = lua_check_tensor(L, 2);
+
+			if (t && t->ndims == 1) {
+				int i_out;
+				int n_in = kann_dim_in(k);
+
+				if (n_in != t->dim[0]) {
+					return luaL_error(L, "invalid params: bad input dimension %d; %d expected",
+									  (int) t->dim[0], n_in);
+				}
+
+				i_out = kann_find(k, KANN_F_OUT, 0);
+
+				if (i_out <= 0) {
+					return luaL_error(L, "invalid ANN: output layer is missing or is "
+										 "at the input pos");
+				}
+
+				kann_set_batch_size(k, 1);
+				kann_feed_bind(k, KANN_F_IN, 0, &t->data);
+				kad_eval_at(k->n, k->v, i_out);
+
+				gint outlen = kad_len(k->v[i_out]);
+				struct rspamd_lua_tensor *out;
+				out = lua_newtensor(L, 1, &outlen, false, false);
+				/* Ensure that kann and tensor have the same understanding of floats */
+				G_STATIC_ASSERT(sizeof(float) == sizeof(rspamd_tensor_num_t));
+				memcpy(out->data, k->v[i_out]->x, outlen * sizeof(float));
+			}
+			else {
+				return luaL_error(L, "invalid arguments: 1D rspamd{tensor} expected");
+			}
+		}
+		else {
+			return luaL_error(L, "invalid arguments: 1D rspamd{tensor} expected");
+		}
+	}
+	else {
+		return luaL_error(L, "invalid arguments: rspamd{kann} expected");
+	}
+
+	return 1;
+}
\ No newline at end of file