def METHOD_NAME(test_case, device):
    input = np.random.randn(3, 4, 3, 5)
    index = np.random.choice(np.arange(3), size=180, replace=True).reshape((3, 4, 3, 5))
    np_out = np.take_along_axis(input, index, 1)
    output = flow.gather(
        flow.tensor(input, dtype=flow.float32, device=flow.device(device)),
        1,
        flow.tensor(index, dtype=flow.int64, device=flow.device(device)),
    )
    test_case.assertTrue(np.allclose(output.numpy(), np_out))
    np_out2 = np.take_along_axis(input, index, 2)
    output2 = flow.gather(
        flow.tensor(input, dtype=flow.float32, device=flow.device(device)),
        2,
        flow.tensor(index, dtype=flow.int64, device=flow.device(device)),
    )
    test_case.assertTrue(np.allclose(output2.numpy(), np_out2))
    np_out3 = np.take_along_axis(input, index, 3)
    output3 = flow.gather(
        flow.tensor(input, dtype=flow.float32, device=flow.device(device)),
        3,
        flow.tensor(index, dtype=flow.int64, device=flow.device(device)),
    )
    test_case.assertTrue(np.allclose(output3.numpy(), np_out3))