Add BcastUntilActxArray

Co-authored-by: Matt Smith <[email protected]>

Author: inducer

arraycontext/__init__.py

@@ -46,6 +46,7 @@
     serialize_container,
 )
 from .container.arithmetic import (
+    BcastUntilActxArray,
     with_container_arithmetic,
 )
 from .container.dataclass import dataclass_array_container
@@ -115,6 +116,7 @@
     "ArrayOrContainerOrScalarT",
     "ArrayOrContainerT",
     "ArrayT",
+    "BcastUntilActxArray",
     "CommonSubexpressionTag",
     "EagerJAXArrayContext",
     "ElementwiseMapKernelTag",

arraycontext/container/arithmetic.py

@@ -6,6 +6,8 @@
 .. currentmodule:: arraycontext
 
 .. autofunction:: with_container_arithmetic
+
+.. autoclass:: BcastUntilActxArray
 """
 
 
@@ -34,12 +36,22 @@
 """
 
 import enum
+import operator
 from collections.abc import Callable
+from dataclasses import dataclass
+from functools import partialmethod
 from typing import Any, TypeVar
 from warnings import warn
 
 import numpy as np
 
+from arraycontext.container import (
+    NotAnArrayContainerError,
+    deserialize_container,
+    serialize_container,
+)
+from arraycontext.context import ArrayContext, ArrayOrArithContainer
+
 
 # {{{ with_container_arithmetic
 
@@ -402,8 +414,9 @@ def wrap(cls: Any) -> Any:
                 warn(
                     f"Broadcasting array context array types across {cls} "
                     "has been explicitly "
-                    "enabled. As of 2025, this will stop working. "
-                    "There is no replacement as of right now. "
+                    "enabled. As of 2026, this will stop working. "
+                    "Use the arraycontext.Bcast* object wrappers for "
+                    "roughly equivalent functionality. "
                     "See the discussion in "
                     "https://github.com/inducer/arraycontext/pull/190. "
                     "To opt out now (and avoid this warning), "
@@ -413,8 +426,9 @@ def wrap(cls: Any) -> Any:
                 warn(
                     f"Broadcasting array context array types across {cls} "
                     "has been implicitly "
-                    "enabled. As of 2025, this will no longer work. "
-                    "There is no replacement as of right now. "
+                    "enabled. As of 2026, this will no longer work. "
+                    "Use the arraycontext.Bcast* object wrappers for "
+                    "roughly equivalent functionality. "
                     "See the discussion in "
                     "https://github.com/inducer/arraycontext/pull/190. "
                     "To opt out now (and avoid this warning), "
@@ -603,8 +617,9 @@ def {fname}(arg1):
                             if isinstance(arg2, {tup_str(bcast_actx_ary_types)}):
                                 warn("Broadcasting {cls} over array "
                                     f"context array type {{type(arg2)}} is deprecated "
-                                    "and will no longer work in 2025. "
-                                    "There is no replacement as of right now. "
+                                    "and will no longer work in 2026. "
+                                    "Use the arraycontext.Bcast* object wrappers for "
+                                    "roughly equivalent functionality. "
                                     "See the discussion in "
                                     "https://github.com/inducer/arraycontext/"
                                     "pull/190. ",
@@ -654,8 +669,10 @@ def {fname}(arg2, arg1):
                                         warn("Broadcasting {cls} over array "
                                             f"context array type {{type(arg1)}} "
                                             "is deprecated "
-                                            "and will no longer work in 2025."
-                                            "There is no replacement as of right now. "
+                                            "and will no longer work in 2026."
+                                            "Use the arraycontext.Bcast* object "
+                                            "rappers for roughly equivalent "
+                                            "functionality. "
                                             "See the discussion in "
                                             "https://github.com/inducer/arraycontext/"
                                             "pull/190. ",
@@ -687,4 +704,71 @@ def {fname}(arg2, arg1):
 # }}}
 
 
+@dataclass(frozen=True)
+class BcastUntilActxArray:
+    array_context: ArrayContext
+    broadcastee: ArrayOrArithContainer
+
+    def _binary_op(self,
+                   op: Callable[
+                       [ArrayOrArithContainer, ArrayOrArithContainer],
+                       ArrayOrArithContainer
+                   ],
+                   right: ArrayOrArithContainer
+               ) -> ArrayOrArithContainer:
+        try:
+            serialized = serialize_container(right)
+        except NotAnArrayContainerError:
+            return op(self.broadcastee, right)
+
+        return deserialize_container(right, [
+            (k, op(self.broadcastee, right_v)
+                if isinstance(right_v, self.array_context.array_types) else
+                self._binary_op(op, right_v)
+            )
+            for k, right_v in serialized])
+
+    def _rev_binary_op(self,
+                   op: Callable[
+                       [ArrayOrArithContainer, ArrayOrArithContainer],
+                       ArrayOrArithContainer
+                   ],
+                   left: ArrayOrArithContainer
+               ) -> ArrayOrArithContainer:
+        try:
+            serialized = serialize_container(left)
+        except NotAnArrayContainerError:
+            return op(left, self.broadcastee)
+
+        return deserialize_container(left, [
+            (k, op(left_v, self.broadcastee)
+                if isinstance(left_v, self.array_context.array_types) else
+                self._rev_binary_op(op, left_v)
+            )
+            for k, left_v in serialized])
+
+    __add__ = partialmethod(_binary_op, operator.add)
+    __radd__ = partialmethod(_rev_binary_op, operator.add)
+    __sub__ = partialmethod(_binary_op, operator.sub)
+    __rsub__ = partialmethod(_rev_binary_op, operator.sub)
+    __mul__ = partialmethod(_binary_op, operator.mul)
+    __rmul__ = partialmethod(_rev_binary_op, operator.mul)
+    __truediv__ = partialmethod(_binary_op, operator.truediv)
+    __rtruediv__ = partialmethod(_rev_binary_op, operator.truediv)
+    __floordiv__ = partialmethod(_binary_op, operator.floordiv)
+    __rfloordiv__ = partialmethod(_rev_binary_op, operator.floordiv)
+    __mod__ = partialmethod(_binary_op, operator.mod)
+    __rmod__ = partialmethod(_rev_binary_op, operator.mod)
+    __pow__ = partialmethod(_binary_op, operator.pow)
+    __rpow__ = partialmethod(_rev_binary_op, operator.pow)
+
+    __lshift__ = partialmethod(_binary_op, operator.lshift)
+    __rlshift__ = partialmethod(_rev_binary_op, operator.lshift)
+    __rshift__ = partialmethod(_binary_op, operator.rshift)
+    __rrshift__ = partialmethod(_rev_binary_op, operator.rshift)
+    __and__ = partialmethod(_binary_op, operator.and_)
+    __rand__ = partialmethod(_rev_binary_op, operator.and_)
+    __or__ = partialmethod(_binary_op, operator.or_)
+    __ror__ = partialmethod(_rev_binary_op, operator.or_)
+
 # vim: foldmethod=marker