diff options
Diffstat (limited to 'src/nn_train.py')
| -rwxr-xr-x | src/nn_train.py | 11 |
1 files changed, 8 insertions, 3 deletions
diff --git a/src/nn_train.py b/src/nn_train.py index f82a369..23cf38b 100755 --- a/src/nn_train.py +++ b/src/nn_train.py @@ -486,9 +486,14 @@ def fast_tanh(x: Tensor) -> Tensor: x2 = x * x poly = x * ( 0.997136032 - + x2 * ( + + x2 + * ( -0.305215124 - + x2 * (0.0832466371 + x2 * (-0.0139947924 + x2 * (0.00123106314 + x2 * -4.29748644e-05))) + + x2 + * ( + 0.0832466371 + + x2 * (-0.0139947924 + x2 * (0.00123106314 + x2 * -4.29748644e-05)) + ) ) ) return torch.where(x > 3.0, 1.0, torch.where(x < -3.0, -1.0, poly)) @@ -543,7 +548,7 @@ class Net(Module): f"#ifndef {guard}\n#define {guard}\n", f'#include "weights_{size}.h"\n\n', f"const float {prefix}_embedding[{len(E)}] =\n{{\n{_fmt_flat(E, per_line=7)}\n}};\n\n", - f"const float {prefix}_dense_weights[{upper}_DENSE_NUM][{upper}_INP_NUM] NN_ALIGNED =\n{{\n{_fmt_matrix(W)}\n}};\n\n", + f"const float {prefix}_dense_weights[{upper}_INP_NUM][{upper}_DENSE_NUM] NN_ALIGNED =\n{{\n{_fmt_matrix(W.T)}\n}};\n\n", f"const float {prefix}_dense_biases[{upper}_DENSE_NUM] =\n{{\n{_fmt_flat(B)}\n}};\n\n", f"const float {prefix}_logit_weights[{upper}_DENSE_NUM] NN_ALIGNED =\n{{\n{_fmt_flat(LW)}\n}};\n\n", f"const float {prefix}_logit_bias = {LB:.9g};\n", |