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

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()])

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?