def get_initial_state(self, batch_size: int):
        """
        Get zeros states

        Returns
        -------
        tf.Tensor, shape [B, num_rnns, nstates, state_size]
            Zero initialized states
        """
        states = []
        for rnn in self.rnns:
            states.append(tf.stack(rnn.get_initial_state(batch_size=batch_size), axis=0))
        return tf.transpose(tf.stack(states, axis=0), perm=[2, 0, 1, 3])

        Get zeros states

        Returns
        -------
        tf.Tensor, shape [B, num_rnns, nstates, state_size]
            Zero initialized states
        
tensorflow_asr/models/transducer/base_transducer.py