[Feature, Testing] Allow nested calls to ConcretizedCallable, more tests

pyproject.toml

@@ -6,7 +6,7 @@ build-backend = "hatchling.build"
 name = "pyqtorch"
 description = "An efficient, large-scale emulator designed for quantum machine learning, seamlessly integrated with a PyTorch backend. Please refer to https://pyqtorch.readthedocs.io/en/latest/ for setup and usage info, along with the full documentation."
 readme = "README.md"
-version = "1.4.9"
+version = "1.5.0"
 requires-python = ">=3.8,<3.13"
 license = { text = "Apache 2.0" }
 keywords = ["quantum"]
@@ -117,7 +117,7 @@ serve = "mkdocs serve --dev-addr localhost:8000"
 
 [tool.ruff]
 lint.select = ["E", "F", "I", "Q"]
-lint.extend-ignore = ["F841"]
+lint.extend-ignore = ["F841", "E731"]
 line-length = 100
 
 [tool.ruff.lint.isort]

pyqtorch/__init__.py

@@ -55,7 +55,16 @@
     Scale,
     Sequence,
 )
-from .embed import ConcretizedCallable, Embedding
+from .embed import (
+    ConcretizedCallable,
+    Embedding,
+    cos,
+    log,
+    sin,
+    sqrt,
+    tan,
+    tanh,
+)
 from .hamiltonians import HamiltonianEvolution, Observable
 from .noise import (
     AmplitudeDamping,

pyqtorch/composite/compose.py

@@ -10,7 +10,7 @@
 from torch.nn import Module, ModuleList, ParameterDict
 
 from pyqtorch.apply import apply_operator
-from pyqtorch.embed import Embedding
+from pyqtorch.embed import ConcretizedCallable, Embedding
 from pyqtorch.matrices import add_batch_dim
 from pyqtorch.primitives import CNOT, RX, RY, Parametric, Primitive
 from pyqtorch.utils import (
@@ -35,7 +35,9 @@ class Scale(Sequence):
     """
 
     def __init__(
-        self, operations: Union[Primitive, Sequence, Add], param_name: str | Tensor
+        self,
+        operations: Union[Primitive, Sequence, Add],
+        param_name: str | float | int | Tensor | ConcretizedCallable,
     ):
         """
         Initializes a Scale object.
@@ -46,6 +48,11 @@ def __init__(
         """
         if not isinstance(operations, (Primitive, Sequence, Add)):
             raise ValueError("Scale only supports a single operation, Sequence or Add.")
+        if not isinstance(param_name, (str, int, float, Tensor, ConcretizedCallable)):
+            raise TypeError(
+                "Only str, int, float, Tensor or ConcretizedCallable types \
+                are supported for param_name"
+            )
         super().__init__([operations])
         self.param_name = param_name
 
@@ -69,12 +76,14 @@ def forward(
         if embedding is not None:
             values = embedding(values)
 
-        scale = (
-            values[self.param_name]
-            if isinstance(self.param_name, str)
-            else self.param_name
-        )
-        return scale * self.operations[0].forward(state, values, embedding)
+        if isinstance(self.param_name, str):
+            scale = values[self.param_name]
+        elif isinstance(self.param_name, Tensor):
+            scale = self.param_name
+        elif isinstance(self.param_name, ConcretizedCallable):
+            scale = self.param_name(values)
+
+        return scale * self.operations[0].forward(state, values)
 
     def tensor(
         self,
@@ -97,12 +106,14 @@ def tensor(
         if embedding is not None:
             values = embedding(values)
 
-        scale = (
-            values[self.param_name]
-            if isinstance(self.param_name, str)
-            else self.param_name
-        )
-        return scale * self.operations[0].tensor(values, embedding, full_support)
+        if isinstance(self.param_name, str):
+            scale = values[self.param_name]
+        elif isinstance(self.param_name, (Tensor, int, float)):
+            scale = self.param_name
+        elif isinstance(self.param_name, ConcretizedCallable):
+            scale = self.param_name(values)
+
+        return scale * self.operations[0].tensor(values, full_support=full_support)
 
     def flatten(self) -> list[Scale]:
         """This method should only be called in the AdjointExpectation,
@@ -121,7 +132,7 @@ def to(self, *args: Any, **kwargs: Any) -> Scale:
             Converted Scale.
         """
         super().to(*args, **kwargs)
-        if not isinstance(self.param_name, str):
+        if not isinstance(self.param_name, (str, float, int)):
             self.param_name = self.param_name.to(*args, **kwargs)
 
         return self

pyqtorch/differentiation/adjoint.py

@@ -130,8 +130,8 @@ def backward(ctx: Any, grad_out: Tensor) -> Tuple[None, ...]:
                         grad_out * 2 * inner_prod(ctx.projected_state, mu).real
                     )
 
-                if values[op.param_name].requires_grad:
-                    grads_dict[op.param_name] = grad_out * 2 * -values[op.param_name]
+                if values[op.param_name].requires_grad:  # type: ignore [index]
+                    grads_dict[op.param_name] = grad_out * 2 * -values[op.param_name]  # type: ignore [index]
                 ctx.projected_state = apply_operator(
                     ctx.projected_state,
                     op.dagger(values, ctx.embedding),

pyqtorch/embed.py

@@ -5,6 +5,7 @@
 from typing import Any, Tuple
 
 from numpy.typing import ArrayLike, DTypeLike
+from torch import Tensor
 
 logger = getLogger(__name__)
 
@@ -21,7 +22,7 @@
     "sub": ("jax.numpy", "subtract"),
     "div": ("jax.numpy", "divide"),
 }
-DEFAULT_TORCH_MAPPING: dict = dict()
+DEFAULT_TORCH_MAPPING = {"hs": ("pyqtorch.utils", "heaviside")}
 DEFAULT_NUMPY_MAPPING = {
     "mul": ("numpy", "multiply"),
     "sub": ("numpy", "subtract"),
@@ -74,8 +75,8 @@ class ConcretizedCallable:
 
     def __init__(
         self,
-        call_name: str,
-        abstract_args: list[str | float | int],
+        call_name: str = "",
+        abstract_args: list[str | float | int | complex | ConcretizedCallable] = ["x"],
         instruction_mapping: dict[str, Tuple[str, str]] = dict(),
         engine_name: str = "torch",
         device: str = "cpu",
@@ -92,6 +93,16 @@ def __init__(
         self._dtype = dtype
         self.engine_call = None
         engine = None
+        if not all(
+            [
+                isinstance(arg, (str, float, int, complex, Tensor, ConcretizedCallable))
+                for arg in abstract_args
+            ]
+        ):
+            raise TypeError(
+                "Only str, float, int, complex, Tensor or ConcretizedCallable type elements \
+                are supported for abstract_args"
+            )
         try:
             engine_name, fn_name = ARRAYLIKE_FN_MAP[engine_name]
             engine = import_module(engine_name)
@@ -113,17 +124,94 @@ def __init__(
     def evaluate(self, inputs: dict[str, ArrayLike] = dict()) -> ArrayLike:
         arraylike_args = []
         for symbol_or_numeric in self.abstract_args:
-            if isinstance(symbol_or_numeric, (float, int)):
+            if isinstance(symbol_or_numeric, ConcretizedCallable):
+                arraylike_args.append(symbol_or_numeric(inputs))
+            if isinstance(symbol_or_numeric, (float, int, Tensor)):
                 arraylike_args.append(
                     self.arraylike_fn(symbol_or_numeric, device=self.device)
                 )
             elif isinstance(symbol_or_numeric, str):
                 arraylike_args.append(inputs[symbol_or_numeric])
         return self.engine_call(*arraylike_args)  # type: ignore[misc]
 
+    @classmethod
+    def _get_independent_args(cls, cc: ConcretizedCallable) -> set:
+        out: set = set()
+        if len(cc.abstract_args) == 1 and isinstance(cc.abstract_args[0], str):
+            return set([cc.abstract_args[0]])
+        else:
+            for arg in cc.abstract_args:
+                if isinstance(arg, ConcretizedCallable):
+                    res = cls._get_independent_args(arg)
+                    out = out.union(res)
+                else:
+                    if isinstance(arg, str):
+                        out.add(arg)
+        return out
+
+    @property
+    def independent_args(self) -> list:
+        return list(self._get_independent_args(self))
+
     def __call__(self, inputs: dict[str, ArrayLike] = dict()) -> ArrayLike:
         return self.evaluate(inputs)
 
+    def __mul__(
+        self, other: str | int | float | complex | ConcretizedCallable
+    ) -> ConcretizedCallable:
+        return ConcretizedCallable("mul", [self, other])
+
+    def __rmul__(
+        self, other: str | int | float | ConcretizedCallable
+    ) -> ConcretizedCallable:
+        return ConcretizedCallable("mul", [other, self])
+
+    def __add__(
+        self, other: str | int | float | ConcretizedCallable
+    ) -> ConcretizedCallable:
+        return ConcretizedCallable("add", [self, other])
+
+    def __radd__(
+        self, other: str | int | float | ConcretizedCallable
+    ) -> ConcretizedCallable:
+        return ConcretizedCallable("add", [other, self])
+
+    def __sub__(
+        self, other: str | int | float | ConcretizedCallable
+    ) -> ConcretizedCallable:
+        return ConcretizedCallable("sub", [self, other])
+
+    def __rsub__(
+        self, other: str | int | float | ConcretizedCallable
+    ) -> ConcretizedCallable:
+        return ConcretizedCallable("sub", [other, self])
+
+    def __pow__(
+        self, other: str | int | float | ConcretizedCallable
+    ) -> ConcretizedCallable:
+        return ConcretizedCallable("pow", [self, other])
+
+    def __rpow__(
+        self, other: str | int | float | ConcretizedCallable
+    ) -> ConcretizedCallable:
+        return ConcretizedCallable("pow", [other, self])
+
+    def __truediv__(
+        self, other: str | int | float | ConcretizedCallable
+    ) -> ConcretizedCallable:
+        return ConcretizedCallable("div", [self, other])
+
+    def __rtruediv__(
+        self, other: str | int | float | ConcretizedCallable
+    ) -> ConcretizedCallable:
+        return ConcretizedCallable("div", [other, self])
+
+    def __repr__(self) -> str:
+        return f"{self.call_name}({self.abstract_args})"
+
+    def __neg__(self) -> ConcretizedCallable:
+        return -1 * self
+
     @property
     def device(self) -> str:
         return self._device
@@ -150,6 +238,30 @@ def init_param(
         return engine.random.uniform(0, 1)
 
 
+def sin(x: str | ConcretizedCallable):
+    return ConcretizedCallable("sin", [x])
+
+
+def cos(x: str | ConcretizedCallable):
+    return ConcretizedCallable("cos", [x])
+
+
+def log(x: str | ConcretizedCallable):
+    return ConcretizedCallable("log", [x])
+
+
+def tan(x: str | ConcretizedCallable):
+    return ConcretizedCallable("tan", [x])
+
+
+def tanh(x: str | ConcretizedCallable):
+    return ConcretizedCallable("tanh", [x])
+
+
+def sqrt(x: str | ConcretizedCallable):
+    return ConcretizedCallable("sqrt", [x])
+
+
 class Embedding:
     """A class relating variational and feature parameters used in ConcretizedCallable instances to
     parameter names used in gates.