Repository: XanaduAI/strawberryfields

Issue:
quantum neural net example output is Nan when using more than 1 mode
#### Issue description

The quantum_neural_network.py example output is Nan when using more than 1 mode. (working as expected for 1 mode)

* *Expected behavior:* Not Nan

* *Actual behavior:* ```Beginning optimization
Rep: 0 Cost: 7.0005 Fidelity: 0.0000 Trace: 1.0000
Rep: 1 Cost: nan Fidelity: nan Trace: nan
Rep: 2 Cost: nan Fidelity: nan Trace: nan
Rep: 3 Cost: nan Fidelity: nan Trace: nan```

* *Reproduces how often:* 100%
* *System information:* 
``` Strawberry Fields: a Python library for continuous-variable quantum circuits.
Copyright 2018-2020 Xanadu Quantum Technologies Inc.

Python version:            3.8.5
Platform info:             Linux-5.8.0-53-generic-x86_64-with-glibc2.10
Installation path:         /home/jonas/anaconda3/envs/strawberry/lib/python3.8/site-packages/strawberryfields
Strawberry Fields version: 0.17.0
Numpy version:             1.19.2
Scipy version:             1.4.1
SymPy version:             1.7.1
NetworkX version:          2.5
The Walrus version:        0.14.0
Blackbird version:         0.3.1-dev
TensorFlow version:        2.2.0
```


#### Source code and tracebacks
update line 135 of the example quantum_neural_network.py from ```modes = 1``` to ```modes = 2``` 




Discussion:
Hi @jonasbrami, thank you for catching this! I could reproduce the issue based on your suggestion and it seems that the issue is with computing the gradient of the
```python
difference = tf.reduce_sum(tf.abs(ket - target_state))
```
line in the cost function. Changing the absolute value computation such that a relatively small constant is added and having `tf.abs(ket - target_state + 1e-10)` makes the resulting `nan` values disappear.

We're considering a fitting fix here.

Resolving patch:
diff --git a/examples/quantum_neural_network.py b/examples/quantum_neural_network.py
--- a/examples/quantum_neural_network.py
+++ b/examples/quantum_neural_network.py
@@ -25,9 +25,9 @@ def interferometer(params, q):
             is to be applied to
     """
     N = len(q)
-    theta = params[:N*(N-1)//2]
-    phi = params[N*(N-1)//2:N*(N-1)]
-    rphi = params[-N+1:]
+    theta = params[: N * (N - 1) // 2]
+    phi = params[N * (N - 1) // 2 : N * (N - 1)]
+    rphi = params[-N + 1 :]
 
     if N == 1:
         # the interferometer is a single rotation
@@ -65,11 +65,11 @@ def layer(params, q):
     M = int(N * (N - 1)) + max(1, N - 1)
 
     int1 = params[:M]
-    s = params[M:M+N]
-    int2 = params[M+N:2*M+N]
-    dr = params[2*M+N:2*M+2*N]
-    dp = params[2*M+2*N:2*M+3*N]
-    k = params[2*M+3*N:2*M+4*N]
+    s = params[M : M + N]
+    int2 = params[M + N : 2 * M + N]
+    dr = params[2 * M + N : 2 * M + 2 * N]
+    dp = params[2 * M + 2 * N : 2 * M + 3 * N]
+    k = params[2 * M + 3 * N : 2 * M + 4 * N]
 
     # begin layer
     interferometer(int1, q)
@@ -115,7 +115,9 @@ def init_weights(modes, layers, active_sd=0.0001, passive_sd=0.1):
     dp_weights = tf.random.normal(shape=[layers, modes], stddev=passive_sd)
     k_weights = tf.random.normal(shape=[layers, modes], stddev=active_sd)
 
-    weights = tf.concat([int1_weights, s_weights, int2_weights, dr_weights, dp_weights, k_weights], axis=1)
+    weights = tf.concat(
+        [int1_weights, s_weights, int2_weights, dr_weights, dp_weights, k_weights], axis=1
+    )
     weights = tf.Variable(weights)
 
     return weights
@@ -166,6 +168,20 @@ def init_weights(modes, layers, active_sd=0.0001, passive_sd=0.1):
         layer(sf_params[k], q)
 
 
+def safe_abs(x):
+    # Helper function to deal with tensor terms near zero
+
+    # Check where we have near zero terms
+    EPS = 1e-15
+    x = tf.where(tf.abs(x) < EPS, tf.zeros_like(x), x)
+    zero = tf.constant(0, dtype=tf.complex64)
+    x_ok = tf.not_equal(x, zero)
+
+    # To make sure, swap out the zeros with ones
+    safe_x = tf.where(x_ok, x, tf.ones_like(x, dtype=tf.complex64))
+    return tf.where(x_ok, tf.abs(safe_x), tf.zeros_like(x, dtype=tf.float32))
+
+
 def cost(weights):
     # Create a dictionary mapping from the names of the Strawberry Fields
     # symbolic gate parameters to the TensorFlow weight values.
@@ -175,7 +191,7 @@ def cost(weights):
     state = eng.run(qnn, args=mapping).state
     ket = state.ket()
 
-    difference = tf.reduce_sum(tf.abs(ket - target_state))
+    difference = tf.reduce_sum(safe_abs(ket - target_state))
     fidelity = tf.abs(tf.reduce_sum(tf.math.conj(ket) * target_state)) ** 2
     return difference, fidelity, ket, tf.math.real(state.trace())
 


Test patch: