feat(tensorflow): Choose NumPy API in forward pass

`tf.function` requires Piquasso to use Tensorflow's NumPy API in the forward
pass. By default, we use regular NumPy for performance reasons, but this patch
enables users to run Piquasso inside `tf.function`, which might be faster with
JIT compilation enabled.

Unfortunately, the JIT compilation is still very slow, and we may need to
reduce the number of Python control flows during the calculation.

benchmarks/tf_cvnn_benchmark.py

@@ -106,6 +106,7 @@ def _pq_state_vector(weights, cutoff):
     return state.get_tensor_representation()
 
 
+@tf.function(jit_compile=True)
 def _calculate_piquasso_results(weights, cutoff):
     with tf.GradientTape() as tape:
         state_vector = _pq_state_vector(weights, cutoff)

piquasso/_backends/calculator.py

@@ -29,6 +29,14 @@ class _BuiltinCalculator(BaseCalculator):
     """Base class for built-in calculators."""
 
 
+def _noop_custom_gradient(func):
+    def noop_grad(*args, **kwargs):
+        result, _ = func(*args, **kwargs)
+        return result
+
+    return noop_grad
+
+
 class NumpyCalculator(_BuiltinCalculator):
     """The calculations for a simulation using NumPy (and SciPy).
 
@@ -38,6 +46,7 @@ class NumpyCalculator(_BuiltinCalculator):
     def __init__(self):
         self.np = np
         self.fallback_np = np
+        self.forward_pass_np = np
         self.block_diag = scipy.linalg.block_diag
         self.block = np.block
         self.logm = scipy.linalg.logm
@@ -51,7 +60,9 @@ def __init__(self):
         self.hafnian = hafnian_with_reduction
         self.loop_hafnian = loop_hafnian_with_reduction
 
-    def maybe_convert_to_numpy(self, value):
+        self.custom_gradient = _noop_custom_gradient
+
+    def preprocess_input_for_custom_gradient(self, value):
         return value
 
     def assign(self, array, index, value):
@@ -112,7 +123,19 @@ class TensorflowCalculator(_BuiltinCalculator):
         gradient = tape.gradient(mean, [r])
     """
 
-    def __init__(self):
+    def __init__(self, no_custom_gradient=False):
+        """
+        Args:
+            no_custom_gradient (bool, optional): Executes the forward pass
+                using Tensorflow, instead of NumPy. This is sometimes necessary, e.g.,
+                when one wants to decorate Piquasso code with
+                `tf.function(jit_compile=True)`. In fact, when eager execution is
+                disabled in Tensorflow, this argument is automatically set to `True`.
+                Otherwise, it defaults to False.
+
+        Raises:
+            ImportError: When TensorFlow is not available.
+        """
         try:
             import tensorflow as tf
         except ImportError:
@@ -131,9 +154,20 @@ def __init__(self):
         self._tf = tf
         self.np = tnp
         self.fallback_np = np
+
+        no_custom_gradient = no_custom_gradient or not tf.executing_eagerly()
+        self.no_custom_gradient = no_custom_gradient
+        self.forward_pass_np = tnp if no_custom_gradient else np
+        self.custom_gradient = (
+            _noop_custom_gradient if no_custom_gradient else tf.custom_gradient
+        )
+
         self.sqrtm = tf.linalg.sqrtm
 
-    def maybe_convert_to_numpy(self, value):
+    def preprocess_input_for_custom_gradient(self, value):
+        if self.no_custom_gradient:
+            return value
+
         return value.numpy() if self._tf.is_tensor(value) else value
 
     def block_diag(self, *arrs):
@@ -147,9 +181,34 @@ def permanent(self, matrix, rows, columns):
         return glynn_gray_permanent(matrix, rows, columns, np=self.np)
 
     def assign(self, array, index, value):
-        # NOTE: This is not as advanced as Numpy's indexing, only supports 1D arrays.
+        """
+        NOTE: This method is very limited, and is a bit hacky, since TF does not support
+        item assignment through its NumPy API.
+        """
 
-        return self._tf.tensor_scatter_nd_update(array, [[index]], [value])
+        if isinstance(array, self.fallback_np.ndarray):
+            array[index] = value
+
+            return array
+
+        if isinstance(index, int):
+            return self._tf.tensor_scatter_nd_update(array, [[index]], [value])
+
+        # NOTE: When using `tf.function`, TensorFlow threw the following error:
+        #
+        # TypeError: Tensors in list passed to 'values' of 'ConcatV2' Op have types [int32, int64] that don't all match.  # noqa: E501
+        #
+        # To make it disappear, I had to convert all the indices to `int32`.
+        index = index.astype(self.fallback_np.int32)
+
+        if not isinstance(index, np.ndarray) or not len(index.shape) == 2:
+            raise RuntimeError(f"Invalid index: {index}")
+
+        array = self._tf.tensor_scatter_nd_update(
+            array, index.reshape(-1, 1), value.reshape(-1)
+        )
+
+        return array
 
     def block(self, arrays):
         # NOTE: This is not as advanced as `numpy.block`, this function only supports

piquasso/_backends/fock/calculations.py

@@ -203,7 +203,7 @@ def calculate_interferometer_helper_indices(space):
     }
 
 
-def calculate_interferometer_on_fock_space(interferometer, index_dict):
+def calculate_interferometer_on_fock_space(interferometer, index_dict, calculator):
     """Calculates finite representation of interferometer in the Fock space.
     The function assumes the knowledge of the 1-particle unitary.
 
@@ -219,6 +219,8 @@ def calculate_interferometer_on_fock_space(interferometer, index_dict):
         numpy.ndarray: Finite representation of interferometer in the Fock space
     """
 
+    np = calculator.forward_pass_np
+
     subspace_representations = []
 
     subspace_representations.append(np.array([[1.0]], dtype=interferometer.dtype))

piquasso/_backends/fock/general/calculations.py

@@ -57,7 +57,7 @@ def _apply_passive_linear(state, interferometer, modes):
     subspace = state._get_subspace(dim=len(interferometer))
 
     subspace_transformations = _get_interferometer_on_fock_space(
-        interferometer, subspace
+        interferometer, subspace, state._calculator
     )
 
     _apply_passive_gate_matrix_to_state(state, subspace_transformations, modes)
@@ -107,10 +107,12 @@ def _calculate_density_matrix_after_interferometer(
     return new_density_matrix
 
 
-def _get_interferometer_on_fock_space(interferometer, space):
+def _get_interferometer_on_fock_space(interferometer, space, calculator):
     index_dict = calculate_interferometer_helper_indices(space)
 
-    return calculate_interferometer_on_fock_space(interferometer, index_dict)
+    return calculate_interferometer_on_fock_space(
+        interferometer, index_dict, calculator
+    )
 
 
 def particle_number_measurement(

piquasso/_backends/fock/pure/calculations/__init__.py

@@ -118,14 +118,14 @@ def _apply_active_gate_matrix_to_state(
 
     @calculator.custom_gradient
     def _apply_active_gate_matrix(state_vector, matrix):
-        state_vector = calculator.maybe_convert_to_numpy(state_vector)
-        matrix = calculator.maybe_convert_to_numpy(matrix)
+        state_vector = calculator.preprocess_input_for_custom_gradient(state_vector)
+        matrix = calculator.preprocess_input_for_custom_gradient(matrix)
 
         state_index_matrix_list = calculate_state_index_matrix_list(
             space, auxiliary_subspace, mode
         )
         new_state_vector = _calculate_state_vector_after_apply_active_gate(
-            state_vector, matrix, state_index_matrix_list
+            state_vector, matrix, state_index_matrix_list, calculator
         )
         grad = _create_linear_active_gate_gradient_function(
             state_vector, matrix, state_index_matrix_list, calculator
@@ -136,8 +136,13 @@ def _apply_active_gate_matrix(state_vector, matrix):
 
 
 def _calculate_state_vector_after_apply_active_gate(
-    state_vector, matrix, state_index_matrix_list
+    state_vector,
+    matrix,
+    state_index_matrix_list,
+    calculator,
 ):
+    np = calculator.forward_pass_np
+
     new_state_vector = np.empty_like(state_vector, dtype=state_vector.dtype)
 
     is_batch = len(state_vector.shape) == 2
@@ -146,8 +151,12 @@ def _calculate_state_vector_after_apply_active_gate(
 
     for state_index_matrix in state_index_matrix_list:
         limit = state_index_matrix.shape[0]
-        new_state_vector[state_index_matrix] = np.einsum(
-            einsum_string, matrix[:limit, :limit], state_vector[state_index_matrix]
+        new_state_vector = calculator.assign(
+            new_state_vector,
+            state_index_matrix,
+            np.einsum(
+                einsum_string, matrix[:limit, :limit], state_vector[state_index_matrix]
+            ),
         )
 
     return new_state_vector

piquasso/_backends/fock/pure/calculations/passive_linear.py

@@ -57,11 +57,11 @@ def _apply_passive_linear(state, interferometer, modes, calculator):
 
 def _get_interferometer_on_fock_space(interferometer, space, calculator):
     def _get_interferometer_with_gradient_callback(interferometer):
-        interferometer = calculator.maybe_convert_to_numpy(interferometer)
+        interferometer = calculator.preprocess_input_for_custom_gradient(interferometer)
         index_dict = calculate_interferometer_helper_indices(space)
 
         subspace_representations = calculate_interferometer_on_fock_space(
-            interferometer, index_dict
+            interferometer, index_dict, calculator
         )
         grad = _calculate_interferometer_gradient_on_fock_space(
             interferometer, calculator, subspace_representations, index_dict
@@ -201,10 +201,10 @@ def _apply_passive_gate_matrix_to_state(
     space = state._space
 
     def _apply_interferometer_matrix(state_vector, subspace_transformations):
-        state_vector = calculator.maybe_convert_to_numpy(state_vector)
+        state_vector = calculator.preprocess_input_for_custom_gradient(state_vector)
 
         subspace_transformations = [
-            calculator.maybe_convert_to_numpy(matrix)
+            calculator.preprocess_input_for_custom_gradient(matrix)
             for matrix in subspace_transformations
         ]
 
@@ -217,6 +217,7 @@ def _apply_interferometer_matrix(state_vector, subspace_transformations):
             state_vector,
             subspace_transformations,
             index_list,
+            calculator,
         )
 
         grad = _create_linear_passive_gate_gradient_function(
@@ -236,16 +237,23 @@ def _calculate_state_vector_after_interferometer(
     state_vector: np.ndarray,
     subspace_transformations: List[np.ndarray],
     index_list: List[np.ndarray],
+    calculator: BaseCalculator,
 ) -> np.ndarray:
+    np = calculator.forward_pass_np
+
     new_state_vector = np.empty_like(state_vector)
 
     is_batch = len(state_vector.shape) == 2
 
     einsum_string = "ij,jkl->ikl" if is_batch else "ij,jk->ik"
 
     for n, indices in enumerate(index_list):
-        new_state_vector[indices] = np.einsum(
-            einsum_string, subspace_transformations[n], state_vector[indices]
+        new_state_vector = calculator.assign(
+            new_state_vector,
+            indices,
+            np.einsum(
+                einsum_string, subspace_transformations[n], state_vector[indices]
+            ),
         )
 
     return new_state_vector

piquasso/_backends/fock/pure/state.py

@@ -205,8 +205,8 @@ def mean_position(self, mode: int) -> np.ndarray:
         state_vector = self._state_vector
 
         accumulator = np.dot(
-            (multipliers * np.take(state_vector, left_indices)),
-            np.take(state_vector, right_indices),
+            (multipliers * np.take(state_vector, np.array(left_indices))),
+            np.take(state_vector, np.array(right_indices)),
         )
 
         return np.real(accumulator) * fallback_np.sqrt(self._config.hbar / 2)

piquasso/_math/fock.py

@@ -168,11 +168,15 @@ def get_single_mode_squeezing_operator(
     ) -> np.ndarray:
         @self.calculator.custom_gradient
         def _single_mode_squeezing_operator(r, phi):
-            r = self.calculator.maybe_convert_to_numpy(r)
-            phi = self.calculator.maybe_convert_to_numpy(phi)
+            r = self.calculator.preprocess_input_for_custom_gradient(r)
+            phi = self.calculator.preprocess_input_for_custom_gradient(phi)
 
             matrix = create_single_mode_squeezing_matrix(
-                r, phi, self.cutoff, complex_dtype=self.config.complex_dtype
+                r,
+                phi,
+                self.cutoff,
+                complex_dtype=self.config.complex_dtype,
+                calculator=self._calculator,
             )
             grad = create_single_mode_squeezing_gradient(
                 r,
@@ -300,11 +304,15 @@ def get_annihilation_operator(self, modes: Tuple[int, ...]) -> np.ndarray:
     def get_single_mode_displacement_operator(self, *, r, phi):
         @self.calculator.custom_gradient
         def _single_mode_displacement_operator(r, phi):
-            r = self.calculator.maybe_convert_to_numpy(r)
-            phi = self.calculator.maybe_convert_to_numpy(phi)
+            r = self.calculator.preprocess_input_for_custom_gradient(r)
+            phi = self.calculator.preprocess_input_for_custom_gradient(phi)
 
             matrix = create_single_mode_displacement_matrix(
-                r, phi, self.cutoff, complex_dtype=self.config.complex_dtype
+                r,
+                phi,
+                self.cutoff,
+                complex_dtype=self.config.complex_dtype,
+                calculator=self._calculator,
             )
             grad = create_single_mode_displacement_gradient(
                 r,