feat(tensorflow): Choose NumPy API in forward pass

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,57 @@ 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.
+        """
+
+        if isinstance(array, self.fallback_np.ndarray):
+            array[index] = value
 
-        return self._tf.tensor_scatter_nd_update(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 len(array.shape) == 1:
+            return self._tf.tensor_scatter_nd_update(
+                array, index.reshape(-1, 1), value.reshape(-1)
+            )
+
+        number_of_batches = array.shape[1]
+        int_dtype = index.dtype
+
+        flattened_index = index.reshape(-1)
+
+        indices = self.fallback_np.column_stack(
+            [
+                self.fallback_np.tile(flattened_index, number_of_batches),
+                self.fallback_np.concatenate(
+                    [
+                        self.fallback_np.full(len(flattened_index), i, dtype=int_dtype)
+                        for i in range(number_of_batches)
+                    ]
+                ),
+            ]
+        )
+
+        values = self.np.concatenate(
+            [value[:, :, i].reshape(-1) for i in range(number_of_batches)]
+        )
+
+        return self._tf.tensor_scatter_nd_update(
+            array,
+            indices,
+            values,
+        )
 
     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,