| """Pack OLMo's document offsets for FLA; no global forward patch."""
|
|
|
| import torch
|
|
|
|
|
| def flatten_documents(x: torch.Tensor, cu_doc_lens: torch.Tensor | None) -> torch.Tensor:
|
| """Flatten before BOTH short convolutions and the recurrent kernel.
|
|
|
| OLMo's offsets already refer to the row-major flattened batch. Returning a
|
| view preserves autograd and lets the caller restore the original shape.
|
| """
|
| if x.ndim != 3:
|
| raise ValueError("Expected x with shape [batch, sequence, hidden]")
|
| if cu_doc_lens is None:
|
| return x
|
| if cu_doc_lens.ndim != 1 or cu_doc_lens.numel() < 2:
|
| raise ValueError("cu_doc_lens must be a one-dimensional offsets tensor")
|
| if cu_doc_lens.dtype not in (torch.int32, torch.int64):
|
| raise ValueError("cu_doc_lens must have integer dtype int32 or int64")
|
| if cu_doc_lens.device != x.device:
|
| raise ValueError("cu_doc_lens and x must be on the same device")
|
| batch, length, hidden = x.shape
|
|
|
| if not torch.compiler.is_compiling():
|
| if int(cu_doc_lens[0]) != 0 or int(cu_doc_lens[-1]) != batch * length:
|
| raise ValueError("cu_doc_lens must start at 0 and end at batch * sequence")
|
| if not bool((cu_doc_lens[1:] > cu_doc_lens[:-1]).all()):
|
| raise ValueError("cu_doc_lens must increase strictly (no empty documents)")
|
| return x.reshape(1, batch * length, hidden)
|
|
|
