Simplify create_diagonal_array

NeilGirdhar · Feb 10, 2025 · 8e5542b · 8e5542b
1 parent 5dd8a8a
commit 8e5542b
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tjax/_src/math_tools.py b/tjax/_src/math_tools.py
@@ -171,5 +171,5 @@ def create_diagonal_array(m: T) -> T:
     retval = xp.zeros((*pre, n ** 2), dtype=m.dtype)
     for index in np.ndindex(*pre):
         target_index = (*index, slice(None, None, n + 1))
-        xpx.at(retval)[target_index].set(m[*index, :])
+        retval = xpx.at(retval)[target_index].set(m[*index, :])
     return xp.reshape(retval, (*m.shape, n))