def step(self, u, state):
        """Must have called self._setup_step() and created state with self.default_state() before calling this."""

        if self._step_mode == 'linear':
            new_state = self._step_state_linear(u, state)
        else:
            new_state = self._step_state(u, state)
        y = torch.einsum(self.output_contraction, self.dC, new_state)
        return y.real, new_state
Must have called self._setup_step() and created state with self.default_state() before calling this.
analyze/convnexts4nd/src/models/sequence/kernels/ssm.py