Repository: PennyLaneAI/pennylane
Pull request #4911 (2023-12-01 18:44:48)

lazy-load tape and operator batch sizes
**Context:**
We compute operator batch size and initialization, and it's oftentimes unnecessary. We should only compute it when it's requested

**Description of the Change:**
Only compute operator/tape batch size when they are requested.
Some details:
- inlined the batch-size computation for `ScalarSymbolicOp`
- removed the `_check_batching` overrides in controlled and adjoint because they are not needed (codecov complained). They both inherit from `SymbolicOp`, and that doesn't call `super().__init__` so it was never computing it anyway

**Benefits:**
- Will run expensive computations less frequently
- the Tensorflow issue when pre-computing batch sizes has disappeared!

**Possible Drawbacks:**
the `_check_batching` function did some data validation as well. operators initialized with poorly-shaped data would fail right away, but now we'll only see the error when the batch size is requested

**Related GitHub Issues:**
[sc-45972]

github-actions[bot]: Hello. You may have forgotten to update the changelog!
Please edit [`doc/releases/changelog-dev.md`](/PennyLaneAI/pennylane/blob/master/doc/releases/changelog-dev.md) with:
* A one-to-two sentence description of the change. You may include a small working example for new features.
* A link back to this PR.
* Your name (or GitHub username) in the contributors section.
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4911?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Comparison is base [(`3c88f12`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/3c88f122101e4982d41457f3196fd6bdc73b750e?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.65% compared to head [(`89d15da`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4911?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.63%.
> Report is 1 commits behind head on master.


<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##           master    #4911      +/-   ##
==========================================
- Coverage   99.65%   99.63%   -0.03%     
==========================================
  Files         387      387              
  Lines       35037    34795     -242     
==========================================
- Hits        34917    34668     -249     
- Misses        120      127       +7     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4911?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).


Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 442410e91..93028f61e 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -195,6 +195,9 @@
 * `qml.ArbitraryUnitary` now supports batching.
   [(#4745)](https://github.com/PennyLaneAI/pennylane/pull/4745)
 
+* Operator and tape batch sizes are evaluated lazily.
+  [(#4911)](https://github.com/PennyLaneAI/pennylane/pull/4911)
+
 <h4>Performance improvements and benchmarking</h4>
 
 * `default.qubit` no longer uses a dense matrix for `MultiControlledX` for more than 8 operation wires.
diff --git a/pennylane/_device.py b/pennylane/_device.py
index 95fc190b6..0402b7a91 100644
--- a/pennylane/_device.py
+++ b/pennylane/_device.py
@@ -100,8 +100,8 @@ def _local_tape_expand(tape, depth, stop_at):
     # Update circuit info
     new_tape.wires = copy.copy(tape.wires)
     new_tape.num_wires = tape.num_wires
-    new_tape._batch_size = tape.batch_size
-    new_tape._output_dim = tape.output_dim
+    new_tape._batch_size = tape._batch_size
+    new_tape._output_dim = tape._output_dim
     return new_tape
 
 
diff --git a/pennylane/operation.py b/pennylane/operation.py
index 18b112923..eab7d266a 100644
--- a/pennylane/operation.py
+++ b/pennylane/operation.py
@@ -267,6 +267,7 @@ from .pytrees import register_pytree
 
 SUPPORTED_INTERFACES = {"numpy", "scipy", "autograd", "torch", "tensorflow", "jax"}
 __use_new_opmath = False
+_UNSET_BATCH_SIZE = -1  # indicates that the (lazy) batch size has not yet been accessed/computed
 
 
 class OperatorPropertyUndefined(Exception):
@@ -662,10 +663,10 @@ class Operator(abc.ABC):
         the ``_check_batching`` method, by comparing the shape of the input data to
         the expected shape. Therefore, it is necessary to call ``_check_batching`` on
         any new input parameters passed to the operator. By default, any class inheriting
-        from :class:`~.operation.Operator` will do so within its ``__init__`` method, and
-        other objects may do so when updating the data of the ``Operator``.
+        from :class:`~.operation.Operator` will do so the first time its
+        ``batch_size`` property is accessed.
 
-        ``_check_batching`` modifies the following class attributes:
+        ``_check_batching`` modifies the following instance attributes:
 
         - ``_ndim_params``: The number of dimensions of the parameters passed to
           ``_check_batching``. For an ``Operator`` that does _not_ set the ``ndim_params``
@@ -678,9 +679,9 @@ class Operator(abc.ABC):
           not support broadcasting will report to not be broadcasted independently of the
           input.
 
-        Both attributes are not defined if ``_check_batching`` is not called. Therefore it
-        *needs to be called* within custom ``__init__`` implementations, either directly
-        or by calling ``Operator.__init__``.
+        These two properties are defined lazily, and accessing the public version of either
+        one of them (in other words, without the leading underscore) for the first time will
+        trigger a call to ``_check_batching``, which validates and sets these properties.
     """
     # pylint: disable=too-many-public-methods, too-many-instance-attributes
 
@@ -1076,39 +1077,25 @@ class Operator(abc.ABC):
                 f"{len(self._wires)} wires given, {self.num_wires} expected."
             )
 
-        self._check_batching(params)
+        self._batch_size = _UNSET_BATCH_SIZE
+        self._ndim_params = _UNSET_BATCH_SIZE
 
         self.data = tuple(np.array(p) if isinstance(p, (list, tuple)) else p for p in params)
 
         self.queue()
 
-    def _check_batching(self, params):
+    def _check_batching(self):
         """Check if the expected numbers of dimensions of parameters coincides with the
         ones received and sets the ``_batch_size`` attribute.
 
-        Args:
-            params (tuple): Parameters with which the operator is instantiated
-
         The check always passes and sets the ``_batch_size`` to ``None`` for the default
         ``Operator.ndim_params`` property but subclasses may overwrite it to define fixed
         expected numbers of dimensions, allowing to infer a batch size.
         """
         self._batch_size = None
+        params = self.data
 
-        try:
-            ndims = tuple(qml.math.ndim(p) for p in params)
-        except ValueError as e:
-            # TODO:[dwierichs] When using tf.function with an input_signature that contains
-            # an unknown-shaped input, ndim() will not be able to determine the number of
-            # dimensions because they are not specified yet. Failing example: Let `fun` be
-            # a single-parameter QNode.
-            # `tf.function(fun, input_signature=(tf.TensorSpec(shape=None, dtype=tf.float32),))`
-            # There might be a way to support batching nonetheless, which remains to be
-            # investigated. For now, the batch_size is left to be `None` when instantiating
-            # an operation with abstract parameters that make `qml.math.ndim` fail.
-            if any(qml.math.is_abstract(p) for p in params):
-                return
-            raise e
+        ndims = tuple(qml.math.ndim(p) for p in params)
 
         if any(len(qml.math.shape(p)) >= 1 and qml.math.shape(p)[0] is None for p in params):
             # if the batch dimension is unknown, then skip the validation
@@ -1192,6 +1179,8 @@ class Operator(abc.ABC):
         Returns:
             tuple: Number of dimensions for each trainable parameter.
         """
+        if self._batch_size is _UNSET_BATCH_SIZE:
+            self._check_batching()
         return self._ndim_params
 
     @property
@@ -1206,6 +1195,8 @@ class Operator(abc.ABC):
         Returns:
             int or None: Size of the parameter broadcasting dimension if present, else ``None``.
         """
+        if self._batch_size is _UNSET_BATCH_SIZE:
+            self._check_batching()
         return self._batch_size
 
     @property
diff --git a/pennylane/ops/op_math/adjoint.py b/pennylane/ops/op_math/adjoint.py
index e9fff121a..1b2fc3942 100644
--- a/pennylane/ops/op_math/adjoint.py
+++ b/pennylane/ops/op_math/adjoint.py
@@ -288,10 +288,6 @@ class Adjoint(SymbolicOp):
     def __repr__(self):
         return f"Adjoint({self.base})"
 
-    # pylint: disable=protected-access
-    def _check_batching(self, params):
-        self.base._check_batching(params)
-
     @property
     def ndim_params(self):
         return self.base.ndim_params
diff --git a/pennylane/ops/op_math/composite.py b/pennylane/ops/op_math/composite.py
index b3bf7750c..f319bda29 100644
--- a/pennylane/ops/op_math/composite.py
+++ b/pennylane/ops/op_math/composite.py
@@ -23,7 +23,7 @@ import numpy as np
 
 import pennylane as qml
 from pennylane import math
-from pennylane.operation import Operator
+from pennylane.operation import Operator, _UNSET_BATCH_SIZE
 from pennylane.wires import Wires
 
 # pylint: disable=too-many-instance-attributes
@@ -67,9 +67,9 @@ class CompositeOp(Operator):
         self._overlapping_ops = None
         self._pauli_rep = self._build_pauli_rep()
         self.queue()
-        self._check_batching(None)  # unused param
+        self._batch_size = _UNSET_BATCH_SIZE
 
-    def _check_batching(self, _):
+    def _check_batching(self):
         batch_sizes = {op.batch_size for op in self if op.batch_size is not None}
         if len(batch_sizes) > 1:
             raise ValueError(
diff --git a/pennylane/ops/op_math/controlled.py b/pennylane/ops/op_math/controlled.py
index f47c1944b..ebb6a55e7 100644
--- a/pennylane/ops/op_math/controlled.py
+++ b/pennylane/ops/op_math/controlled.py
@@ -365,10 +365,6 @@ class Controlled(SymbolicOp):
     def has_matrix(self):
         return self.base.has_matrix
 
-    # pylint: disable=protected-access
-    def _check_batching(self, params):
-        self.base._check_batching(params)
-
     @property
     def batch_size(self):
         return self.base.batch_size
diff --git a/pennylane/ops/op_math/symbolicop.py b/pennylane/ops/op_math/symbolicop.py
index 3ad23f2e7..916bbc8a4 100644
--- a/pennylane/ops/op_math/symbolicop.py
+++ b/pennylane/ops/op_math/symbolicop.py
@@ -20,7 +20,7 @@ from copy import copy
 import numpy as np
 
 import pennylane as qml
-from pennylane.operation import Operator
+from pennylane.operation import Operator, _UNSET_BATCH_SIZE
 from pennylane.queuing import QueuingManager
 
 
@@ -166,10 +166,24 @@ class ScalarSymbolicOp(SymbolicOp):
     def __init__(self, base, scalar: float, id=None):
         self.scalar = np.array(scalar) if isinstance(scalar, list) else scalar
         super().__init__(base, id=id)
-        self._batch_size = self._check_and_compute_batch_size(scalar)
+        self._batch_size = _UNSET_BATCH_SIZE
 
     @property
     def batch_size(self):
+        if self._batch_size is _UNSET_BATCH_SIZE:
+            base_batch_size = self.base.batch_size
+            if qml.math.ndim(self.scalar) == 0:
+                # coeff is not batched
+                self._batch_size = base_batch_size
+            else:
+                # coeff is batched
+                scalar_size = qml.math.size(self.scalar)
+                if base_batch_size is not None and base_batch_size != scalar_size:
+                    raise ValueError(
+                        "Broadcasting was attempted but the broadcasted dimensions "
+                        f"do not match: {scalar_size}, {base_batch_size}."
+                    )
+                self._batch_size = scalar_size
         return self._batch_size
 
     @property
@@ -181,20 +195,6 @@ class ScalarSymbolicOp(SymbolicOp):
         self.scalar = new_data[0]
         self.base.data = new_data[1:]
 
-    def _check_and_compute_batch_size(self, scalar):
-        batched_scalar = qml.math.ndim(scalar) > 0
-        scalar_size = qml.math.size(scalar)
-        if not batched_scalar:
-            # coeff is not batched
-            return self.base.batch_size
-        # coeff is batched
-        if self.base.batch_size is not None and self.base.batch_size != scalar_size:
-            raise ValueError(
-                "Broadcasting was attempted but the broadcasted dimensions "
-                f"do not match: {scalar_size}, {self.base.batch_size}."
-            )
-        return scalar_size
-
     @property
     def has_matrix(self):
         return self.base.has_matrix or isinstance(self.base, qml.Hamiltonian)
diff --git a/pennylane/ops/qubit/hamiltonian.py b/pennylane/ops/qubit/hamiltonian.py
index 1f0e1dff9..f1d991536 100644
--- a/pennylane/ops/qubit/hamiltonian.py
+++ b/pennylane/ops/qubit/hamiltonian.py
@@ -228,7 +228,7 @@ class Hamiltonian(Observable):
         # while H.coeffs is the original tensor
         super().__init__(*coeffs_flat, wires=self._wires, id=id)
 
-    def _check_batching(self, params):
+    def _check_batching(self):
         """Override for Hamiltonian, batching is not yet supported."""
 
     def label(self, decimals=None, base_label=None, cache=None):
diff --git a/pennylane/tape/qscript.py b/pennylane/tape/qscript.py
index 4e3f5d9bb..e10957c67 100644
--- a/pennylane/tape/qscript.py
+++ b/pennylane/tape/qscript.py
@@ -35,7 +35,7 @@ from pennylane.measurements import (
     Shots,
 )
 from pennylane.typing import TensorLike
-from pennylane.operation import Observable, Operator, Operation
+from pennylane.operation import Observable, Operator, Operation, _UNSET_BATCH_SIZE
 from pennylane.pytrees import register_pytree
 from pennylane.queuing import AnnotatedQueue, process_queue
 from pennylane.wires import Wires
@@ -198,8 +198,8 @@ class QuantumScript:
         self._trainable_params = trainable_params
         self._graph = None
         self._specs = None
-        self._output_dim = 0
-        self._batch_size = None
+        self._output_dim = None
+        self._batch_size = _UNSET_BATCH_SIZE
 
         self.wires = _empty_wires
         self.num_wires = 0
@@ -347,11 +347,15 @@ class QuantumScript:
         Returns:
             int or None: The batch size of the quantum script if present, else ``None``.
         """
+        if self._batch_size is _UNSET_BATCH_SIZE:
+            self._update_batch_size()
         return self._batch_size
 
     @property
     def output_dim(self):
         """The (inferred) output dimension of the quantum script."""
+        if self._output_dim is None:
+            self._update_output_dim()  # this will set _batch_size if it isn't already
         return self._output_dim
 
     @property
@@ -434,10 +438,6 @@ class QuantumScript:
         self._update_par_info()  # Updates _par_info; O(ops+obs)
 
         self._update_observables()  # Updates _obs_sharing_wires and _obs_sharing_wires_id
-        self._update_batch_size()  # Updates _batch_size; O(ops)
-
-        # The following line requires _batch_size to be up to date
-        self._update_output_dim()  # Updates _output_dim; O(obs)
 
     def _update_circuit_info(self):
         """Update circuit metadata
@@ -890,8 +890,8 @@ class QuantumScript:
         new_qscript._update_par_info()
         new_qscript._obs_sharing_wires = self._obs_sharing_wires
         new_qscript._obs_sharing_wires_id = self._obs_sharing_wires_id
-        new_qscript._batch_size = self.batch_size
-        new_qscript._output_dim = self.output_dim
+        new_qscript._batch_size = self._batch_size
+        new_qscript._output_dim = self._output_dim
 
         return new_qscript
 
diff --git a/pennylane/tape/tape.py b/pennylane/tape/tape.py
index 7b58abe9f..17ab58f13 100644
--- a/pennylane/tape/tape.py
+++ b/pennylane/tape/tape.py
@@ -227,8 +227,8 @@ def expand_tape(tape, depth=1, stop_at=None, expand_measurements=False):
     # Update circuit info
     new_tape.wires = copy.copy(tape.wires)
     new_tape.num_wires = tape.num_wires
-    new_tape._batch_size = tape.batch_size
-    new_tape._output_dim = tape.output_dim
+    new_tape._batch_size = tape._batch_size
+    new_tape._output_dim = tape._output_dim
     return new_tape
 
 
@@ -278,8 +278,8 @@ def expand_tape_state_prep(tape, skip_first=True):
     # Update circuit info
     new_tape.wires = copy.copy(tape.wires)
     new_tape.num_wires = tape.num_wires
-    new_tape._batch_size = tape.batch_size
-    new_tape._output_dim = tape.output_dim
+    new_tape._batch_size = tape._batch_size
+    new_tape._output_dim = tape._output_dim
     return new_tape
 
 


Test patch:
diff --git a/tests/devices/qubit/test_apply_operation.py b/tests/devices/qubit/test_apply_operation.py
index a21ad2b3a..9e61aba35 100644
--- a/tests/devices/qubit/test_apply_operation.py
+++ b/tests/devices/qubit/test_apply_operation.py
@@ -20,6 +20,7 @@ import pytest
 import numpy as np
 from scipy.stats import unitary_group
 import pennylane as qml
+from pennylane.operation import _UNSET_BATCH_SIZE, Operation
 
 
 from pennylane.devices.qubit.apply_operation import (
@@ -51,7 +52,7 @@ def test_custom_operator_with_matrix():
     )
 
     # pylint: disable=too-few-public-methods
-    class CustomOp(qml.operation.Operation):
+    class CustomOp(Operation):
         num_wires = 1
 
         def matrix(self):
@@ -800,7 +801,7 @@ class TestBroadcasting:  # pylint: disable=too-few-public-methods
         param = qml.math.asarray([0.1, 0.2, 0.3], like=ml_framework)
         state = np.ones((2, 2)) / 2
         op = qml.RX(param, 0)
-        op._batch_size = None  # pylint:disable=protected-access
+        assert op._batch_size is _UNSET_BATCH_SIZE  # pylint:disable=protected-access
         state = method(op, state)
         assert state.shape == (3, 2, 2)
         assert op.batch_size == 3
diff --git a/tests/ops/op_math/test_composite.py b/tests/ops/op_math/test_composite.py
index 89ed5d737..f974ef3dc 100644
--- a/tests/ops/op_math/test_composite.py
+++ b/tests/ops/op_math/test_composite.py
@@ -143,10 +143,11 @@ class TestConstruction:
     def test_different_batch_sizes_raises_error(self):
         """Test that an error is raised if the operands have different batch sizes."""
         base = qml.RX(np.array([1.2, 2.3, 3.4]), 0)
+        op = ValidOp(base, qml.RY(1, 0), qml.RZ(np.array([1, 2, 3, 4]), wires=2))
         with pytest.raises(
             ValueError, match="Broadcasting was attempted but the broadcasted dimensions"
         ):
-            _ = ValidOp(base, qml.RY(1, 0), qml.RZ(np.array([1, 2, 3, 4]), wires=2))
+            _ = op.batch_size
 
     def test_decomposition_raises_error(self):
         """Test that calling decomposition() raises a ValueError."""
diff --git a/tests/ops/op_math/test_exp.py b/tests/ops/op_math/test_exp.py
index 70da3f7fe..dcb1128ff 100644
--- a/tests/ops/op_math/test_exp.py
+++ b/tests/ops/op_math/test_exp.py
@@ -160,10 +160,11 @@ class TestProperties:
     def test_different_batch_sizes_raises_error(self):
         """Test that using different batch sizes for base and scalar raises an error."""
         base = qml.RX(np.array([1.2, 2.3, 3.4]), 0)
+        op = Exp(base, np.array([0.1, 1.2, 2.3, 3.4]))
         with pytest.raises(
             ValueError, match="Broadcasting was attempted but the broadcasted dimensions"
         ):
-            _ = Exp(base, np.array([0.1, 1.2, 2.3, 3.4]))
+            _ = op.batch_size
 
 
 class TestMatrix:
diff --git a/tests/ops/op_math/test_pow_op.py b/tests/ops/op_math/test_pow_op.py
index 7ed54391d..828f453b7 100644
--- a/tests/ops/op_math/test_pow_op.py
+++ b/tests/ops/op_math/test_pow_op.py
@@ -374,10 +374,11 @@ class TestProperties:
     def test_different_batch_sizes_raises_error(self, power_method):
         """Test that using different batch sizes for base and scalar raises an error."""
         base = qml.RX(np.array([1.2, 2.3, 3.4]), 0)
+        op = power_method(base, np.array([0.1, 1.2, 2.3, 3.4]))
         with pytest.raises(
             ValueError, match="Broadcasting was attempted but the broadcasted dimensions"
         ):
-            _ = power_method(base, np.array([0.1, 1.2, 2.3, 3.4]))
+            _ = op.batch_size
 
     op_pauli_reps = (
         (qml.PauliZ(wires=0), 1, qml.pauli.PauliSentence({qml.pauli.PauliWord({0: "Z"}): 1})),
diff --git a/tests/ops/op_math/test_sprod.py b/tests/ops/op_math/test_sprod.py
index 0024e573b..bc219a491 100644
--- a/tests/ops/op_math/test_sprod.py
+++ b/tests/ops/op_math/test_sprod.py
@@ -771,10 +771,11 @@ class TestProperties:
     def test_different_batch_sizes_raises_error(self):
         """Test that using different batch sizes for base and scalar raises an error."""
         base = qml.RX(np.array([1.2, 2.3, 3.4]), 0)
+        op = qml.s_prod(np.array([0.1, 1.2, 2.3, 3.4]), base)
         with pytest.raises(
             ValueError, match="Broadcasting was attempted but the broadcasted dimensions"
         ):
-            _ = qml.s_prod(np.array([0.1, 1.2, 2.3, 3.4]), base)
+            _ = op.batch_size
 
 
 class TestSimplify:
diff --git a/tests/tape/test_qscript.py b/tests/tape/test_qscript.py
index 1da39172d..a1daf08b4 100644
--- a/tests/tape/test_qscript.py
+++ b/tests/tape/test_qscript.py
@@ -19,6 +19,7 @@ import pytest
 
 import pennylane as qml
 from pennylane.measurements import MutualInfo, State, VnEntropy, Shots
+from pennylane.operation import _UNSET_BATCH_SIZE
 from pennylane.tape import QuantumScript
 
 # pylint: disable=protected-access, unused-argument, too-few-public-methods
@@ -40,7 +41,8 @@ class TestInitialization:
         assert qs._graph is None
         assert qs._specs is None
         assert qs._shots.total_shots is None
-        assert qs._batch_size is None
+        assert qs._batch_size is _UNSET_BATCH_SIZE
+        assert qs.batch_size is None
         assert qs.wires == qml.wires.Wires([])
         assert qs.num_wires == 0
         assert qs.samples_computational_basis is False
@@ -270,11 +272,11 @@ class TestUpdate:
     def test_error_inconsistent_batch_sizes(self, x, rot, y):
         """Tests that an error is raised if inconsistent batch sizes exist."""
         ops = [qml.RX(x, wires=0), qml.Rot(*rot, 1), qml.RX(y, wires=1)]
-
+        tape = QuantumScript(ops)
         with pytest.raises(
             ValueError, match="batch sizes of the quantum script operations do not match."
         ):
-            _ = QuantumScript(ops)
+            _ = tape.batch_size
 
     @pytest.mark.parametrize(
         "m, output_dim",
@@ -292,6 +294,37 @@ class TestUpdate:
         qs = QuantumScript(ops, m)
         assert qs.output_dim == output_dim * factor
 
+    def test_lazy_batch_size_and_output_dim(self):
+        """Test that batch_size and output_dim are computed lazily."""
+        qs = QuantumScript([qml.RX([1.1, 2.2], 0)], [qml.expval(qml.PauliZ(0))])
+        copied = qs.copy()
+        assert qs._batch_size is _UNSET_BATCH_SIZE
+        assert qs._output_dim is None
+        # copying did not evaluate them either
+        assert copied._batch_size is _UNSET_BATCH_SIZE
+        assert copied._output_dim is None
+
+        # now evaluate it
+        assert qs.batch_size == 2
+        assert qs._output_dim is None  # setting batch_size didn't set output_dim
+        assert qs.output_dim == 2
+        copied = qs.copy()
+        assert qs._batch_size == 2
+        assert qs._output_dim == 2
+        # copied tape has it pre-evaluated
+        assert copied._batch_size == 2
+        assert copied._output_dim == 2
+
+    def test_lazy_setting_output_dim_sets_batch_size(self):
+        """Test that setting the output_dim also sets the batch_size."""
+        qs = QuantumScript([qml.RX([1.1, 2.2], 0)], [qml.expval(qml.PauliZ(0))])
+        assert qs._batch_size is _UNSET_BATCH_SIZE
+        assert qs._output_dim is None
+
+        assert qs.output_dim == 2  # getting this sets both _output_dim and _batch_size
+        assert qs._output_dim == 2
+        assert qs._batch_size == 2
+
 
 class TestIteration:
     """Test the capabilities related to iterating over quantum script."""
diff --git a/tests/tape/test_tape.py b/tests/tape/test_tape.py
index 92469a5d4..16f00f149 100644
--- a/tests/tape/test_tape.py
+++ b/tests/tape/test_tape.py
@@ -31,7 +31,7 @@ from pennylane.measurements import (
     sample,
     var,
 )
-from pennylane.tape import QuantumTape, expand_tape_state_prep
+from pennylane.tape import QuantumTape, QuantumScript, expand_tape_state_prep
 
 
 def TestOperationMonkeypatching():
@@ -91,7 +91,6 @@ class TestConstruction:
 
         assert tape.wires == qml.wires.Wires([0, "a", 4])
         assert tape._output_dim == len(obs[0].wires) + 2 ** len(obs[1].wires)
-        assert tape._batch_size is None
 
     def test_observable_processing(self, make_tape):
         """Test that observables are processed correctly"""
@@ -358,26 +357,24 @@ class TestConstruction:
         """Test that the batch size is correctly inferred from all operation's
         batch_size, when creating and when using `bind_new_parameters`."""
 
+        tape = QuantumScript(
+            [qml.RX(x, wires=0), qml.Rot(*rot, wires=1), qml.RX(y, wires=1)],
+            [qml.expval(qml.PauliZ(0))],
+        )
         with pytest.raises(
             ValueError, match="batch sizes of the quantum script operations do not match."
         ):
-            with qml.queuing.AnnotatedQueue() as q:
-                qml.RX(x, wires=0)
-                qml.Rot(*rot, wires=1)
-                qml.RX(y, wires=1)
-                qml.apply(qml.expval(qml.PauliZ(0)))
+            _ = tape.batch_size
 
-            tape = qml.tape.QuantumScript.from_queue(q)
-        with qml.queuing.AnnotatedQueue() as q:
-            qml.RX(0.2, wires=0)
-            qml.Rot(1.0, 0.2, -0.3, wires=1)
-            qml.RX(0.2, wires=1)
-            qml.apply(qml.expval(qml.PauliZ(0)))
-        tape = qml.tape.QuantumScript.from_queue(q)
+        tape = QuantumScript(
+            [qml.RX(0.2, wires=0), qml.Rot(1.0, 0.2, -0.3, wires=1), qml.RX(0.2, wires=1)],
+            [qml.expval(qml.PauliZ(0))],
+        )
+        tape = tape.bind_new_parameters([x] + rot + [y], [0, 1, 2, 3, 4])
         with pytest.raises(
             ValueError, match="batch sizes of the quantum script operations do not match."
         ):
-            tape.bind_new_parameters([x] + rot + [y], [0, 1, 2, 3, 4])
+            _ = tape.batch_size
 
 
 class TestIteration:
diff --git a/tests/test_operation.py b/tests/test_operation.py
index 06e9f0276..038ae960b 100644
--- a/tests/test_operation.py
+++ b/tests/test_operation.py
@@ -25,7 +25,14 @@ from numpy.linalg import multi_dot
 
 import pennylane as qml
 from pennylane import numpy as pnp
-from pennylane.operation import Operation, Operator, StatePrepBase, Tensor, operation_derivative
+from pennylane.operation import (
+    Operation,
+    Operator,
+    StatePrepBase,
+    Tensor,
+    operation_derivative,
+    _UNSET_BATCH_SIZE,
+)
 from pennylane.ops import Prod, SProd, Sum, cv
 from pennylane.wires import Wires
 
@@ -147,8 +154,9 @@ class TestOperatorConstruction:
             grad_method = "A"
             ndim_params = (0,)
 
+        op = DummyOp([[[0.5], [0.1]]], wires=0)
         with pytest.raises(ValueError, match=r"wrong number\(s\) of dimensions in parameters"):
-            DummyOp([[[0.5], [0.1]]], wires=0)
+            _ = op.batch_size
 
         op = DummyOp(0.5, wires=0)
         assert op.ndim_params == (0,)
@@ -159,8 +167,9 @@ class TestOperatorConstruction:
             num_wires = 1
             grad_method = "A"
 
+        op = DummyOp2([0.5], 0.6, wires=0)
         with pytest.raises(ValueError, match=r"wrong number\(s\) of dimensions in parameters"):
-            DummyOp2([0.5], 0.6, wires=0)
+            _ = op.batch_size
 
         op2 = DummyOp2([0.1], [[0.4, 0.1], [0.2, 1.2]], wires=0)
         assert op2.ndim_params == (1, 2)
@@ -183,8 +192,9 @@ class TestOperatorConstruction:
             num_wires = 1
 
         # Test with mismatching batch dimensions
+        op = DummyOp4([0.3] * 4, [[[0.3, 1.2]]] * 3, wires=0)
         with pytest.raises(ValueError, match="Broadcasting was attempted but the broadcasted"):
-            DummyOp4([0.3] * 4, [[[0.3, 1.2]]] * 3, wires=0)
+            _ = op.batch_size
 
     def test_default_pauli_rep(self):
         """Test that the default _pauli_rep attribute is None"""
@@ -299,6 +309,22 @@ class TestOperatorConstruction:
         assert hash(op2) == op2.hash
         assert hash(op1) == hash(op2)
 
+    @pytest.mark.parametrize("data,batch_size,ndim_params", [(1.1, None, 0), ([1.1, 2.2], 2, 1)])
+    def test_lazy_ndim_params_and_batch_size(self, data, batch_size, ndim_params):
+        """Test that ndim_params and batch_size are lazy properties."""
+
+        class DummyOp(Operator):
+            num_wires = 1
+            num_params = 1
+            ndim_params = (0,)
+
+        op = DummyOp(data, wires=[0])
+        assert op._batch_size is _UNSET_BATCH_SIZE
+        assert op._ndim_params is _UNSET_BATCH_SIZE
+        assert op.batch_size == batch_size
+        assert op._ndim_params == (ndim_params,)
+        assert op.ndim_params == (0,)
+
 
 class TestPytreeMethods:
     def test_pytree_defaults(self):
@@ -357,7 +383,7 @@ class TestBroadcasting:
 
         op = DummyOp(*params, wires=0)
         assert op.ndim_params == (0, 2)
-        assert op._batch_size == exp_batch_size
+        assert op.batch_size == exp_batch_size
 
     @pytest.mark.autograd
     @pytest.mark.parametrize("params, exp_batch_size", broadcasted_params_test_data)
@@ -373,7 +399,7 @@ class TestBroadcasting:
         params = tuple(pnp.array(p, requires_grad=True) for p in params)
         op = DummyOp(*params, wires=0)
         assert op.ndim_params == (0, 2)
-        assert op._batch_size == exp_batch_size
+        assert op.batch_size == exp_batch_size
 
     @pytest.mark.jax
     @pytest.mark.parametrize("params, exp_batch_size", broadcasted_params_test_data)
@@ -390,7 +416,7 @@ class TestBroadcasting:
         params = tuple(jax.numpy.array(p) for p in params)
         op = DummyOp(*params, wires=0)
         assert op.ndim_params == (0, 2)
-        assert op._batch_size == exp_batch_size
+        assert op.batch_size == exp_batch_size
 
     @pytest.mark.tf
     @pytest.mark.parametrize("params, exp_batch_size", broadcasted_params_test_data)
@@ -407,7 +433,7 @@ class TestBroadcasting:
         params = tuple(tf.Variable(p) for p in params)
         op = DummyOp(*params, wires=0)
         assert op.ndim_params == (0, 2)
-        assert op._batch_size == exp_batch_size
+        assert op.batch_size == exp_batch_size
 
     @pytest.mark.torch
     @pytest.mark.parametrize("params, exp_batch_size", broadcasted_params_test_data)
@@ -424,7 +450,7 @@ class TestBroadcasting:
         params = tuple(torch.tensor(p, requires_grad=True) for p in params)
         op = DummyOp(*params, wires=0)
         assert op.ndim_params == (0, 2)
-        assert op._batch_size == exp_batch_size
+        assert op.batch_size == exp_batch_size
 
     @pytest.mark.tf
     @pytest.mark.parametrize("jit_compile", [True, False])
@@ -436,6 +462,7 @@ class TestBroadcasting:
         class MyRX(qml.RX):
             @property
             def ndim_params(self):
+                self._check_batching()
                 return self._ndim_params
 
         def fun(x):