SlowGuess/NeurIPS_2023_OCR huggingface dataset:
nu_x -- lr * psi_prime
print(f'stepe=step,loss=loss:.4g},rate=rate:.4g},distortion=  $\equiv$  distortion:.4g}）
return nu_x, nu_w
Run a toy example on 2D Gaussian samples.
rng =jax.random.PRNGKey(0)
X = jax.random.normal(rng, [10, 2])
nu_x, nu_w = wgd(X, n=4, rd_lambda=2., num_steps=100, lr=0.1, lng=rng)
60   
61   
62   
63   
64   
65   
66   
67   
68   
69   
70   
71   
72   
73   
74   
75   
76   
77   
78   
79   
80   
81   
82   
83   
84   
85   
86