Improve Documentation for PQC and Expectation Layers (GH#238)

tensorflow_quantum/python/layers/circuit_executors/expectation.py

@@ -30,6 +30,30 @@
 class Expectation(tf.keras.layers.Layer):
     """A Layer that calculates an expectation value.
 
+    Args:
+        backend (Union[str, cirq.Sampler,
+        cirq.sim.simulator.SimulatesExpectationValues], optional):
+            Backend simulator to use. Default is 'noiseless'.
+        differentiator (Optional[tfq.differentiators.Differentiator]):
+            Differentiation scheme for gradients.
+        **kwargs: Additional keyword arguments for the parent class.
+
+    Input shape:
+        - circuits: tf.Tensor or list of shape [batch_size], each entry a
+          serialized circuit (from tfq.convert_to_tensor).
+        - symbol_names: list or tf.Tensor of shape [n_symbols], names of
+          circuit parameters (optional).
+        - symbol_values: tf.Tensor of shape [batch_size, n_symbols], values for
+          circuit parameters (optional).
+        - operators: list or tf.Tensor of shape [n_ops], observables to measure
+          (optional).
+        - repetitions: int or tf.Tensor, number of measurement repetitions
+          (optional, only for noisy backend).
+
+    Output shape:
+        tf.Tensor of shape [batch_size, n_ops], expectation values for each
+        circuit and operator.
+
     Given an input circuit and set of parameter values, prepare a quantum state
     and output expectation values taken on that state with respect to some
     observables to the tensorflow graph.
@@ -199,10 +223,9 @@ class Expectation(tf.keras.layers.Layer):
     something like `tfq.layers.Expectation()(cirq.Circuit(...), ...)` please
     be sure to instead use `tfq.layers.Expectation()(circuit_input, ...)`
     where `circuit_input` is a `tf.keras.Input` that is filled with
-    `tfq.conver_to_tensor([cirq.Circuit(..)] * batch_size)` at runtime. This
+    `tfq.convert_to_tensor([cirq.Circuit(..)] * batch_size)` at runtime. This
     is because compiled Keras models require non keyword layer `call` inputs to
     be traceable back to a `tf.keras.Input`.
-
     """
 
     def __init__(self, backend='noiseless', differentiator=None, **kwargs):
@@ -225,6 +248,7 @@ def __init__(self, backend='noiseless', differentiator=None, **kwargs):
                 which uses `tfq.differentiators.ParameterShift()`. If
                 `backend` is also 'noiseless' then default is
                 `tfq.differentiators.Adjoint`.
+            **kwargs: Additional keyword arguments for the parent class.
 
         """
         super().__init__(**kwargs)
@@ -273,15 +297,23 @@ def call(self,
              initializer=tf.keras.initializers.RandomUniform(0, 2 * np.pi)):
         """Keras call function.
 
-        Input options:
-            `inputs`, `symbol_names`, `symbol_values`:
-                see `input_checks.expand_circuits`
-            `operators`: see `input_checks.expand_operators`
-
-        Output shape:
-            `tf.Tensor` with shape [batch_size, n_ops] that holds the
-                expectation value for each circuit with each op applied to it
-                (after resolving the corresponding parameters in).
+        Args:
+            inputs (tf.Tensor or list): Circuits to execute, shape [batch_size].
+            symbol_names (list or tf.Tensor, optional): Names of circuit
+                parameters, shape [n_symbols].
+            symbol_values (tf.Tensor, optional): Values for circuit parameters,
+                shape [batch_size, n_symbols].
+            operators (list or tf.Tensor, optional): Observables to measure,
+                shape [n_ops].
+            repetitions (int or tf.Tensor, optional): Number of measurement
+                repetitions (for noisy backend).
+            initializer (tf.keras.initializers.Initializer, optional):
+                Initializer for circuit parameters.
+
+        Returns:
+            tf.Tensor: Tensor of shape [batch_size, n_ops] that holds the
+            expectation value for each circuit with each op applied to it
+            (after resolving the corresponding parameters in).
         """
         values_empty = False
         if symbol_values is None:

tensorflow_quantum/python/layers/high_level/pqc.py

@@ -28,6 +28,35 @@
 class PQC(tf.keras.layers.Layer):
     """Parametrized Quantum Circuit (PQC) Layer.
 
+    Args:
+        model_circuit (cirq.Circuit):
+            Circuit with sympy.Symbols to be trained.
+        operators (Union[cirq.PauliSum, cirq.PauliString, list]):
+            Observable(s) to measure.
+        repetitions (Optional[int]):
+            Number of measurement repetitions. If None, analytic
+            expectation is used.
+        backend (Union[str, cirq.Sampler,
+            cirq.sim.simulator.SimulatesExpectationValues]):
+            Backend simulator.
+        differentiator (Optional[tfq.differentiators.Differentiator]):
+            Differentiation scheme.
+        initializer (tf.keras.initializers.Initializer):
+            Initializer for circuit parameters.
+        regularizer (Optional[tf.keras.regularizers.Regularizer]):
+            Regularizer for circuit parameters.
+        constraint (Optional[tf.keras.constraints.Constraint]):
+            Constraint for circuit parameters.
+        **kwargs: Additional keyword arguments for the parent class.
+
+    Input shape:
+        tf.Tensor of shape [batch_size], each entry a serialized circuit
+        (from tfq.convert_to_tensor).
+
+    Output shape:
+        tf.Tensor of shape [batch_size, n_operators], expectation values
+        for each operator.
+
     This layer is for training parameterized quantum models.
     Given a parameterized circuit, this layer initializes the parameters
     and manages them in a Keras native way.
@@ -293,7 +322,16 @@ def build(self, input_shape):
         super().build(input_shape)
 
     def call(self, inputs):
-        """Keras call function."""
+        """Keras call function.
+
+        Args:
+            inputs (tf.Tensor): Tensor of shape [batch_size], each entry a
+            serialized circuit (from tfq.convert_to_tensor).
+
+        Returns:
+            tf.Tensor: Tensor of shape [batch_size, n_operators],
+            expectation values for each operator.
+        """
         circuit_batch_dim = tf.gather(tf.shape(inputs), 0)
         tiled_up_model = tf.tile(self._model_circuit, [circuit_batch_dim])
         model_appended = self._append_layer(inputs, append=tiled_up_model)