Added support for output_dim to be a tuple

[kessler-frost]

Context:
Until now, only int and flat tuples like (5,) were acceptable values by the output_dim parameter in qnn.KerasLayer. After the addition of the density_matrix() return type by a QNode in PR #878 the output from a QNode may not necessarily be flat. This PR adds the support for tuples to be passed as an acceptable value to the output_dim parameter.

Description of the Change:
Now values like (5, 5) as well as integers like 2 or flat tuples like (5,) can be passed to the output_dim parameter. It is also to be noted that if an Iterable is passed (e.g, a list like [5, 5]) instead, it will be converted to a tuple. Internally the value of output_dim will either be a tuple or an integer.

Example:

Here, c is a circuit that outputs a density_matrix() qnode with the shape output_dim(a square tuple), and w is the shape of weights to be used in the circuit. We need to flatten the outputs if we are expecting the following layer to receive inputs with only 1 axis (excluding batch size as it remains unchanged when flattening). We also add a lambda layer so that only real values in the density matrix are used for processing (read the discussion below to know more).

qml.enable_tape()

custom_layer = KerasLayer(c, w, output_dim)

model = tf.keras.models.Sequential(
        [
            tf.keras.layers.Dense(n_qubits),
            layer1,
            # Adding a lambda layer to take only the real values from density matrix
            tf.keras.layers.Lambda(lambda x: tf.abs(x)),
            tf.keras.layers.Flatten(),
            tf.keras.layers.Dense(n_qubits),
            layer2,
            # Adding a lambda layer to take only the real values from density matrix
            tf.keras.layers.Lambda(lambda x: tf.abs(x)),
            tf.keras.layers.Flatten(),
            tf.keras.layers.Dense(output_dim[0] * output_dim[1])
        ]
    )

Benefits:

• Tuples are now supported as acceptable values by output_dim without breaking the current usage patterns.
• This may also enable the use of other Keras Layers like Conv2D with our KerasLayer in the future.

Related GitHub Issues:
Issue #936

[kessler-frost]

The two tests that are failing are because they indirectly make sure that the output_dim is an integer value, and since this improvement enables the use of tuples in output_dim, I am unsure of what should be done next.

[josh146]

Thanks @kessler-frost for the PR! Let us know when it's ready for final code review, and someone from the team will be able to have a look :)

The two tests that are failing are because they indirectly make sure that the output_dim is an integer value, and since this improvement enables the use of tuples in output_dim, I am unsure of what should be done next.

If the logical changes cause assumptions made in previous tests to fail, then it is fine to update the tests to match the newer logic.

[codecov]

Codecov Report

Merging #1070 (554b887) into master (1206dbd) will increase coverage by 0.00%.
The diff coverage is 100.00%.

@@           Coverage Diff           @@
##           master    #1070   +/-   ##
=======================================
  Coverage   97.74%   97.74%           
=======================================
  Files         154      154           
  Lines       11726    11728    +2     
=======================================
+ Hits        11461    11463    +2     
  Misses        265      265           

Impacted Files	Coverage Δ
pennylane/qnn/keras.py	97.95% <100.00%> (+0.04%)	⬆️

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 1206dbd...554b887. Read the comment docs.

[kessler-frost]

I was writing the test test_train_model_for_density_matrix() needed to verify if a model can be trained using the density_matrix() type QNode but am unable to resolve the following input incompatibility error:

self = <test_keras.TestKerasLayerIntegration object at 0x7f22d44e00d0>, model = <tensorflow.python.keras.engine.sequential.Sequential object at 0x7f22f83bb1c0>, batch_size = 2
n_qubits = 4, output_dim = (4, 4)

    @pytest.mark.parametrize("n_qubits, output_dim", indices_up_to(3, has_tuple=True))
    @pytest.mark.parametrize("batch_size", [2])
    def test_train_model_for_density_matrix(self, model, batch_size, n_qubits, output_dim):
        """Test if a model can train using the KerasLayer when QNode returns a density_matrix().
        The model is composed of two KerasLayers sandwiched between Dense neural network layers,
        and the dataset is simply input and output vectors of zeros."""

        if not qml.tape_mode_active():
            pytest.skip("This functionality is only supported in tape mode.")

        x = np.zeros((batch_size, n_qubits))
        y = np.zeros((batch_size, output_dim[0] * output_dim[1]))

        model.compile(optimizer="sgd", loss="mse")

>       model.fit(x, y, batch_size=batch_size, verbose=0)

....

ValueError: Input 0 of layer dense_52 is incompatible with the layer: expected axis -1 of input shape to have value 16 but received input with shape (2, 4)

I've been trying to figure out the output shape of our KerasLayer if a tuple is passed in output_dim for a while now, so any help would be appreciated. Thanks!

Edit: Turns out I was sampling output_dim incorrectly and not in accordance with the circuit used for testing.

[kessler-frost]

As discussed in the related issue, if the layer next to our KerasLayer is a Dense layer, it will be expecting a single-dimension tensor hence it is necessary to use the Flatten layer before passing it to the Dense layer when creating a model with a circuit returning density_matrix() output type. This update to output_dim may also allow the support of Keras' Conv2D layer.

[kessler-frost]

@trbromley The PR is now ready for review.

[josh146]

The PR is now ready for review.

Awesome, thanks @kessler-frost 💪

Someone from the team will be able to have a look shortly.

[trbromley]

Thanks @kessler-frost! Looks great! I have left some comments and can give the ✔️ once resolved.

One thing to bear in mind is I suggested we add a note in the docstring to warn users about casting - returning the state or dm is complex and we need to be careful. Also, we only support differentiating the state or dm if the device is default.qubit.tf using diff_method=backprop, but this was recently made the default so we probabaly don't need to add a comment for that.

[trbromley]

Thanks @kessler-frost! Apologies for being a bit slow to get back to this, we've been having fun with QHack this week and last.

This looks good and I approve! Although I still think we should consider adding the comment about complex/real dtypes, but would be interested to get your thoughts too.

Finally, I saw this commit: 1a2354c for updating the changelog, but couldn't see the changelog in the latest updates. Please add the change as well as your name 😄 Thanks!!

[trbromley]

Does this work without
@pytest.mark.usefixtures("get_circuit_dm")
?

(we have @pytest.mark.usefixtures("get_circuit") for model())

[kessler-frost]

Yes, apparently it is mentioned in pytest's documentation that fixtures cannot use other fixtures in this way. I suggest we remove the @pytest.mark.usefixtures("get_circuit") for model() too.

[kessler-frost]

AFAIK the fixture get_circuit or get_circuit_dm is getting used when we call the @pytest.mark.usefixture decorator on TestKerasLayerIntegration and TestKerasLayerIntegrationDM classes respectively. Thus, we do not need to use it for model() and model_dm().

[kessler-frost]

Thanks for the review @trbromley! I have added the cautionary comment and updated the changelog as well. The PR may finally be ready for merging 😅.

[trbromley]

Thanks @kessler-frost for the contribution! Merging in shortly!! 💯