def _step_state(self, u, state):
        """Must be called after self.default_state() is used to construct an initial state!"""
        next_state = (torch.einsum(self.state_contraction, self.dA, state)
                     + torch.einsum(self.input_contraction, self.dB, u))
        return next_state
Must be called after self.default_state() is used to construct an initial state!
analyze/convnexts4nd/src/models/sequence/kernels/ssm.py