Upload own-solver/neurogolf_solver/gather_helpers.py

own-solver/neurogolf_solver/gather_helpers.py

@@ -0,0 +1,63 @@
+#!/usr/bin/env python3
+"""Gather-based model building utilities."""
+
+import numpy as np
+from onnx import numpy_helper, helper
+from .onnx_helpers import mk
+from .constants import GH, GW
+
+
+def _build_gather_model(OH, OW, idx):
+    """Build gather model from index mapping."""
+    flat_idx = np.zeros((GH * GW,), dtype=np.int64)
+    mask = np.zeros((1, 1, GH, GW), dtype=np.float32)
+    for oi in range(OH):
+        for oj in range(OW):
+            flat_idx[oi * GW + oj] = idx[oi, oj, 0] * GW + idx[oi, oj, 1]
+            mask[0, 0, oi, oj] = 1.0
+    inits = [
+        numpy_helper.from_array(np.array([1, 10, GH * GW], dtype=np.int64), 'fs'),
+        numpy_helper.from_array(flat_idx, 'idx'),
+        numpy_helper.from_array(np.array([1, 10, GH, GW], dtype=np.int64), 'os'),
+        numpy_helper.from_array(mask, 'mask'),
+    ]
+    nodes = [
+        helper.make_node('Reshape', ['input', 'fs'], ['flat']),
+        helper.make_node('Gather', ['flat', 'idx'], ['g'], axis=2),
+        helper.make_node('Reshape', ['g', 'os'], ['raw']),
+        helper.make_node('Mul', ['raw', 'mask'], ['output']),
+    ]
+    return mk(nodes, inits)
+
+
+def _build_gather_model_with_const(IH, IW, OH, OW, idx, cst):
+    """Build gather model with constant values."""
+    flat_idx = np.zeros((GH * GW,), dtype=np.int64)
+    gather_mask = np.zeros((1, 1, GH, GW), dtype=np.float32)
+    const_oh = np.zeros((1, 10, GH, GW), dtype=np.float32)
+    for oi in range(OH):
+        for oj in range(OW):
+            if idx[oi, oj, 0] >= 0:
+                flat_idx[oi * GW + oj] = idx[oi, oj, 0] * GW + idx[oi, oj, 1]
+                gather_mask[0, 0, oi, oj] = 1.0
+            elif cst[oi, oj] >= 0:
+                const_oh[0, cst[oi, oj], oi, oj] = 1.0
+    has_const = np.any(const_oh > 0)
+    inits = [
+        numpy_helper.from_array(np.array([1, 10, GH * GW], dtype=np.int64), 'fs'),
+        numpy_helper.from_array(flat_idx, 'idx'),
+        numpy_helper.from_array(np.array([1, 10, GH, GW], dtype=np.int64), 'os'),
+        numpy_helper.from_array(gather_mask, 'gmask'),
+    ]
+    nodes = [
+        helper.make_node('Reshape', ['input', 'fs'], ['flat']),
+        helper.make_node('Gather', ['flat', 'idx'], ['g'], axis=2),
+        helper.make_node('Reshape', ['g', 'os'], ['raw']),
+        helper.make_node('Mul', ['raw', 'gmask'], ['masked']),
+    ]
+    if has_const:
+        inits.append(numpy_helper.from_array(const_oh, 'cst'))
+        nodes.append(helper.make_node('Add', ['masked', 'cst'], ['output']))
+    else:
+        nodes[-1] = helper.make_node('Mul', ['raw', 'gmask'], ['output'])
+    return mk(nodes, inits)