How to take gradient of the `tfq.layers.State` output?

Question

I am using the following code for building a quantum circuit as a custom tf.keras.layers.Layer:

import tensorflow as tf
import tensorflow_quantum as tfq
import numpy as np
import sympy
import cirq


class QuantumLayer(tf.keras.layers.Layer):

  def __init__(self) -> None:
    super(QuantumLayer, self).__init__()
    self.qubits = [cirq.GridQubit(1, 0), cirq.GridQubit(1, 1)]
    self.num_params = 2
    self.params = sympy.symbols("params0:%d"%self.num_params)
    self.theta = tf.Variable(initial_value=np.random.uniform(0, 2*np.pi, (1, self.num_params)), dtype="float32", trainable=True)
    self.operation = tfq.layers.State()
  
  def quantum_circ(self, param):
    c = cirq.Circuit()
    for i in range(len(self.qubits)):
      c += cirq.ry(param[i]).on(self.qubits[i])
    return c

  def __call__(self, inputs):
    res = self.operation(self.quantum_circ(self.params), symbol_names=self.params, 
                         symbol_values=self.theta)
    out = tf.squeeze(tf.abs(res.to_tensor() ** 2))

    return out

layer = QuantumLayer()
inputs = tfq.convert_to_tensor([cirq.Circuit()])
with tf.GradientTape() as tape:
  result = layer(inputs)
grad = tape.gradient(result[1], layer.trainable_variables)
print(grad)


>>> [None]

I think this is mostly because of the fact that the output res in the class QuantumLayer is a tf.RaggedTensor which cannot be differentiated (Also, mentioned here that the output quantum state is not differentiable, but here I am using the probability of the state of the quantum state to get the output (as shown in the out of the class). Where am I wrong in this case to get the gradient right and how can I achieve this?