def METHOD_NAME():
    op = ProxylessMixedLayer(
        {
            'a': MultiInputLayer(1),
            'b': MultiInputLayer(3)
        },
        nn.Parameter(torch.randn(2)),
        nn.Softmax(-1), 'proxyless'
    )
    optimizer = torch.optim.SGD(op.arch_parameters(), 0.1)
    for retry in range(10):
        q = torch.randn(1, 3, 9, 9).requires_grad_()
        k = torch.randn(1, 3, 9, 8).requires_grad_()
        v = None if retry < 3 else torch.randn(1, 3, 9, 7).requires_grad_()
        mask = None if retry < 5 else torch.randn(1, 3, 9, 6).requires_grad_()
        op.resample({})
        if mask is None:
            if v is None:
                y = op(q, k)
            else:
                y = op(q, k, v)
        else:
            y = op(q, k, v, mask)
        y = y[0].sum() + y[1].sum()
        optimizer.zero_grad()
        y.backward()
        assert op._arch_alpha.grad.abs().sum().item() != 0, op._arch_alpha.grad