WIP: GraphCast_small in SQL, with DeepMind's published weights

benchmarks/ml/01_attention.py

@@ -185,7 +185,7 @@ def main():
         want = attention(q, k, v)
         want_grads = jax.grad(loss, argnums=(0, 1, 2))(q, k, v)
 
-        for host in hosts(args.hosts):
+        for host in hosts(args.hosts, args.memory):
             for name, data in zip("qkv", (q, k, v)):
                 host.load(name, table(data))
 

benchmarks/ml/02_nanogpt.py

@@ -569,7 +569,7 @@ def main():
             )
         )
 
-        for host in hosts(args.hosts):
+        for host in hosts(args.hosts, args.memory):
             reference = dict(init)
             moments = {
                 n: (np.zeros_like(w), np.zeros_like(w)) for n, w in init.items()