Migrate AbstractArray to qp.math

Author: mudit2812

pennylane/math/__init__.py

@@ -35,6 +35,14 @@
 
 import autoray as ar
 
+from .abstract_types import (
+    AbstractArray,
+    AbstractBool,
+    AbstractComplex,
+    AbstractFloat,
+    AbstractInt,
+    AbstractWires,
+)
 from .binary_linalg import (
     binary_decimals,
     binary_finite_reduced_row_echelon,
@@ -210,6 +218,12 @@ def __getattr__(name):
 
 
 __all__ = [
+    "AbstractArray",
+    "AbstractBool",
+    "AbstractComplex",
+    "AbstractFloat",
+    "AbstractInt",
+    "AbstractWires",
     "add",
     "allclose",
     "allequal",

pennylane/math/abstract_types.py

@@ -0,0 +1,94 @@
+# Copyright 2026 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""This module contains data structures to represent abstract arrays."""
+
+from dataclasses import dataclass
+from math import prod
+from numbers import Number
+
+import numpy as np
+
+
+@dataclass(frozen=True)
+class AbstractArray:
+    """Abstract array type."""
+
+    shape: tuple[int, ...]
+    dtype: np.dtype | type[Number] = np.int64
+
+    def __post_init__(self):
+        object.__setattr__(self, "shape", tuple(self.shape))
+        object.__setattr__(self, "dtype", np.dtype(self.dtype))
+
+    @property
+    def size(self) -> int:
+        """Total number of elements."""
+        return prod(self.shape)
+
+    @property
+    def T(self) -> "AbstractArray":
+        """Transpose view of the array."""
+        return AbstractArray(self.shape[::-1], self.dtype)
+
+    @property
+    def ndim(self) -> int:
+        """Number of dimensions."""
+        return len(self.shape)
+
+    def __getitem__(self, *_, **__):
+        raise IndexError("Cannot index into an abstract array.")
+
+    def __setitem__(self, *_, **__):
+        raise IndexError("Cannot index into an abstract array.")
+
+    def __len__(self) -> int:
+        if not self.shape:
+            raise TypeError("len() of unsized object.")
+        return self.shape[0]
+
+    def __eq__(self, other: "AbstractArray") -> bool:
+        # This should probably just raise an error
+        if isinstance(other, AbstractArray):
+            return self.shape == other.shape and self.dtype == other.dtype
+
+        raise TypeError("Tried to check equality against an abstract array.")
+
+    def __hash__(self) -> int:
+        return hash((self.shape, self.dtype))
+
+
+AbstractBool = AbstractArray((), bool)
+AbstractInt = AbstractArray((), int)
+AbstractFloat = AbstractArray((), float)
+AbstractComplex = AbstractArray((), complex)
+
+
+@dataclass(frozen=True)
+class AbstractWires(AbstractArray):
+    """Abstract wires."""
+
+    num_wires: int
+
+    def __post_init__(self):
+        object.__setattr__(self, "shape", (self.num_wires,))
+        object.__setattr__(self, "dtype", int)
+
+    def __eq__(self, other: "AbstractWires"):
+        if isinstance(other, AbstractWires):
+            return self.num_wires == other.num_wires
+
+        raise TypeError("Tried to check equality against an abstract wire register.")
+
+    def __hash__(self):
+        return hash(("AbstractWires", self.num_wires))

pennylane/math/utils.py

@@ -24,6 +24,8 @@
 
 from pennylane import math
 
+from .abstract_types import AbstractArray
+
 
 def allequal(tensor1, tensor2, **kwargs):
     """Returns True if two tensors are element-wise equal along a given axis.
@@ -417,6 +419,9 @@ def function(x):
     Abstract: True
     <tf.Tensor: shape=(), dtype=float32, numpy=0.26>
     """
+    if isinstance(tensor, AbstractArray):
+        return True
+
     interface = like or math.get_interface(tensor)
 
     if interface == "jax":

pennylane/templates/core.py

@@ -21,7 +21,6 @@
 
     ~Subroutine
     ~SubroutineOp
-    ~AbstractArray
     ~change_op_basis_subroutine_resource_rep
     ~adjoint_subroutine_resource_rep
     ~subroutine_resource_rep
@@ -32,8 +31,7 @@
 from collections import defaultdict
 from collections.abc import Callable
 from copy import deepcopy
-from dataclasses import dataclass
-from functools import lru_cache, reduce, update_wrapper
+from functools import lru_cache, update_wrapper
 from importlib.util import find_spec
 from inspect import BoundArguments, Signature, signature
 from typing import Any, ParamSpec
@@ -58,34 +56,6 @@
 has_jax = find_spec("jax") is not None
 
 
-@dataclass(frozen=True)
-class AbstractArray:
-    """An abstract representation of an array that contains the shape and dtype
-    attributes necessary for resource calculations.
-
-    This class is used with :func:`~pennylane.templates.subroutine_resource_rep`
-    for specifying abstract information about a :class:`~.Subroutine` for
-    purposes of resource calculations used with graph decompositions.
-
-    Args:
-        shape (tuple(int)): the dimensions of the array. ``()`` corresponds to a scalar.
-        dtype (type): the data type of the array. Defaults to ``np.dtype(int)`` for easier use in specifying
-        wires.
-    """
-
-    shape: tuple[int, ...]
-    dtype: np.dtype = np.dtype(int)
-
-    def __len__(self):
-        return reduce(lambda a, b: a * b, self.shape)
-
-    def __post_init__(self):
-        if math.get_interface(self.dtype) == "torch":
-            dummy = math.array((), dtype=self.dtype, like="torch")
-            object.__setattr__(self, "dtype", dummy.numpy().dtype)
-        object.__setattr__(self, "dtype", np.dtype(self.dtype))
-
-
 def _make_signature_key(subroutine: "Subroutine", *args, **kwargs):
     bound = subroutine.signature.bind(*args, **kwargs)
     bound.apply_defaults()
@@ -207,7 +177,8 @@ def S(params, wires, rotation):
 
     .. code-block:: python
 
-        from pennylane.templates import AbstractArray, subroutine_resource_rep
+        from pennylane.math import AbstractArray
+        from pennylane.templates import subroutine_resource_rep
 
         class MyOp(qp.operation.Operation):
             pass
@@ -254,12 +225,12 @@ def _create_signature_key(
         if arg in static_argnames:
             key.append(val)
         elif arg in wire_argnames:
-            key.append(AbstractArray(shape=(len(val),), dtype=int))
+            key.append(math.AbstractArray(shape=(len(val),), dtype=int))
         else:
             leaves, struct = flatten(val)
 
             shapes = (
-                AbstractArray(shape=math.shape(l), dtype=getattr(l, "dtype", type(l)))
+                math.AbstractArray(shape=math.shape(l), dtype=getattr(l, "dtype", type(l)))
                 for l in leaves
             )
             key.append((struct, tuple(shapes)))
@@ -427,11 +398,11 @@ def _default_resources(subroutine: "Subroutine", *args, **kwargs) -> defaultdict
     sig = subroutine.signature.bind(*args, **kwargs)
     for arg in subroutine.dynamic_argnames:
         avals, struct = flatten(sig.arguments[arg])
-        if avals and isinstance(avals[0], AbstractArray):
+        if avals and isinstance(avals[0], math.AbstractArray):
             params = (np.empty(shape=aval.shape, dtype=aval.dtype) for aval in avals)
             sig.arguments[arg] = unflatten(params, struct)
     for arg in subroutine.wire_argnames:
-        if isinstance(sig.arguments[arg], AbstractArray):
+        if isinstance(sig.arguments[arg], math.AbstractArray):
             sig.arguments[arg] = list(range(sig.arguments[arg].shape[0]))
     with queuing.AnnotatedQueue() as q:
         subroutine.definition(**sig.arguments)
@@ -633,7 +604,7 @@ def RXLayer(params, wires):
     For example, we should be able to calculate the resources using the :class:`~.AbstractArray`
     class.
 
-    >>> from pennylane.templates import AbstractArray
+    >>> from pennylane.math import AbstractArray
     >>> abstract_params = AbstractArray((10,), float)
     >>> abstract_wires = AbstractArray((10,))
     >>> RXLayer.compute_resources(abstract_params, abstract_wires)
@@ -644,7 +615,8 @@ def RXLayer(params, wires):
 
     .. code-block:: python
 
-        from pennylane.templates import AbstractArray, subroutine_resource_rep
+        from pennylane.math import AbstractArray
+        from pennylane.templates import subroutine_resource_rep
 
         class MyOp(qp.operation.Operation):
             pass
@@ -918,7 +890,6 @@ def _(*args, **kwargs):
 __all__ = [
     "Subroutine",
     "SubroutineOp",
-    "AbstractArray",
     "subroutine_resource_rep",
     "CollectedSubroutine",
     "adjoint_subroutine_resource_rep",