Function do_test_softmax

python/oneflow/test/modules/test_activation.py:456–470

(batch_size: int, log_softmax: bool = False)

def do_test_softmax(batch_size: int, log_softmax: bool = False):
    num_dims = random(low=1, high=5).to(int)
    m = torch.nn.Softmax(dim=random(low=0, high=num_dims).to(int) | nothing())
    if log_softmax:
        m = torch.nn.LogSoftmax(dim=random(low=0, high=num_dims).to(int) | nothing())
    m.train(random())
    device = random_device()
    m.to(device)
    x = (
        random_tensor(ndim=num_dims).to(device)
        if batch_size < 0
        else random_tensor(ndim=num_dims, dim0=batch_size).to(device)
    )
    y = m(x)
    return y


@flow.unittest.skip_unless_1n1d()