Allow scalars for actx.tag, document it is best-effort

Author: inducer

arraycontext/context.py

@@ -125,7 +125,7 @@
 )
 from warnings import warn
 
-from typing_extensions import Self
+from typing_extensions import Self, override
 
 from pytools import memoize_method
 
@@ -146,7 +146,6 @@
         ArrayOrArithContainerOrScalarT,
         ArrayOrContainerOrScalar,
         ArrayOrContainerOrScalarT,
-        ArrayOrContainerT,
         ContainerOrScalarT,
         NumpyOrContainerOrScalar,
         ScalarLike,
@@ -217,6 +216,7 @@ def __init__(self) -> None:
     def _get_fake_numpy_namespace(self) -> BaseFakeNumpyNamespace:
         ...
 
+    @override
     def __hash__(self) -> int:
         raise TypeError(f"unhashable type: '{type(self).__name__}'")
 
@@ -333,12 +333,14 @@ def freeze_thaw(
     @abstractmethod
     def tag(self,
             tags: ToTagSetConvertible,
-            array: ArrayOrContainerT) -> ArrayOrContainerT:
+            array: ArrayOrContainerOrScalarT) -> ArrayOrContainerOrScalarT:
         """If the array type used by the array context is capable of capturing
         metadata, return a version of *array* with the *tags* applied. *array*
         itself is not modified. When working with array containers, the
         tags are applied to each leaf of the container.
 
+        Tagging is best-effort. Untaggable types will be returned as-is.
+
         See :ref:`metadata` as well as application-specific metadata types.
 
         .. versionadded:: 2021.2

arraycontext/impl/pytato/__init__.py

@@ -76,7 +76,6 @@
     Array,
     ArrayOrArithContainerOrScalarT,
     ArrayOrContainerOrScalarT,
-    ArrayOrContainerT,
     ArrayOrScalar,
     ScalarLike,
     is_scalar_like,
@@ -1031,8 +1030,8 @@ def transform_dag(self, dag: pytato.AbstractResultWithNamedArrays
     @override
     def tag(self,
                 tags: ToTagSetConvertible,
-                array: ArrayOrContainerT,
-            ) -> ArrayOrContainerT:
+                array: ArrayOrContainerOrScalarT,
+            ) -> ArrayOrContainerOrScalarT:
         def _tag(ary: Array) -> Array:
             import jax.numpy as jnp
             if isinstance(ary, jnp.ndarray):