def _zero_state(self, input: torch.Tensor) -> torch.Tensor:
        """Returns a zeros state matching dtype and batch size of `input`.

        Arguments:
          input: Tensor, to specify the device and dtype of the returned tensors.

        Returns:
          zero_state: a nested structure of zero Tensors.
        """
        batch_dim = input.shape[1]
        state = torch.zeros(
            (self.config.num_states, batch_dim, self.config.hidden_size),
            dtype=input.dtype,
            device=input.device,
        )
        return state
Returns a zeros state matching dtype and batch size of `input`.

        Arguments:
          input: Tensor, to specify the device and dtype of the returned tensors.

        Returns:
          zero_state: a nested structure of zero Tensors.
        
xlstm/blocks/slstm/cell.py