Function _inputs

tests/test_stateless_executor.py:80–105

()

def _inputs() -> StatelessForwardInputs:
    input_ids = torch.tensor(
        [
            [0, 1, 2, 3, 0],
            [0, 2, 1, 4, 5],
        ],
        dtype=torch.long,
    )
    attention_mask = torch.tensor(
        [
            [1, 1, 1, 1, 0],
            [1, 1, 1, 1, 1],
        ],
        dtype=torch.bool,
    )
    completion_mask = torch.tensor(
        [
            [False, False, True, True, False],
            [False, False, True, True, True],
        ]
    )
    return StatelessForwardInputs(
        input_ids=input_ids,
        attention_mask=attention_mask,
        completion_mask=completion_mask,
    )


def _logits_for(inputs: StatelessForwardInputs, vocab_size: int = 8) -> torch.Tensor: