fix tests and code

Author: tinker495

tests/xnp/update_on_condition_test.py

@@ -277,7 +277,7 @@ def test_update_on_condition_condition_shape_mismatch_raises():
     indices = jnp.array([0, 1, 2], dtype=jnp.int32)
     condition = jnp.array([[True, False, True]], dtype=jnp.bool_)
 
-    with pytest.raises(ValueError, match="`condition` shape .* must match `true_values` shape"):
+    with pytest.raises(ValueError, match="`condition` shape .* must match `indices` shape"):
         xnp.update_on_condition(original, indices, condition, 5.0)
 
 

xtructure/core/xtructure_numpy/array_ops.py

@@ -39,6 +39,11 @@ def _update_array_on_condition(
         return original_array
 
     indices_array = jnp.asarray(indices)
+    if condition.shape != indices_array.shape:
+        raise ValueError(
+            f"`condition` shape {condition.shape} must match `indices` shape {indices_array.shape}."
+        )
+
     indices_array = jnp.reshape(indices_array, (num_updates,))
     index_dtype = indices_array.dtype
     invalid_index = jnp.array(original_array.shape[0], dtype=index_dtype)