Implement deepsets and deepsets policy, add tests

tests/nets/test_deepsets.py

@@ -0,0 +1,66 @@
+import torch
+import random
+from src.nets import deepsets as ds
+
+ds_config = {
+    "input_dim": 5,
+    "phi": {
+        "hidden_n": 1,
+        "hidden_dim": 10,
+    },
+    "latent_dim": 8,
+    "rho": {
+        "hidden_n": 1,
+        "hidden_dim": 10,
+    },
+    "output_dim" : 1,
+}
+
+def test_constructor():
+    m = ds.DeepSetsModule.from_config(ds_config)
+
+def test_phi():
+    input_dim = 5
+    phi = ds.Phi(input_dim, 1, 10, 2)
+    x = torch.rand(7, input_dim)
+    y = phi(x)
+    assert y.shape == torch.Size([7, 2])
+
+    y = phi(torch.rand(input_dim))
+    y = phi(torch.rand(7,7,7,input_dim))
+
+def test_deepsets():
+    m = ds.DeepSetsModule.from_config(ds_config)
+
+    input_dim = ds_config["input_dim"]
+    n_dynamic = random.randint(5, 15)
+    x = torch.rand(n_dynamic, input_dim)
+
+    n_batch = 20
+    x = x.unsqueeze(0).expand(n_batch, n_dynamic, input_dim)
+
+    y = m(x)
+    assert y.shape == torch.Size([n_batch, ds_config["output_dim"]])
+
+    for i in range(n_batch):
+        assert torch.allclose(y[i], y[0])
+
+def test_deepsets_computation():
+    n_dynamic = random.randint(5,15)
+    n_batch = 7
+    input_dim = 5
+    output_dim = 3
+    x = torch.rand(n_dynamic, input_dim)
+    x = x.unsqueeze(0).expand(n_batch, n_dynamic, input_dim)
+    assert x.shape == torch.Size([n_batch, n_dynamic, input_dim])
+
+    phi = torch.nn.Linear(input_dim, output_dim)
+
+    y = torch.stack(tuple(phi(instance) for instance in x.unbind(-2)), dim=-2)
+    assert y.shape == torch.Size([n_batch, n_dynamic, output_dim])
+
+    y = y.sum(dim=-2)
+    assert y.shape == torch.Size([n_batch, output_dim])
+
+    for i in range(n_batch):
+        assert torch.allclose(y[i], y[0])