def test_to_numpy(np_torch_tuple, dim):
    """ Test torch --> np conversion (right angles)"""
    from_np, from_torch = np_torch_tuple
    if dim == 0:
        np_array = np.array(from_np)
        torch_tensor = torch.tensor(from_torch)
    elif dim == 1:
        np_array = np.array([from_np])
        torch_tensor = torch.tensor([from_torch])
    elif dim == 2:
        np_array = np.array([[from_np]])
        torch_tensor = torch.tensor([[from_torch]])
    else:
        return
    np_from_torch = transforms.to_numpy(torch_tensor, dim=dim)
    np.testing.assert_allclose(np_array, np_from_torch) 

def test_from_numpy(np_torch_tuple, dim):
    """ Test np --> torch conversion (right angles)"""
    from_np, from_torch = np_torch_tuple
    if dim == 0:
        np_array = np.array(from_np)
        torch_tensor = torch.tensor(from_torch)
    elif dim == 1:
        np_array = np.array([from_np])
        torch_tensor = torch.tensor([from_torch])
    elif dim == 2:
        np_array = np.array([[from_np]])
        torch_tensor = torch.tensor([[from_torch]])
    else:
        return
    torch_from_np = transforms.from_numpy(np_array, dim=dim)
    np.testing.assert_allclose(torch_tensor, torch_from_np) 

def test_last_pooling_with_mask():
    lp = pooling.LastPooling()
    out = lp(TENSOR, padding_mask=MASK)

    assert out.size() == torch.Size([2, 4])

    expected = Tensor(
        [
            [5, 6, 7, 8],
            [9, 10, 11, 12]
        ]
    )
    torch.testing.assert_allclose(out, expected) 

def test_first_pooling_with_mask():
    lp = pooling.FirstPooling()
    out = lp(TENSOR, padding_mask=MASK)

    assert out.size() == torch.Size([2, 4])

    expected = Tensor(
        [
            [1, 2, 3, 4],
            [1, 2, 3, 4]
        ]
    )
    torch.testing.assert_allclose(out, expected) 

def test_sum_pooling_with_mask():
    lp = pooling.SumPooling()
    out = lp(TENSOR, padding_mask=MASK)

    assert out.size() == torch.Size([2, 4])

    expected = Tensor(
        [
            [6, 8, 10, 12],
            [15, 18, 21, 24]
        ]
    )
    torch.testing.assert_allclose(out, expected) 

def test_avg_pooling_with_mask():
    lp = pooling.AvgPooling()
    out = lp(TENSOR, padding_mask=MASK)

    assert out.size() == torch.Size([2, 4])

    expected = Tensor(
        [
            [3, 4, 5, 6],
            [5, 6, 7, 8]
        ]
    )
    torch.testing.assert_allclose(out, expected) 

def test_model_updates(inputs, model, target, whitelist=[]):
    args = Args()
    args.balanceloss = False
    args.window_smooth = 0
    criterion = default_criterion.DefaultCriterion(args)
    optimizer = torch.optim.SGD(model.parameters(), lr=1.)
    optimizer.zero_grad()
    params = list(model.named_parameters())
    initial_params = [(name, p.clone()) for (name, p) in params]
    output = model(inputs)
    meta = {}
    _, loss, _ = criterion(output, target, meta)
    loss.backward()
    optimizer.step()
    for (_, p0), (name, p1) in zip(initial_params, params):
        if name in whitelist:
            continue
        try:
            np.testing.assert_raises(AssertionError, torch.testing.assert_allclose, p0, p1)
        except AssertionError:
            if 'bias' in name:
                print('Warning: {} not updating'.format(name))
                continue
            if p1.grad.norm() > 1e-6:
                print('Warning: {} not significantly updating'.format(name))
                continue
            print('{} not updating'.format(name))
            for (nn1, pp1) in params:
                print('{} grad: {}'.format(nn1, pp1.grad.norm().item()))
            import pdb
            pdb.set_trace()
            raise 

def test_return_ticket_np_torch(dim):
    """ Test torch --> np --> torch --> np conversion"""
    max_tested_ndim = 4
    # Random tensor shape
    tensor_shape = [random.randint(1, 10) for _ in range(max_tested_ndim)]
    # Make sure complex dimension has even shape
    tensor_shape[dim] = 2 * tensor_shape[dim]
    complex_tensor = torch.randn(tensor_shape)
    np_array = transforms.to_numpy(complex_tensor, dim=dim)
    tensor_back = transforms.from_numpy(np_array, dim=dim)
    np_back = transforms.to_numpy(tensor_back, dim=dim)
    # Check torch --> np --> torch
    assert_allclose(complex_tensor, tensor_back)
    # Check np --> torch --> np
    np.testing.assert_allclose(np_array, np_back)