def test_input_shape_and_dtype(
    gru_model: GRUModel,
    batch_prev_tkids: torch.Tensor,
):
    r"""Input must be long tensor."""

    try:
        gru_model = gru_model.eval()
        gru_model.pred(batch_prev_tkids)
    except Exception:
        assert False
def test_value_range(
    gru_model: GRUModel,
    batch_prev_tkids: torch.Tensor,
):
    r"""Return values are probabilities."""
    gru_model = gru_model.eval()
    out = gru_model.pred(batch_prev_tkids)

    # Probabilities are values within range [0, 1].
    assert torch.all(0 <= out).item()
    assert torch.all(out <= 1).item()

    # Sum of the probabilities equals to 1.
    accum_out = out.sum(dim=-1)
    assert torch.allclose(accum_out, torch.ones_like(accum_out))
def test_return_shape_and_dtype(
    gru_model: GRUModel,
    batch_prev_tkids: torch.Tensor,
):
    r"""Return float tensor with correct shape."""
    gru_model = gru_model.eval()
    out = gru_model.pred(batch_prev_tkids)

    # Output float tensor.
    assert out.dtype == torch.float

    # Input shape: (B, S).
    # Output shape: (B, S, V).
    assert out.shape == (
        batch_prev_tkids.shape[0],
        batch_prev_tkids.shape[1],
        gru_model.emb.num_embeddings,
    )