broadcast_shapes

Author: crusaderky

src/array_api_extra/_lib/_lazy.py

@@ -9,6 +9,7 @@
 from types import ModuleType
 from typing import TYPE_CHECKING, Any, cast, overload
 
+from ._funcs import broadcast_shapes
 from ._utils import _compat
 from ._utils._compat import (
     array_namespace,
@@ -207,9 +208,7 @@ def lazy_apply(  # type: ignore[valid-type]  # numpydoc ignore=GL07,SA04
     multi_output = False
 
     if shape is None:
-        # FIXME https://github.com/data-apis/array-api-extra/pull/133
-        # shapes = [broadcast_shapes(*(arg.shape for arg in args))]
-        shapes = [xp.broadcast_arrays(*args)[0].shape]
+        shapes = [broadcast_shapes(*(arg.shape for arg in args))]
     elif all(isinstance(s, int | None) for s in shape):
         # Do not test for shape to be a tuple
         # https://github.com/data-apis/array-api/issues/891#issuecomment-2637430522