1D matrix fixing

Author: ItCodinTime

examples/keras_manual_apply.py

@@ -5,6 +5,7 @@
 """
 
 import keras
+import tensorflow as tf
 
 from neat_optim import NEAT
 
@@ -14,7 +15,7 @@ def main() -> None:
     gradient = keras.ops.array([0.5, -0.25], dtype="float32")
     optimizer = NEAT(learning_rate=0.1, alpha=0.25, beta=0.9)
     optimizer.apply_gradients([(gradient, variable)])
-    print(keras.ops.convert_to_numpy(variable))
+    print(tf.convert_to_tensor(variable).numpy())
 
 
 if __name__ == "__main__":

src/neat_optim/engine/reference.py

@@ -11,21 +11,25 @@ def _as_float32(array: np.ndarray) -> np.ndarray:
     return np.asarray(array, dtype=np.float32)
 
 
+def _flat_dot(left: np.ndarray, right: np.ndarray) -> float:
+    return float(np.dot(left.reshape(-1), right.reshape(-1)))
+
+
 def _safe_projection(
     gradient: np.ndarray, vector: np.ndarray, eps: float
 ) -> np.ndarray:
-    denom = float(np.dot(vector, vector))
+    denom = _flat_dot(vector, vector)
     if denom <= eps:
         return np.zeros_like(gradient)
-    return (float(np.dot(gradient, vector)) / (denom + eps)) * vector
+    return (_flat_dot(gradient, vector) / (denom + eps)) * vector
 
 
 def _conflict_ratio(gradient: np.ndarray, vector: np.ndarray, eps: float) -> float:
     grad_norm = l2_norm(gradient)
     vec_norm = l2_norm(vector)
     if grad_norm <= eps or vec_norm <= eps:
         return 0.0
-    cosine = float(np.dot(gradient, vector) / ((grad_norm * vec_norm) + eps))
+    cosine = _flat_dot(gradient, vector) / ((grad_norm * vec_norm) + eps)
     return max(0.0, -cosine)
 
 

tests/integration/test_keras_optimizer.py

@@ -10,6 +10,10 @@
 from neat_optim.state import ArrayState  # noqa: E402
 
 
+def _to_numpy(tensor) -> np.ndarray:
+    return tensorflow.convert_to_tensor(tensor).numpy()
+
+
 def _run_reference(
     initial_param: np.ndarray,
     gradients: list[np.ndarray],
@@ -38,7 +42,7 @@ def _run_keras(
         optimizer.apply_gradients(
             [(keras.ops.array(gradient, dtype="float32"), variable)]
         )
-    return optimizer, keras.ops.convert_to_numpy(variable)
+    return optimizer, _to_numpy(variable)
 
 
 def test_keras_optimizer_can_apply_gradients() -> None:
@@ -49,7 +53,7 @@ def test_keras_optimizer_can_apply_gradients() -> None:
 
     optimizer.apply_gradients([(gradient, variable)])
 
-    values = keras.ops.convert_to_numpy(variable)
+    values = _to_numpy(variable)
     assert values.shape == (2,)
     assert optimizer.variables
 
@@ -99,12 +103,12 @@ def test_keras_optimizer_matches_reference_projection_mode() -> None:
 
     np.testing.assert_allclose(keras_param, reference.param, atol=1e-6)
     np.testing.assert_allclose(
-        keras.ops.convert_to_numpy(optimizer.momentums[0]),
+        _to_numpy(optimizer.momentums[0]),
         reference.state.momentum,
         atol=1e-6,
     )
     np.testing.assert_allclose(
-        keras.ops.convert_to_numpy(optimizer.nces[0]),
+        _to_numpy(optimizer.nces[0]),
         reference.state.nce,
         atol=1e-6,
     )
@@ -134,7 +138,7 @@ def test_keras_optimizer_matches_reference_with_weight_decay_modes() -> None:
     )
     np.testing.assert_allclose(decoupled_param, decoupled_reference.param, atol=1e-6)
     np.testing.assert_allclose(
-        keras.ops.convert_to_numpy(decoupled_optimizer.momentums[0]),
+        _to_numpy(decoupled_optimizer.momentums[0]),
         decoupled_reference.state.momentum,
         atol=1e-6,
     )
@@ -159,7 +163,7 @@ def test_keras_optimizer_matches_reference_with_weight_decay_modes() -> None:
     )
     np.testing.assert_allclose(coupled_param, coupled_reference.param, atol=1e-6)
     np.testing.assert_allclose(
-        keras.ops.convert_to_numpy(coupled_optimizer.momentums[0]),
+        _to_numpy(coupled_optimizer.momentums[0]),
         coupled_reference.state.momentum,
         atol=1e-6,
     )
@@ -191,7 +195,7 @@ def test_keras_optimizer_matches_reference_when_nce_is_disabled() -> None:
 
     np.testing.assert_allclose(keras_param, reference.param, atol=1e-6)
     np.testing.assert_allclose(
-        keras.ops.convert_to_numpy(optimizer.nces[0]),
+        _to_numpy(optimizer.nces[0]),
         np.zeros_like(initial_param),
         atol=1e-6,
     )

tests/unit/test_reference_step.py

@@ -76,3 +76,18 @@ def test_correction_is_clipped_to_gradient_norm() -> None:
     result = neat_step_reference(param, grad, state, config)
 
     assert np.linalg.norm(result.state.nce) == pytest.approx(0.25, abs=1e-6)
+
+
+def test_reference_step_supports_matrix_parameters() -> None:
+    param = np.array([[1.0, -2.0], [0.5, 3.0]], dtype=np.float32)
+    grad = np.array([[0.5, -0.25], [0.1, -0.2]], dtype=np.float32)
+    state = ArrayState.zeros_like(param)
+    config = NEATConfig(learning_rate=0.05, alpha=0.25, beta=0.9)
+
+    result = neat_step_reference(param, grad, state, config)
+
+    assert result.state.step == 1
+    assert result.param.shape == param.shape
+    assert result.state.momentum.shape == param.shape
+    assert result.state.nce.shape == param.shape
+    assert np.isfinite(result.param).all()