Resolve TODO

NeilGirdhar · Dec 25, 2024 · 4b190a2 · 4b190a2
1 parent 711d60e
commit 4b190a2
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/tjax/_src/math_tools.py b/tjax/_src/math_tools.py
@@ -17,11 +17,11 @@ def abs_square(x: JaxComplexArray) -> JaxRealArray: ...
 def abs_square(x: ComplexArray) -> RealArray: ...
 def abs_square(x: ComplexArray) -> RealArray:
     xp = get_namespace(x)
-    # TODO: remove workaround when Jax is 0.4.27.
-    return xp.square(x.real) + xp.square(xp.asarray(x.imag))
+    return xp.square(x.real) + xp.square(x.imag)
 
 
-# TODO: Remove this when the Array API has it with broadcasting under xp.linalg.norm.
+# TODO: Remove when it's added to the Array API:
+# https://github.com/data-apis/array-api/issues/242
 @overload
 def outer_product(x: JaxRealArray, y: JaxRealArray) -> JaxRealArray: ...
 @overload