1 files changed, 14 insertions, 13 deletions

diff --git a/include/nn.c b/include/nn.c
index 8c85299..909b870 100644
--- a/include/nn.c
+++ b/include/nn.c
@@ -62,7 +62,8 @@ inline uint8_t cell_idx_at(const tak_state_p state, const uint8_t loc) {
     default: {
       // |{flat, standing, cap} * {black, white}| = 6
       // index 1 .. 6
-      return (uint8_t)(1u + 3u * TOP_COLOUR_AT(state, loc) + (d & STONE_MASK));
+      const uint8_t top_colour = TOP_COLOUR_AT(state, loc);
+      return (uint8_t)(1u + 3u * ((state->ply & 1) ? 1-top_colour : top_colour)  + (d & STONE_MASK));
     }
   }
 }
@@ -73,12 +74,12 @@ inline void nn_cell_idxs(uint8_t *out, tak_state_p state) {
     out[loc] = cell_idx_at(state, loc);
 }
 
-inline void nn_prepare_head(float out_head[3], tak_state_p state) {
+inline void nn_prepare_head(float out_head[2], tak_state_p state) {
   const float inv_max_flats =
     (state->board_size == 6) ? (1.0f / 30.0f) : (1.0f / 21.0f);
-  out_head[0] = (state->ply & 1) ? 1 : -1;
-  out_head[1] = (float)(state->white_count & 127) * inv_max_flats;
-  out_head[2] = (float)(state->black_count & 127) * inv_max_flats;
+  const uint8_t idx = state->ply & 1;
+  out_head[idx] = (float)(state->white_count & 127) * inv_max_flats;
+  out_head[1-idx] = (float)(state->black_count & 127) * inv_max_flats;
 }
 
 // The cell indices computed by the most recent nn_prepare_input, and the hash
@@ -93,7 +94,7 @@ void nn_prepare_input(float *cur_board, tak_state_p state,
   nn_cell_idxs(prepared_idxs, state);
   prepared_hash = state->hash;
   for (unsigned int loc = 0; loc < n_cells; loc++)
-    cur_board[3 + loc] = embedding[prepared_idxs[loc]];
+    cur_board[2 + loc] = embedding[prepared_idxs[loc]];
 }
 
 // Copy the cell indices produced by the most recent nn_prepare_input into out
@@ -111,10 +112,10 @@ int nn_last_cell_idxs(uint8_t *out, tak_state_p state) {
 // Implementation of small neural networks
 // ===================================================================
 
-float eval_net(const float *board, const unsigned int n_inp,
-               const float *w, const float *biases,
-               const float *logit_w, const float logit_b,
-               const unsigned int n_dense) {
+float eval_net(const float *board, const unsigned int n_inp, const float *w,
+               const float *biases, const float *logit_w, const float logit_b,
+               const unsigned int n_dense,
+              const float parity) {
   float dense1[FIVE_DENSE_NUM]; // FIVE_DENSE_NUM == SIX_DENSE_NUM
   float z = logit_b;
 
@@ -128,7 +129,7 @@ float eval_net(const float *board, const unsigned int n_inp,
   for (unsigned int d1 = 0; d1 < n_dense; d1++)
     z += dense1[d1] * logit_w[d1];
 
-  return fast_tanh(z);
+  return parity * fast_tanh(z);
 }
 
 // ===================================================================
@@ -141,7 +142,7 @@ float five_nn_fn(tak_state_p state) {
   nn_prepare_input(cur_board, state, five_embedding);
   return eval_net(cur_board, FIVE_INP_NUM, &five_dense_weights[0][0],
                   five_dense_biases, five_logit_weights, five_logit_bias,
-                  FIVE_DENSE_NUM);
+                  FIVE_DENSE_NUM, state->ply & 1 ? 1.0f : -1.0f);
 }
 
 float six_nn_fn(tak_state_p state) {
@@ -150,5 +151,5 @@ float six_nn_fn(tak_state_p state) {
   nn_prepare_input(cur_board, state, six_embedding);
   return eval_net(cur_board, SIX_INP_NUM, &six_dense_weights[0][0],
                   six_dense_biases, six_logit_weights, six_logit_bias,
-                  SIX_DENSE_NUM);
+                  SIX_DENSE_NUM, state->ply & 1 ? 1.0f : -1.0f);
 }