Question about Building a Custom Op for Beam Search

[StarWang]

Hi Team,

I'm trying to build a custom op for beam search and have a question on specifying the input type. My question is: can we take RNNCell as input? If so, can you share a reference or example on how to do this?

Thanks!