1import torch
 2from random_attracting_rnn import GeometricallyAttractingRandomRNN
 3
 4device = 'cuda' if torch.cuda.is_available() else 'cpu'
 5try:
 6    torch.manual_seed(2803)
 7    m = GeometricallyAttractingRandomRNN(5, 7, candidates=3, target_rho=.9).to(device)
 8    x = torch.randn(11, 4, 5, device=device)
 9    states, routes = m(x, stochastic=False)
10    loss = states.square().mean() + m.contraction_penalty(routes.mean(0))
11    loss.backward()
12    assert states.shape == (11, 4, 7)
13    assert routes.shape == (11, 4, 3)
14    assert torch.isfinite(loss) and m.W.grad is not None
15    # Make the response of the regularizer explicit.
16    with torch.no_grad():
17        m.W.mul_(.5)
18    low = float(m.contraction_penalty(torch.full((4, 3), 1/3, device=device)))
19    with torch.no_grad():
20        m.W.mul_(3.)
21    high = float(m.contraction_penalty(torch.full((4, 3), 1/3, device=device)))
22    print({'device': device, 'loss': float(loss), 'gain_norms': m.gains().detach().cpu().tolist(),
23           'low_gain_penalty': low, 'high_gain_penalty': high})
24except Exception as e:
25    if device == 'cuda':
26        print('CUDA failed, retrying CPU:', repr(e))
27        m = GeometricallyAttractingRandomRNN(5, 7, candidates=3)
28        x = torch.randn(11, 4, 5)
29        states, routes = m(x, stochastic=False)
30        print({'device':'cpu', 'states':tuple(states.shape), 'routes':tuple(routes.shape)})
31    else:
32        raise