Clean up GPU backend: remove dead matmul, redundant branches, and duplicate logic

- Remove wasted cp.matmul call in CTRNN loop that computed and discarded result;
  use pre-allocated s_buf via out= parameter instead
- Remove per_genome_inputs if/else branches where both paths were identical
  (CuPy broadcasting handles both input shapes)
- Collapse duplicate input-building if/else in both GPU evaluator classes using
  np.zeros((num_steps, *first_input.shape))
- Eliminate N = len(genomes) that could diverge from input_fn's actual shape

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: CodeReclaimers

neat/gpu/_cupy_backend.py

@@ -131,14 +131,8 @@ def evaluate_ctrnn_batch(packed, inputs_cpu, dt):
     scale[:, :num_inputs] = 0.0
 
     # Transfer inputs to GPU.
-    if inputs_cpu.ndim == 2:
-        # [num_steps, num_inputs] → broadcast to all genomes
-        inputs_gpu = cp.asarray(inputs_cpu.astype(np.float32))
-        per_genome_inputs = False
-    else:
-        # [num_steps, N, num_inputs]
-        inputs_gpu = cp.asarray(inputs_cpu.astype(np.float32))
-        per_genome_inputs = True
+    # Shape is either [num_steps, num_inputs] (broadcast) or [num_steps, N, num_inputs].
+    inputs_gpu = cp.asarray(inputs_cpu.astype(np.float32))
 
     # Initialize state.
     u = cp.zeros((N, M), dtype=cp.float32)
@@ -165,15 +159,11 @@ def evaluate_ctrnn_batch(packed, inputs_cpu, dt):
 
     for step in range(num_steps):
         # Step 1: Set input node states.
-        if per_genome_inputs:
-            u[:, :num_inputs] = inputs_gpu[step]  # [N, num_inputs]
-        else:
-            u[:, :num_inputs] = inputs_gpu[step]  # [num_inputs] broadcast
+        u[:, :num_inputs] = inputs_gpu[step]
 
         # Step 2: Batched matrix-vector multiply.
         # s = W @ u → [N, M] (treat u as [N, M, 1], squeeze result)
-        cp.matmul(W, u[:, :, None], out=s_buf[:, :, None] if False else None)
-        s_buf = cp.matmul(W, u[:, :, None]).squeeze(-1)  # [N, M]
+        cp.matmul(W, u[:, :, None], out=s_buf[:, :, None])
 
         # Step 3: Apply activation function via custom kernel.
         s_flat = s_buf.ravel()
@@ -187,10 +177,7 @@ def evaluate_ctrnn_batch(packed, inputs_cpu, dt):
         u += scale * z_buf
 
         # Step 5: Re-clamp input nodes.
-        if per_genome_inputs:
-            u[:, :num_inputs] = inputs_gpu[step]
-        else:
-            u[:, :num_inputs] = inputs_gpu[step]
+        u[:, :num_inputs] = inputs_gpu[step]
 
         # Step 6: Record output node states.
         trajectory[:, step, :] = u[:, out_start:out_end]
@@ -235,13 +222,9 @@ def evaluate_iznn_batch(packed, inputs_cpu, dt, num_steps):
     d = cp.asarray(packed['d'])              # [N, M]
     node_mask = cp.asarray(packed['node_mask'])  # [N, M]
 
-    # Transfer inputs.
-    if inputs_cpu.ndim == 2:
-        inputs_gpu = cp.asarray(inputs_cpu.astype(np.float32))
-        per_genome_inputs = False
-    else:
-        inputs_gpu = cp.asarray(inputs_cpu.astype(np.float32))
-        per_genome_inputs = True
+    # Transfer inputs to GPU.
+    # Shape is either [num_steps, num_inputs] (broadcast) or [num_steps, N, num_inputs].
+    inputs_gpu = cp.asarray(inputs_cpu.astype(np.float32))
 
     # Initialize state: v = c, u_recov = b * v, fired = 0.
     v = c.copy()
@@ -263,10 +246,7 @@ def evaluate_iznn_batch(packed, inputs_cpu, dt, num_steps):
     for step in range(num_steps):
         # Build source vector: fired for neuron slots, external input for input slots.
         source[:] = fired
-        if per_genome_inputs:
-            source[:, :num_inputs] = inputs_gpu[step]
-        else:
-            source[:, :num_inputs] = inputs_gpu[step]
+        source[:, :num_inputs] = inputs_gpu[step]
 
         # Compute synaptic current: I = bias + W @ source
         I = bias + cp.matmul(W, source[:, :, None]).squeeze(-1)

neat/gpu/evaluator.py

@@ -67,23 +67,13 @@ def evaluate(self, genomes, config):
         num_steps = int(self.t_max / self.dt)
 
         # Precompute input trajectory on CPU.
+        # first_input shape is either [num_inputs] or [N, num_inputs].
         first_input = np.asarray(self.input_fn(0.0, self.dt), dtype=np.float32)
-        if first_input.ndim == 1:
-            # Scalar input: [num_inputs]
-            inputs = np.zeros((num_steps, len(first_input)), dtype=np.float32)
-            inputs[0] = first_input
-            for step in range(1, num_steps):
-                inputs[step] = np.asarray(
-                    self.input_fn(step * self.dt, self.dt), dtype=np.float32)
-        else:
-            # Per-genome input: [N, num_inputs]
-            N = len(genomes)
-            num_inputs = first_input.shape[-1]
-            inputs = np.zeros((num_steps, N, num_inputs), dtype=np.float32)
-            inputs[0] = first_input
-            for step in range(1, num_steps):
-                inputs[step] = np.asarray(
-                    self.input_fn(step * self.dt, self.dt), dtype=np.float32)
+        inputs = np.zeros((num_steps, *first_input.shape), dtype=np.float32)
+        inputs[0] = first_input
+        for step in range(1, num_steps):
+            inputs[step] = np.asarray(
+                self.input_fn(step * self.dt, self.dt), dtype=np.float32)
 
         # Pack genomes into padded arrays.
         packed = pack_ctrnn_population(genomes, config)
@@ -131,21 +121,13 @@ def evaluate(self, genomes, config):
         num_steps = int(self.t_max / self.dt)
 
         # Precompute input trajectory on CPU.
+        # first_input shape is either [num_inputs] or [N, num_inputs].
         first_input = np.asarray(self.input_fn(0.0, self.dt), dtype=np.float32)
-        if first_input.ndim == 1:
-            inputs = np.zeros((num_steps, len(first_input)), dtype=np.float32)
-            inputs[0] = first_input
-            for step in range(1, num_steps):
-                inputs[step] = np.asarray(
-                    self.input_fn(step * self.dt, self.dt), dtype=np.float32)
-        else:
-            N = len(genomes)
-            num_inputs = first_input.shape[-1]
-            inputs = np.zeros((num_steps, N, num_inputs), dtype=np.float32)
-            inputs[0] = first_input
-            for step in range(1, num_steps):
-                inputs[step] = np.asarray(
-                    self.input_fn(step * self.dt, self.dt), dtype=np.float32)
+        inputs = np.zeros((num_steps, *first_input.shape), dtype=np.float32)
+        inputs[0] = first_input
+        for step in range(1, num_steps):
+            inputs[step] = np.asarray(
+                self.input_fn(step * self.dt, self.dt), dtype=np.float32)
 
         packed = pack_iznn_population(genomes, config)
         trajectory = evaluate_iznn_batch(packed, inputs, self.dt, num_steps)