diff --git a/code/.gitignore b/liar_dice/.gitignore
similarity index 100%
rename from code/.gitignore
rename to liar_dice/.gitignore
diff --git a/liar_dice/__init__.py b/liar_dice/__init__.py
new file mode 100644
index 0000000..11e54d0
--- /dev/null
+++ b/liar_dice/__init__.py
@@ -0,0 +1,38 @@
+from dataclasses import dataclass
+from typing import Literal, TypeAlias, TypedDict
+
+from torch import Tensor
+from torch.optim.optimizer import StateDict
+
+Roll: TypeAlias = tuple[int, ...]
+PlayerId: TypeAlias = Literal[0, 1]
+MatchResult: TypeAlias = Literal[-1, 1]
+
+State = Tensor  # Public state. Includes previous actions. Initialized as `zeros(self.D_PUB)`.
+Priv = Tensor
+# Private state of a player. Initialized as `zeros(self.D_PRI)`. including the perspective for the scores
+
+Variant = Literal["normal", "joker", "stairs"]
+PruneType = Literal["zero", "upper", "lower", "us", "them", "avg"]
+
+
+@dataclass
+class TrainArg:
+    d1: int
+    d2: int
+    sides: int = 6
+    variant: Variant = "normal"
+    eps: float = 1e-2
+    layers: int = 4
+    layer_size: int = 100
+    lr: float = 1e-3
+    w: float = 1e-2
+    path: str = "model.pt"
+    device: Literal["cpu", "cuda"] = "cpu"
+
+
+class CheckpointFile(TypedDict):
+    epoch: int
+    model_state_dict: StateDict
+    optimizer_state_dict: StateDict
+    args: TrainArg
diff --git a/liar_dice/__main__.py b/liar_dice/__main__.py
new file mode 100644
index 0000000..e69de29
diff --git a/code/compare_models.py b/liar_dice/compare_models.py
similarity index 100%
rename from code/compare_models.py
rename to liar_dice/compare_models.py
diff --git a/code/convert_to_onnx.py b/liar_dice/convert_to_onnx.py
similarity index 100%
rename from code/convert_to_onnx.py
rename to liar_dice/convert_to_onnx.py
diff --git a/code/elo.py b/liar_dice/elo.py
similarity index 100%
rename from code/elo.py
rename to liar_dice/elo.py
diff --git a/code/lbr.py b/liar_dice/lbr.py
similarity index 70%
rename from code/lbr.py
rename to liar_dice/lbr.py
index 0373f8c..ab6ae37 100644
--- a/code/lbr.py
+++ b/liar_dice/lbr.py
@@ -2,17 +2,29 @@
 Computes the (possibly Local) Best Response strategy and exploitability of a model
 """
 
-import torch
 import argparse
-
-from snyd import Net, Game, calc_args
 from collections import Counter
+from typing import cast
+
+import torch
+
+from liar_dice import PlayerId, PruneType, Roll
+from liar_dice.snyd import Game, Net, calc_args
 
 # in normal form games, given strategy matrix A for player 1 (hand -> action_probs)
 # we will play a deterministic strategy b (hand -> action)
 
 
-def lbr(game, roll, us, prune=0, prune_type="zero"):
+def lbr(
+    game: Game,
+    roll: Roll,
+    us: PlayerId,
+    prune: float = 0,
+    prune_type: PruneType = "zero",
+) -> float:
+    """
+    Compute the Local Best Response strategy for a player in a game.
+    """
     # All possible opponent rolls - with multiplicity
     roll_cnt = Counter(game.rolls(1 - us))
     op_rolls = roll_cnt.keys()
@@ -23,7 +35,13 @@ def lbr(game, roll, us, prune=0, prune_type="zero"):
     # The opponents private state in each initial information state
     op_privs = [game.make_priv(op_roll, 1 - us) for op_roll in op_rolls]
 
-    def inner(state, reach_probs, likelihood):
+    def inner(
+        state: torch.Tensor, reach_probs: torch.Tensor, likelihood: float
+    ) -> float:
+        """
+        Compute the Local Best Response strategy for a player in a game.
+        """
+
         # basically simple max-chance algorithm.
         # possibly using our learned values for reducing the tree height
         # But I could start by making a version that just goes to the bottom
@@ -32,42 +50,50 @@ def inner(state, reach_probs, likelihood):
         cur = len(calls) % 2
 
         if likelihood < prune:
+
+            def calc_our_guess():
+                """
+                Even using our guess here, we will probably get a lower bound,
+                since we are replacing a "perfect" (best response) strategy
+                with whatever strategy the model suggests.
+                However it's not guaranteed to be (a lower bound), since the
+                value network could be overly optimistic."""
+                return game.model(game.make_priv(roll, us), state).item()
+
+            def calc_op_guess():
+                """
+                This is the same as calc_our_guess, but for the opponent.
+                Opponent guess is relative to them, not us, so it's negated.
+                """
+                return -sum(
+                    prob * game.model(op_priv, state).item()
+                    for prob, op_priv in zip(reach_probs, op_privs)
+                )
+
             # print('Pruning', likelihood)
             if prune_type == "zero":
                 return 0
             # I can return +1 or -1 here to get either an upper or lower
             # bound on exploitability
-            if prune_type == "upper":
+            elif prune_type == "upper":
                 return 1
-            if prune_type == "lower":
+            elif prune_type == "lower":
                 return -1
-            if prune_type in ["us", "avg"]:
-                # Even using our guess here, we will probably get a lower bound,
-                # since we are replacing a "perfect" (best response) strategy
-                # with whatever strategy the model suggests.
-                # However it's not guaranteed to be (a lower bound), since the
-                # value network could be overly optimistic.
-                our_guess = game.model(game.make_priv(roll, us), state).item()
-            if prune_type in ["them", "avg"]:
-                op_guess = sum(
-                    prob * game.model(op_priv, state).item()
-                    for prob, op_priv in zip(reach_probs, op_privs)
-                )
-                # Remember opponent guess is relative to them, not us
-                op_guess = -op_guess
-            if prune_type == "us":
-                return our_guess
-            if prune_type == "them":
-                return op_guess
-            if prune_type == "avg":
-                return (our_guess + op_guess) / 2
-            assert False
+            elif prune_type == "us":
+                return calc_our_guess()
+            elif prune_type == "them":
+                return calc_op_guess()
+            elif prune_type == "avg":
+                return (calc_our_guess() + calc_op_guess()) / 2
+            else:
+                raise ValueError(f"Unknown prune type: {prune_type}")
 
         if calls and calls[-1] == game.LIE_ACTION:
             prev_call = calls[-2] if len(calls) >= 2 else -1
             res = 0
             # Take the average outcome over our opponents possible rolls
             for prob, op_roll in zip(reach_probs, op_rolls):
+                prob = cast(float, prob)
                 r1, r2 = (roll, op_roll) if us == 0 else (op_roll, roll)
                 correct = game.evaluate_call(r1, r2, prev_call)
                 # If prev_call is good, and we are now, it mean we won
@@ -106,9 +132,12 @@ def inner(state, reach_probs, likelihood):
             score = 0
             for action in range(last_call + 1, game.N_ACTIONS):
                 ai = action - last_call - 1
-                pa = reach_probs @ policies[:, ai]
-                if torch.isclose(pa, torch.tensor(0.0)):
+                pa_mat = reach_probs @ policies[:, ai]
+
+                if torch.isclose(pa_mat, torch.tensor(0.0)):
                     continue
+
+                pa = float(pa_mat)
                 # Bayes: P(R|A) = P(R)P(A|R)/P(A)
                 # P(A) = P(R)P(A|R)/sum(P(A|r)P(r) for r in Rs)
                 new_reach_probs = (reach_probs * policies[:, ai]) / pa
@@ -135,7 +164,7 @@ def main():
     # options=['upper', 'lower', 'zero', 'us', 'them', 'avg'])
     args = parser.parse_args()
 
-    checkpoint = torch.load(args.path)
+    checkpoint = torch.load(args.path)  # type: ignore
     train_args = checkpoint["args"]
 
     D_PUB, D_PRI, *_ = calc_args(
@@ -147,12 +176,14 @@ def main():
         model, train_args.d1, train_args.d2, train_args.sides, train_args.variant
     )
 
-    for player in range(2):
+    for player in (0, 1):
         print(f"Testing exploitability of player {1-player}")
 
         total_val = 0
         total_cnt = 0
         for roll, cnt in Counter(game.rolls(player)).items():
+            if cnt != 1:
+                print("Warning: Duplicate roll", roll)
             print("Exploiting with roll", roll)
             total_val += cnt * lbr(game, roll, player, args.prune, args.type)
             total_cnt += cnt
diff --git a/code/play_model.py b/liar_dice/play_model.py
similarity index 86%
rename from code/play_model.py
rename to liar_dice/play_model.py
index 0ebffa2..3d911db 100644
--- a/code/play_model.py
+++ b/liar_dice/play_model.py
@@ -2,23 +2,19 @@
 # This script allows you to play against the model from the terminal
 ################################################################################
 
-import random
-import torch
-from torch import nn
-import itertools
-import numpy as np
-import math
-from collections import Counter
 import argparse
+import random
 import re
 
-from snyd import *
+import torch
+
+from liar_dice.snyd import *
 
 parser = argparse.ArgumentParser()
 parser.add_argument("path", type=str, help="Path of model")
 args = parser.parse_args()
 
-checkpoint = torch.load(args.path, map_location=torch.device("cpu"))
+checkpoint = torch.load(args.path, map_location=torch.device(device="cpu"))  # type: ignore
 train_args = checkpoint["args"]
 
 D_PUB, D_PRI, *_ = calc_args(
@@ -30,7 +26,7 @@
 
 
 class Human:
-    def get_action(self, state):
+    def get_action(self, state: State) -> int:
         last_call = game.get_last_call(state)
         while True:
             call = input('Your call [e.g. 24 for 2 fours, or "lie" to call a bluff]: ')
@@ -53,10 +49,10 @@ def __repr__(self):
 
 
 class Robot:
-    def __init__(self, priv):
+    def __init__(self, priv: Priv):
         self.priv = priv
 
-    def get_action(self, state):
+    def get_action(self, state: State) -> int:
         last_call = game.get_last_call(state)
         return game.sample_action(self.priv, state, last_call, eps=0)
 
@@ -64,7 +60,7 @@ def __repr__(self):
         return "robot"
 
 
-def repr_action(action):
+def repr_action(action: int) -> str:
     action = int(action)
     if action == -1:
         return "nothing"
@@ -77,9 +73,9 @@ def repr_action(action):
 
 while True:
     while (ans := input("Do you want to go first? [y/n/r] ")) not in ["y", "n", "r"]:
-        pass
+        pass  # r: bot vs bot
 
-    r1 = random.choice(list(game.rolls(0)))
+    r1 = random.choice(list(game.rolls(player=0)))
     r2 = random.choice(list(game.rolls(1)))
     privs = [game.make_priv(r1, 0), game.make_priv(r2, 1)]
     state = game.make_state()
@@ -92,10 +88,12 @@ def repr_action(action):
         players = [Robot(privs[0]), Human()]
     elif ans == "r":
         players = [Robot(privs[0]), Robot(privs[1])]
+    else:
+        raise ValueError(f"Invalid answer: {ans}")
 
     cur = 0
     while True:
-        action = players[cur].get_action(state)
+        action: int = players[cur].get_action(state)
         print()
         print(f"> The {players[cur]} called {repr_action(action)}!")
 
diff --git a/liar_dice/py.typed b/liar_dice/py.typed
new file mode 100644
index 0000000..e69de29
diff --git a/code/snyd.py b/liar_dice/snyd.py
similarity index 57%
rename from code/snyd.py
rename to liar_dice/snyd.py
index 1bc77f9..ce0a738 100644
--- a/code/snyd.py
+++ b/liar_dice/snyd.py
@@ -2,91 +2,94 @@
 Contains the Game class, as well as various simple netural network architectures.
 """
 
-import random
-import torch
-from torch import nn
 import itertools
-import numpy as np
-import math
 from collections import Counter
 
+from torch import cat, nn, utils, zeros
 
-class Net(torch.nn.Module):
-    def __init__(self, d_pri, d_pub):
-        super().__init__()
+from liar_dice import PlayerId, Priv, Roll, State, Variant
 
-        hiddens = (100,) * 4
+
+class Net(nn.Module):
+    def __init__(self, d_pri: int, d_pub: int):
+        super().__init__()  # type: ignore # pytorch convention
+
+        hidden = (100,) * 4
 
         # Bilinear can't be used inside nn.Sequantial
         # https://github.com/pytorch/pytorch/issues/37092
-        self.layer0 = torch.nn.Bilinear(d_pri, d_pub, hiddens[0])
+        self.layer0 = nn.Bilinear(d_pri, d_pub, hidden[0])
 
-        layers = [torch.nn.ReLU()]
-        for size0, size1 in zip(hiddens, hiddens[1:]):
-            layers += [torch.nn.Linear(size0, size1), torch.nn.ReLU()]
-        layers += [torch.nn.Linear(hiddens[-1], 1), nn.Tanh()]
+        layers = [nn.ReLU()]
+        for size0, size1 in zip(hidden, hidden[1:]):
+            layers += [nn.Linear(size0, size1), nn.ReLU()]
+        layers += [nn.Linear(hidden[-1], 1), nn.Tanh()]
         self.seq = nn.Sequential(*layers)
 
-    def forward(self, priv, pub):
+    def forward(self, priv: Priv, pub: State):
         joined = self.layer0(priv, pub)
         return self.seq(joined)
 
 
-class NetConcat(torch.nn.Module):
-    def __init__(self, d_pri, d_pub):
-        super().__init__()
+class NetConcat(nn.Module):
+    def __init__(self, d_pri: int, d_pub: int):
+        super().__init__()  # type: ignore # pytorch convention
 
         hiddens = (500, 400, 300, 200, 100)
 
-        layers = [torch.nn.Linear(d_pri + d_pub, hiddens[0]), torch.nn.ReLU()]
+        layers = [nn.Linear(d_pri + d_pub, hiddens[0]), nn.ReLU()]
         for size0, size1 in zip(hiddens, hiddens[1:]):
-            layers += [torch.nn.Linear(size0, size1), torch.nn.ReLU()]
-        layers += [torch.nn.Linear(hiddens[-1], 1), nn.Tanh()]
+            layers += [nn.Linear(size0, size1), nn.ReLU()]
+        layers += [nn.Linear(hiddens[-1], 1), nn.Tanh()]
         self.seq = nn.Sequential(*layers)
 
-    def forward(self, priv, pub):
+    def forward(self, priv: Priv, pub: State):
         if len(priv.shape) == 1:
-            joined = torch.cat((priv, pub), dim=0)
+            joined = cat((priv, pub), dim=0)
         else:
-            joined = torch.cat((priv, pub), dim=1)
+            joined = cat((priv, pub), dim=1)
         return self.seq(joined)
 
 
-class NetCompBilin(torch.nn.Module):
-    def __init__(self, d_pri, d_pub):
-        super().__init__()
+class NetCompBilin(nn.Module):
+    def __init__(self, d_pri: int, d_pub: int):
+        super().__init__()  # type: ignore # pytorch convention
 
         hiddens = (100,) * 4
 
         middle = 500
-        self.layer_pri = torch.nn.Linear(d_pri, middle)
-        self.layer_pub = torch.nn.Linear(d_pub, middle)
+        self.layer_pri = nn.Linear(d_pri, middle)
+        self.layer_pub = nn.Linear(d_pub, middle)
 
-        layers = [torch.nn.ReLU(), torch.nn.Linear(middle, hiddens[0]), torch.nn.ReLU()]
+        layers = [nn.ReLU(), nn.Linear(middle, hiddens[0]), nn.ReLU()]
         for size0, size1 in zip(hiddens, hiddens[1:]):
-            layers += [torch.nn.Linear(size0, size1), torch.nn.ReLU()]
-        layers += [torch.nn.Linear(hiddens[-1], 1), nn.Tanh()]
+            layers += [nn.Linear(size0, size1), nn.ReLU()]
+        layers += [nn.Linear(hiddens[-1], 1), nn.Tanh()]
         self.seq = nn.Sequential(*layers)
 
-    def forward(self, priv, pub):
+    def forward(self, priv: Priv, pub: State):
         joined = self.layer_pri(priv) * self.layer_pub(pub)
         return self.seq(joined)
 
 
+"""
+These are never used
+
 class Resid(nn.Module):
-    def __init__(self, in_channels, out_channels, kernel_size):
-        super().__init__()
+    def __init__(self, in_channels: int, out_channels: int, kernel_size: int):
+        super().__init__()  # type: ignore # pytorch convention
         self.conv = nn.Conv1d(in_channels, out_channels, kernel_size)
 
+    @no_type_check
     def forward(self, x):
         y = self.conv(y)
 
 
-class Net3(torch.nn.Module):
-    def __init__(self, d_pri, d_pub):
-        super().__init__()
+class Net3(nn.Module):
+    def __init__(self, d_pri: int, d_pub: int):
+        super().__init__()  # type: ignore # pytorch convention
 
-        def conv(channels, size):
+        def conv(channels: int, size: int):
             return nn.Sequential(
                 nn.Conv1d(1, channels, kernel_size=size),
                 nn.BatchNorm1d(channels),
@@ -95,24 +98,25 @@ def conv(channels, size):
                 nn.BatchNorm1d(channels),
             )
 
-        # Bilinear can't be used inside nn.Sequantial
+        # Bilinear can't be used inside nn.Sequential
         # https://github.com/pytorch/pytorch/issues/37092
-        self.layer0 = torch.nn.Bilinear(d_pri, d_pub, hiddens[0])
+        self.layer0 = nn.Bilinear(d_pri, d_pub, hiddens[0])
 
-        layers = [torch.nn.ReLU()]
+        layers = [nn.ReLU()]
         for size0, size1 in zip(hiddens, hiddens[1:]):
-            layers += [torch.nn.Linear(size0, size1), torch.nn.ReLU()]
-        layers += [torch.nn.Linear(hiddens[-1], 1), nn.Tanh()]
+            layers += [nn.Linear(size0, size1), nn.ReLU()]
+        layers += [nn.Linear(hiddens[-1], 1), nn.Tanh()]
         self.seq = nn.Sequential(*layers)
 
     def forward(self, priv, pub):
         joined = self.layer0(priv, pub)
         return self.seq(joined)
+"""
 
 
-class Net2(torch.nn.Module):
-    def __init__(self, d_pri, d_pub):
-        super().__init__()
+class Net2(nn.Module):
+    def __init__(self, d_pri: int, d_pub: int):
+        super().__init__()  # type: ignore # pytorch convention
 
         channels = 20
         self.left = nn.Sequential(
@@ -130,15 +134,15 @@ def __init__(self, d_pri, d_pub):
         )
 
         layer_size = 100
-        self.bilin = torch.nn.Bilinear(d_pri, d_pub, layer_size)
+        self.bilin = nn.Bilinear(d_pri, d_pub, layer_size)
 
-        layers = [torch.nn.ReLU()]
-        for i in range(3):
-            layers += [torch.nn.Linear(layer_size, layer_size), torch.nn.ReLU()]
-        layers += [torch.nn.Linear(layer_size, 1), nn.Tanh()]
+        layers = [nn.ReLU()]
+        for _ in range(3):
+            layers += [nn.Linear(layer_size, layer_size), nn.ReLU()]
+        layers += [nn.Linear(layer_size, 1), nn.Tanh()]
         self.seq = nn.Sequential(*layers)
 
-    def forward(self, priv, pub):
+    def forward(self, priv: Priv, pub: State):
         if len(priv.shape) == 1:
             assert len(pub.shape) == 1
             priv = priv.unsqueeze(0)
@@ -150,45 +154,49 @@ def forward(self, priv, pub):
         return self.seq(mixed)
 
 
-def calc_args(d1, d2, sides, variant):
-    # Maximum call is (D1+D2) 6s
-    D_PUB = (d1 + d2) * sides
+def calc_args(
+    d1: int, d2: int, sides: int, variant: Variant
+) -> tuple[int, int, int, int, int, int, int]:
+    # Maximum call is (D1+D2) 6's (6 as the default number of sides)
+    d_pub = (d1 + d2) * sides
     if variant == "stairs":
-        D_PUB = 2 * (d1 + d2) * sides
+        d_pub = 2 * (d1 + d2) * sides
 
     # We also have a feature that means "game has been called"
-    LIE_ACTION = D_PUB
-    D_PUB += 1
+    lie_action = d_pub
+    d_pub += 1
 
     # Total number of actions in the game
-    N_ACTIONS = D_PUB
+    n_actions = d_pub
 
     # Add two extra features to keep track of who's to play next
-    CUR_INDEX = D_PUB
-    D_PUB += 1
+    cur_index = d_pub
+    d_pub += 1
 
     # Duplicate for other player
-    D_PUB_PER_PLAYER = D_PUB
-    D_PUB *= 2
+    d_pub_per_player = d_pub
+    d_pub *= 2
 
     # One player technically may have a smaller private space than the other,
     # but we just take the maximum for simplicity
-    D_PRI = max(d1, d2) * sides
+    d_pri = max(d1, d2) * sides
 
-    # And then two features to describe from whos perspective we
+    # And then two features to describe from who's perspective we
     # are given the private information
-    PRI_INDEX = D_PRI
-    D_PRI += 2
+    pri_index = d_pri
+    d_pri += 2
 
-    return D_PUB, D_PRI, N_ACTIONS, LIE_ACTION, CUR_INDEX, PRI_INDEX, D_PUB_PER_PLAYER
+    return d_pub, d_pri, n_actions, lie_action, cur_index, pri_index, d_pub_per_player
 
 
 class Game:
-    def __init__(self, model, d1, d2, sides, variant):
-        self.model = model
-        self.D1 = d1
-        self.D2 = d2
-        self.SIDES = sides
+    def __init__(
+        self, model: nn.Module, d1: int, d2: int, sides: int, variant: Variant
+    ):
+        self.model: nn.Module = model
+        self.D1: int = d1
+        self.D2: int = d2
+        self.SIDES: int = sides
         self.VARIANT = variant
 
         (
@@ -201,10 +209,10 @@ def __init__(self, model, d1, d2, sides, variant):
             self.D_PUB_PER_PLAYER,
         ) = calc_args(d1, d2, sides, variant)
 
-    def make_regrets(self, priv, state, last_call):
+    def make_regrets(self, priv: Priv, state: State, last_call: int) -> list[float]:
         """
         priv: Private state, including the perspective for the scores
-        state: Public statte
+        state: Public state
         last_call: Last action taken by a player. Returned regrets will be for actions after this one.
         """
 
@@ -227,7 +235,7 @@ def make_regrets(self, priv, state, last_call):
         # The Hedge method
         # return [math.exp(10*(vi - v)) for vi in vs]
 
-    def evaluate_call(self, r1, r2, last_call):
+    def evaluate_call(self, r1: Roll, r2: Roll, last_call: int) -> bool:
         # Players have rolled r1, and r2.
         # Previous actions are `state`
         # Player `caller` just called lie. (This is not included in last_call)
@@ -244,17 +252,22 @@ def evaluate_call(self, r1, r2, last_call):
         cnt = Counter(r1 + r2)
         if self.VARIANT == "normal":
             actual = cnt[d]
-        if self.VARIANT == "joker":
+        elif self.VARIANT == "joker":
             actual = cnt[d] + cnt[1] if d != 1 else cnt[d]
-        if self.VARIANT == "stairs":
+        elif self.VARIANT == "stairs":
+            actual = 0
             if all(r == i + 1 for r, i in zip(r1, range(self.SIDES))):
                 actual += 2 * len(r1) - r1.count(d)
             if all(r == i + 1 for r, i in zip(r2, range(self.SIDES))):
                 actual += 2 * len(r2) - r1.count(d)
+        else:
+            raise ValueError(f"Unknown variant: {self.VARIANT}")
         # print(f'{r1=}, {r2=}, {last_call=}, {(n, d)=}, {actual=}', actual >= n)
         return actual >= n
 
-    def policy(self, priv, state, last_call, eps=0):
+    def policy(
+        self, priv: Priv, state: State, last_call: int, eps: float = 0
+    ) -> list[float]:
         regrets = self.make_regrets(priv, state, last_call)
         for i in range(len(regrets)):
             regrets[i] += eps
@@ -264,27 +277,29 @@ def policy(self, priv, state, last_call, eps=0):
             s = sum(regrets)
             return [r / s for r in regrets]
 
-    def sample_action(self, priv, state, last_call, eps):
+    def sample_action(
+        self, priv: Priv, state: State, last_call: int, eps: float
+    ) -> int:
         pi = self.policy(priv, state, last_call, eps)
-        action = next(iter(torch.utils.data.WeightedRandomSampler(pi, num_samples=1)))
+        action = next(iter(utils.data.WeightedRandomSampler(pi, num_samples=1)))
         return action + last_call + 1
 
-    def apply_action(self, state, action):
+    def apply_action(self, state: State, action: int) -> State:
         new_state = state.clone()
         self._apply_action(new_state, action)
         return new_state
 
-    def _apply_action(self, state, action):
+    def _apply_action(self, state: State, action: int) -> State:
         # Inplace
-        cur = self.get_cur(state)
+        cur = self.get_current_player(state)
         state[action + cur * self.D_PUB_PER_PLAYER] = 1
         state[self.CUR_INDEX + cur * self.D_PUB_PER_PLAYER] = 0
         state[self.CUR_INDEX + (1 - cur) * self.D_PUB_PER_PLAYER] = 1
         return state
 
-    def make_priv(self, roll, player):
+    def make_priv(self, roll: Roll, player: int) -> Priv:
         assert player in [0, 1]
-        priv = torch.zeros(self.D_PRI)
+        priv = zeros(self.D_PRI)
         priv[self.PRI_INDEX + player] = 1
         # New method inspired by Chinese poker paper
         cnt = Counter(roll)
@@ -296,16 +311,16 @@ def make_priv(self, roll, player):
         #     priv[i * self.SIDES + r - 1] = 1
         return priv
 
-    def make_state(self):
-        state = torch.zeros(self.D_PUB)
+    def make_state(self) -> State:
+        state = zeros(self.D_PUB)
         state[self.CUR_INDEX] = 1
         return state
 
-    def get_cur(self, state):
-        # Player index in {0,1} is equal to one-hot encoding of player 2
+    def get_current_player(self, state: State) -> int:
+        # Player index in {0,1} is equal to 1 - (hot encoding of player 2)
         return 1 - int(state[self.CUR_INDEX])
 
-    def rolls(self, player):
+    def rolls(self, player: PlayerId) -> list[tuple[int, ...]]:
         assert player in [0, 1]
         n_faces = self.D1 if player == 0 else self.D2
         return [
@@ -313,14 +328,19 @@ def rolls(self, player):
             for r in itertools.product(range(1, self.SIDES + 1), repeat=n_faces)
         ]
 
-    def get_calls(self, state):
+    def get_calls(self, state: State) -> list[int]:
+        """
+        Returns the indices of the calls.
+        find the first 1 in the first part of the state, and return the index
+        the length of the returned list is same as the dimensions of the state
+        """
         merged = (
             state[: self.CUR_INDEX]
             + state[self.D_PUB_PER_PLAYER : self.D_PUB_PER_PLAYER + self.CUR_INDEX]
         )
-        return (merged == 1).nonzero(as_tuple=True)[0].tolist()
+        return (merged == 1).nonzero(as_tuple=True)[0].tolist()  # type: ignore # pytorch stub problem
 
-    def get_last_call(self, state):
+    def get_last_call(self, state: State) -> int:
         ids = self.get_calls(state)
         if not ids:
             return -1
diff --git a/code/train.py b/liar_dice/train.py
similarity index 66%
rename from code/train.py
rename to liar_dice/train.py
index e507814..257e774 100644
--- a/code/train.py
+++ b/liar_dice/train.py
@@ -2,76 +2,39 @@
 Train a new model from self-play.
 """
 
-import random
-import torch
-from torch import nn
-import itertools
-import math
-from collections import Counter
 import argparse
+import itertools
 import os
+import random
+from collections import Counter
+from typing import cast
 
-from snyd import *
-
-parser = argparse.ArgumentParser()
-parser.add_argument("d1", type=int, help="Number of dice for player 1")
-parser.add_argument("d2", type=int, help="Number of dice for player 2")
-parser.add_argument("--sides", type=int, default=6, help="Number of sides on the dice")
-parser.add_argument(
-    "--variant", type=str, default="normal", help="one of normal, joker, stairs"
-)
-parser.add_argument(
-    "--eps", type=float, default=1e-2, help="Added to regrets for exploration"
-)
-parser.add_argument(
-    "--layers", type=int, default=4, help="Number of fully connected layers"
-)
-parser.add_argument(
-    "--layer-size", type=int, default=100, help="Number of neurons per layer"
-)
-parser.add_argument("--lr", type=float, default=1e-3, help="LR = lr/t")
-parser.add_argument("--w", type=float, default=1e-2, help="weight decay")
-parser.add_argument(
-    "--path", type=str, default="model.pt", help="Where to save checkpoints"
-)
-
-args = parser.parse_args()
-
-
-# Check if there is a model we should continue training
-if os.path.isfile(args.path):
-    checkpoint = torch.load(args.path)
-    print(f"Using args from {args.path}")
-    old_path = args.path
-    args = checkpoint["args"]
-    args.path = old_path
-else:
-    checkpoint = None
-
-# Model : (private state, public state) -> value
-D_PUB, D_PRI, *_ = calc_args(args.d1, args.d2, args.sides, args.variant)
-model = NetConcat(D_PRI, D_PUB)
-# model = Net(D_PRI, D_PUB)
-# model = Net2(D_PRI, D_PUB)
-game = Game(model, args.d1, args.d2, args.sides, args.variant)
-
-if checkpoint is not None:
-    print("Loading previous model for continued training")
-    model.load_state_dict(checkpoint["model_state_dict"])
+import torch
 
+from liar_dice import CheckpointFile, MatchResult, Priv, Roll, State, TrainArg
+from liar_dice.snyd import Game, NetConcat, calc_args
 
-device = torch.device("cuda")
-model.to(device)
 
+@torch.no_grad()  # type: ignore # pytorch stub problem
+def play(
+    r1: Roll, r2: Roll, replay_buffer: list[tuple[Priv, State, MatchResult]]
+) -> None:
+    privs = (
+        game.make_priv(r1, 0).to(device),
+        game.make_priv(r2, 1).to(device),
+    )
 
-@torch.no_grad()
-def play(r1, r2, replay_buffer):
-    privs = [game.make_priv(r1, 0).to(device), game.make_priv(r2, 1).to(device)]
+    def play_inner(state: State) -> MatchResult:
+        """
+        Play a game from the current state and return the result.
+        The result is from the perspective of the current player.
 
-    def play_inner(state):
-        cur = game.get_cur(state)
-        calls = game.get_calls(state)
-        assert cur == len(calls) % 2
+        Returns:
+            1 if we won, -1 if we lost
+        """
+        current_player = game.get_current_player(state)
+        calls = game.get_calls(state=state)
+        assert current_player == len(calls) % 2
 
         if calls and calls[-1] == game.LIE_ACTION:
             prev_call = calls[-2] if len(calls) >= 2 else -1
@@ -80,14 +43,16 @@ def play_inner(state):
 
         else:
             last_call = calls[-1] if calls else -1
-            action = game.sample_action(privs[cur], state, last_call, args.eps)
+            action = game.sample_action(
+                privs[current_player], state, last_call, args.eps
+            )
             new_state = game.apply_action(state, action)
             # Just classic min/max stuff
             res = -play_inner(new_state)
 
         # Save the result from the perspective of both sides
-        replay_buffer.append((privs[cur], state, res))
-        replay_buffer.append((privs[1 - cur], state, -res))
+        replay_buffer.append((privs[current_player], state, res))
+        replay_buffer.append((privs[1 - current_player], state, -res))
 
         return res
 
@@ -96,7 +61,7 @@ def play_inner(state):
         play_inner(state)
 
 
-def print_strategy(state):
+def print_strategy(state: State):
     total_v = 0
     total_cnt = 0
     for r1, cnt in sorted(Counter(game.rolls(0)).items()):
@@ -105,28 +70,37 @@ def print_strategy(state):
         rs = torch.tensor(game.make_regrets(priv, state, last_call=-1))
         if rs.sum() != 0:
             rs /= rs.sum()
-        strat = []
+        strategy: list[str] = []
         for action, prob in enumerate(rs):
             n, d = divmod(action, game.SIDES)
             n, d = n + 1, d + 1
             if d == 1:
-                strat.append(f"{n}:")
-            strat.append(f"{prob:.2f}")
-        print(r1, f"{float(v):.4f}".rjust(7), f"({cnt})", " ".join(strat))
+                strategy.append(f"{n}:")
+            strategy.append(f"{prob:.2f}")
+        print(r1, f"{float(v):.4f}".rjust(7), f"({cnt})", " ".join(strategy))
         total_v += v
         total_cnt += cnt
     print(f"Mean value: {total_v / total_cnt}")
 
 
-class ReciLR(torch.optim.lr_scheduler._LRScheduler):
-    def __init__(self, optimizer, gamma=1, last_epoch=-1, verbose=False):
+class ReciprocalLR(torch.optim.lr_scheduler.LRScheduler):
+    """
+    Reciprocal learning rate scheduler. Adjusts the learning rate according to
+    """
+
+    def __init__(
+        self,
+        optimizer: torch.optim.Optimizer,
+        gamma: float = 1,
+        last_epoch: int = -1,
+    ) -> None:
         self.gamma = gamma
-        super(ReciLR, self).__init__(optimizer, last_epoch, verbose)
+        super(ReciprocalLR, self).__init__(optimizer, last_epoch)
 
-    def get_lr(self):
+    def get_lr(self) -> list[float]:
         return [
             base_lr / (self.last_epoch + 1) ** self.gamma
-            for base_lr, group in zip(self.base_lrs, self.optimizer.param_groups)
+            for base_lr, _group in zip(self.base_lrs, self.optimizer.param_groups)
         ]
 
     def _get_closed_form_lr(self):
@@ -137,17 +111,18 @@ def _get_closed_form_lr(self):
 
 def train():
     optimizer = torch.optim.AdamW(model.parameters(), weight_decay=args.w)
-    scheduler = ReciLR(optimizer, gamma=0.5)
+    scheduler = ReciprocalLR(optimizer, gamma=0.5)
     value_loss = torch.nn.MSELoss()
     all_rolls = list(itertools.product(game.rolls(0), game.rolls(1)))
     for t in range(100_000):
-        replay_buffer = []
+        replay_buffer = list[tuple[Priv, State, MatchResult]]()
 
         BS = 100  # Number of rolls to include
+        # BS = 60466176  # 6^5 * 6^5, all possible rolls
         for r1, r2 in (
             all_rolls if len(all_rolls) <= BS else random.sample(all_rolls, BS)
         ):
-            play(r1, r2, replay_buffer)
+            play(r1, r2, replay_buffer=replay_buffer)
 
         random.shuffle(replay_buffer)
         privs, states, y = zip(*replay_buffer)
@@ -162,19 +137,19 @@ def train():
         loss = value_loss(y_pred, y)
         print(t, loss.item())
 
-        if t % 5 == 0:
-            with torch.no_grad():
-                print_strategy(game.make_state().to(device))
+        # if t % 5 == 0:
+        #     with torch.no_grad():
+        #         print_strategy(game.make_state().to(device))
 
         # Zero gradients, perform a backward pass, and update the weights.
         optimizer.zero_grad()
         loss.backward()
-        optimizer.step()
+        optimizer.step()  # type: ignore # pytorch stub problem
         scheduler.step()
 
         if (t + 1) % 10 == 0:
             print(f"Saving to {args.path}")
-            torch.save(
+            torch.save(  # type: ignore # pytorch stub problem
                 {
                     "epoch": t,
                     "model_state_dict": model.state_dict(),
@@ -184,7 +159,7 @@ def train():
                 args.path,
             )
         if (t + 1) % 1000 == 0:
-            torch.save(
+            torch.save(  # type: ignore # pytorch stub problem
                 {
                     "epoch": t,
                     "model_state_dict": model.state_dict(),
@@ -195,4 +170,54 @@ def train():
             )
 
 
+parser = argparse.ArgumentParser()
+parser.add_argument("d1", type=int, help="Number of dice for player 1")
+parser.add_argument("d2", type=int, help="Number of dice for player 2")
+parser.add_argument("--sides", type=int, default=6, help="Number of sides on the dice")
+parser.add_argument(
+    "--variant", type=str, default="normal", help="one of normal, joker, stairs"
+)
+parser.add_argument(
+    "--eps", type=float, default=1e-2, help="Added to regrets for exploration"
+)
+parser.add_argument(
+    "--layers", type=int, default=4, help="Number of fully connected layers"
+)
+parser.add_argument(
+    "--layer-size", type=int, default=100, help="Number of neurons per layer"
+)
+parser.add_argument("--lr", type=float, default=1e-3, help="LR = lr/t")
+parser.add_argument("--w", type=float, default=1e-2, help="weight decay")
+parser.add_argument(
+    "--path", type=str, default="model.pt", help="Where to save checkpoints"
+)
+
+args = TrainArg(**vars(parser.parse_args()))
+
+
+# Check if there is a model we should continue training
+if os.path.isfile(path=args.path):
+    checkpoint = cast(CheckpointFile, torch.load(args.path))  # type: ignore
+    print(f"Using args from {args.path}")
+    old_path = args.path
+    args = checkpoint["args"]
+    args.path = old_path
+else:
+    checkpoint = None
+
+# Model : (private state, public state) -> value
+D_PUB, D_PRI, *_ = calc_args(args.d1, args.d2, args.sides, args.variant)
+model = NetConcat(D_PRI, D_PUB)
+# model = Net(D_PRI, D_PUB)
+# model = Net2(D_PRI, D_PUB)
+game = Game(model, args.d1, args.d2, args.sides, args.variant)
+
+if checkpoint is not None:
+    print("Loading previous model for continued training")
+    model.load_state_dict(checkpoint["model_state_dict"])
+
+
+# device = torch.device("cuda")
+device = torch.device("cpu")
+model.to(device)
 train()
diff --git a/poetry.lock b/poetry.lock
new file mode 100644
index 0000000..c8eadd0
--- /dev/null
+++ b/poetry.lock
@@ -0,0 +1,603 @@
+# This file is automatically @generated by Poetry 1.7.1 and should not be changed by hand.
+
+[[package]]
+name = "colorama"
+version = "0.4.6"
+description = "Cross-platform colored terminal text."
+optional = false
+python-versions = "!=3.0.*,!=3.1.*,!=3.2.*,!=3.3.*,!=3.4.*,!=3.5.*,!=3.6.*,>=2.7"
+files = [
+    {file = "colorama-0.4.6-py2.py3-none-any.whl", hash = "sha256:4f1d9991f5acc0ca119f9d443620b77f9d6b33703e51011c16baf57afb285fc6"},
+    {file = "colorama-0.4.6.tar.gz", hash = "sha256:08695f5cb7ed6e0531a20572697297273c47b8cae5a63ffc6d6ed5c201be6e44"},
+]
+
+[[package]]
+name = "filelock"
+version = "3.16.1"
+description = "A platform independent file lock."
+optional = false
+python-versions = ">=3.8"
+files = [
+    {file = "filelock-3.16.1-py3-none-any.whl", hash = "sha256:2082e5703d51fbf98ea75855d9d5527e33d8ff23099bec374a134febee6946b0"},
+    {file = "filelock-3.16.1.tar.gz", hash = "sha256:c249fbfcd5db47e5e2d6d62198e565475ee65e4831e2561c8e313fa7eb961435"},
+]
+
+[package.extras]
+docs = ["furo (>=2024.8.6)", "sphinx (>=8.0.2)", "sphinx-autodoc-typehints (>=2.4.1)"]
+testing = ["covdefaults (>=2.3)", "coverage (>=7.6.1)", "diff-cover (>=9.2)", "pytest (>=8.3.3)", "pytest-asyncio (>=0.24)", "pytest-cov (>=5)", "pytest-mock (>=3.14)", "pytest-timeout (>=2.3.1)", "virtualenv (>=20.26.4)"]
+typing = ["typing-extensions (>=4.12.2)"]
+
+[[package]]
+name = "fsspec"
+version = "2024.10.0"
+description = "File-system specification"
+optional = false
+python-versions = ">=3.8"
+files = [
+    {file = "fsspec-2024.10.0-py3-none-any.whl", hash = "sha256:03b9a6785766a4de40368b88906366755e2819e758b83705c88cd7cb5fe81871"},
+    {file = "fsspec-2024.10.0.tar.gz", hash = "sha256:eda2d8a4116d4f2429db8550f2457da57279247dd930bb12f821b58391359493"},
+]
+
+[package.extras]
+abfs = ["adlfs"]
+adl = ["adlfs"]
+arrow = ["pyarrow (>=1)"]
+dask = ["dask", "distributed"]
+dev = ["pre-commit", "ruff"]
+doc = ["numpydoc", "sphinx", "sphinx-design", "sphinx-rtd-theme", "yarl"]
+dropbox = ["dropbox", "dropboxdrivefs", "requests"]
+full = ["adlfs", "aiohttp (!=4.0.0a0,!=4.0.0a1)", "dask", "distributed", "dropbox", "dropboxdrivefs", "fusepy", "gcsfs", "libarchive-c", "ocifs", "panel", "paramiko", "pyarrow (>=1)", "pygit2", "requests", "s3fs", "smbprotocol", "tqdm"]
+fuse = ["fusepy"]
+gcs = ["gcsfs"]
+git = ["pygit2"]
+github = ["requests"]
+gs = ["gcsfs"]
+gui = ["panel"]
+hdfs = ["pyarrow (>=1)"]
+http = ["aiohttp (!=4.0.0a0,!=4.0.0a1)"]
+libarchive = ["libarchive-c"]
+oci = ["ocifs"]
+s3 = ["s3fs"]
+sftp = ["paramiko"]
+smb = ["smbprotocol"]
+ssh = ["paramiko"]
+test = ["aiohttp (!=4.0.0a0,!=4.0.0a1)", "numpy", "pytest", "pytest-asyncio (!=0.22.0)", "pytest-benchmark", "pytest-cov", "pytest-mock", "pytest-recording", "pytest-rerunfailures", "requests"]
+test-downstream = ["aiobotocore (>=2.5.4,<3.0.0)", "dask-expr", "dask[dataframe,test]", "moto[server] (>4,<5)", "pytest-timeout", "xarray"]
+test-full = ["adlfs", "aiohttp (!=4.0.0a0,!=4.0.0a1)", "cloudpickle", "dask", "distributed", "dropbox", "dropboxdrivefs", "fastparquet", "fusepy", "gcsfs", "jinja2", "kerchunk", "libarchive-c", "lz4", "notebook", "numpy", "ocifs", "pandas", "panel", "paramiko", "pyarrow", "pyarrow (>=1)", "pyftpdlib", "pygit2", "pytest", "pytest-asyncio (!=0.22.0)", "pytest-benchmark", "pytest-cov", "pytest-mock", "pytest-recording", "pytest-rerunfailures", "python-snappy", "requests", "smbprotocol", "tqdm", "urllib3", "zarr", "zstandard"]
+tqdm = ["tqdm"]
+
+[[package]]
+name = "jinja2"
+version = "3.1.4"
+description = "A very fast and expressive template engine."
+optional = false
+python-versions = ">=3.7"
+files = [
+    {file = "jinja2-3.1.4-py3-none-any.whl", hash = "sha256:bc5dd2abb727a5319567b7a813e6a2e7318c39f4f487cfe6c89c6f9c7d25197d"},
+    {file = "jinja2-3.1.4.tar.gz", hash = "sha256:4a3aee7acbbe7303aede8e9648d13b8bf88a429282aa6122a993f0ac800cb369"},
+]
+
+[package.dependencies]
+MarkupSafe = ">=2.0"
+
+[package.extras]
+i18n = ["Babel (>=2.7)"]
+
+[[package]]
+name = "markupsafe"
+version = "3.0.2"
+description = "Safely add untrusted strings to HTML/XML markup."
+optional = false
+python-versions = ">=3.9"
+files = [
+    {file = "MarkupSafe-3.0.2-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:7e94c425039cde14257288fd61dcfb01963e658efbc0ff54f5306b06054700f8"},
+    {file = "MarkupSafe-3.0.2-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:9e2d922824181480953426608b81967de705c3cef4d1af983af849d7bd619158"},
+    {file = "MarkupSafe-3.0.2-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:38a9ef736c01fccdd6600705b09dc574584b89bea478200c5fbf112a6b0d5579"},
+    {file = "MarkupSafe-3.0.2-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:bbcb445fa71794da8f178f0f6d66789a28d7319071af7a496d4d507ed566270d"},
+    {file = "MarkupSafe-3.0.2-cp310-cp310-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:57cb5a3cf367aeb1d316576250f65edec5bb3be939e9247ae594b4bcbc317dfb"},
+    {file = "MarkupSafe-3.0.2-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:3809ede931876f5b2ec92eef964286840ed3540dadf803dd570c3b7e13141a3b"},
+    {file = "MarkupSafe-3.0.2-cp310-cp310-musllinux_1_2_i686.whl", hash = "sha256:e07c3764494e3776c602c1e78e298937c3315ccc9043ead7e685b7f2b8d47b3c"},
+    {file = "MarkupSafe-3.0.2-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:b424c77b206d63d500bcb69fa55ed8d0e6a3774056bdc4839fc9298a7edca171"},
+    {file = "MarkupSafe-3.0.2-cp310-cp310-win32.whl", hash = "sha256:fcabf5ff6eea076f859677f5f0b6b5c1a51e70a376b0579e0eadef8db48c6b50"},
+    {file = "MarkupSafe-3.0.2-cp310-cp310-win_amd64.whl", hash = "sha256:6af100e168aa82a50e186c82875a5893c5597a0c1ccdb0d8b40240b1f28b969a"},
+    {file = "MarkupSafe-3.0.2-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:9025b4018f3a1314059769c7bf15441064b2207cb3f065e6ea1e7359cb46db9d"},
+    {file = "MarkupSafe-3.0.2-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:93335ca3812df2f366e80509ae119189886b0f3c2b81325d39efdb84a1e2ae93"},
+    {file = "MarkupSafe-3.0.2-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:2cb8438c3cbb25e220c2ab33bb226559e7afb3baec11c4f218ffa7308603c832"},
+    {file = "MarkupSafe-3.0.2-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:a123e330ef0853c6e822384873bef7507557d8e4a082961e1defa947aa59ba84"},
+    {file = "MarkupSafe-3.0.2-cp311-cp311-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:1e084f686b92e5b83186b07e8a17fc09e38fff551f3602b249881fec658d3eca"},
+    {file = "MarkupSafe-3.0.2-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:d8213e09c917a951de9d09ecee036d5c7d36cb6cb7dbaece4c71a60d79fb9798"},
+    {file = "MarkupSafe-3.0.2-cp311-cp311-musllinux_1_2_i686.whl", hash = "sha256:5b02fb34468b6aaa40dfc198d813a641e3a63b98c2b05a16b9f80b7ec314185e"},
+    {file = "MarkupSafe-3.0.2-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:0bff5e0ae4ef2e1ae4fdf2dfd5b76c75e5c2fa4132d05fc1b0dabcd20c7e28c4"},
+    {file = "MarkupSafe-3.0.2-cp311-cp311-win32.whl", hash = "sha256:6c89876f41da747c8d3677a2b540fb32ef5715f97b66eeb0c6b66f5e3ef6f59d"},
+    {file = "MarkupSafe-3.0.2-cp311-cp311-win_amd64.whl", hash = "sha256:70a87b411535ccad5ef2f1df5136506a10775d267e197e4cf531ced10537bd6b"},
+    {file = "MarkupSafe-3.0.2-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:9778bd8ab0a994ebf6f84c2b949e65736d5575320a17ae8984a77fab08db94cf"},
+    {file = "MarkupSafe-3.0.2-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:846ade7b71e3536c4e56b386c2a47adf5741d2d8b94ec9dc3e92e5e1ee1e2225"},
+    {file = "MarkupSafe-3.0.2-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:1c99d261bd2d5f6b59325c92c73df481e05e57f19837bdca8413b9eac4bd8028"},
+    {file = "MarkupSafe-3.0.2-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:e17c96c14e19278594aa4841ec148115f9c7615a47382ecb6b82bd8fea3ab0c8"},
+    {file = "MarkupSafe-3.0.2-cp312-cp312-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:88416bd1e65dcea10bc7569faacb2c20ce071dd1f87539ca2ab364bf6231393c"},
+    {file = "MarkupSafe-3.0.2-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:2181e67807fc2fa785d0592dc2d6206c019b9502410671cc905d132a92866557"},
+    {file = "MarkupSafe-3.0.2-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:52305740fe773d09cffb16f8ed0427942901f00adedac82ec8b67752f58a1b22"},
+    {file = "MarkupSafe-3.0.2-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:ad10d3ded218f1039f11a75f8091880239651b52e9bb592ca27de44eed242a48"},
+    {file = "MarkupSafe-3.0.2-cp312-cp312-win32.whl", hash = "sha256:0f4ca02bea9a23221c0182836703cbf8930c5e9454bacce27e767509fa286a30"},
+    {file = "MarkupSafe-3.0.2-cp312-cp312-win_amd64.whl", hash = "sha256:8e06879fc22a25ca47312fbe7c8264eb0b662f6db27cb2d3bbbc74b1df4b9b87"},
+    {file = "MarkupSafe-3.0.2-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:ba9527cdd4c926ed0760bc301f6728ef34d841f405abf9d4f959c478421e4efd"},
+    {file = "MarkupSafe-3.0.2-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:f8b3d067f2e40fe93e1ccdd6b2e1d16c43140e76f02fb1319a05cf2b79d99430"},
+    {file = "MarkupSafe-3.0.2-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:569511d3b58c8791ab4c2e1285575265991e6d8f8700c7be0e88f86cb0672094"},
+    {file = "MarkupSafe-3.0.2-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:15ab75ef81add55874e7ab7055e9c397312385bd9ced94920f2802310c930396"},
+    {file = "MarkupSafe-3.0.2-cp313-cp313-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:f3818cb119498c0678015754eba762e0d61e5b52d34c8b13d770f0719f7b1d79"},
+    {file = "MarkupSafe-3.0.2-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:cdb82a876c47801bb54a690c5ae105a46b392ac6099881cdfb9f6e95e4014c6a"},
+    {file = "MarkupSafe-3.0.2-cp313-cp313-musllinux_1_2_i686.whl", hash = "sha256:cabc348d87e913db6ab4aa100f01b08f481097838bdddf7c7a84b7575b7309ca"},
+    {file = "MarkupSafe-3.0.2-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:444dcda765c8a838eaae23112db52f1efaf750daddb2d9ca300bcae1039adc5c"},
+    {file = "MarkupSafe-3.0.2-cp313-cp313-win32.whl", hash = "sha256:bcf3e58998965654fdaff38e58584d8937aa3096ab5354d493c77d1fdd66d7a1"},
+    {file = "MarkupSafe-3.0.2-cp313-cp313-win_amd64.whl", hash = "sha256:e6a2a455bd412959b57a172ce6328d2dd1f01cb2135efda2e4576e8a23fa3b0f"},
+    {file = "MarkupSafe-3.0.2-cp313-cp313t-macosx_10_13_universal2.whl", hash = "sha256:b5a6b3ada725cea8a5e634536b1b01c30bcdcd7f9c6fff4151548d5bf6b3a36c"},
+    {file = "MarkupSafe-3.0.2-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:a904af0a6162c73e3edcb969eeeb53a63ceeb5d8cf642fade7d39e7963a22ddb"},
+    {file = "MarkupSafe-3.0.2-cp313-cp313t-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:4aa4e5faecf353ed117801a068ebab7b7e09ffb6e1d5e412dc852e0da018126c"},
+    {file = "MarkupSafe-3.0.2-cp313-cp313t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:c0ef13eaeee5b615fb07c9a7dadb38eac06a0608b41570d8ade51c56539e509d"},
+    {file = "MarkupSafe-3.0.2-cp313-cp313t-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:d16a81a06776313e817c951135cf7340a3e91e8c1ff2fac444cfd75fffa04afe"},
+    {file = "MarkupSafe-3.0.2-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:6381026f158fdb7c72a168278597a5e3a5222e83ea18f543112b2662a9b699c5"},
+    {file = "MarkupSafe-3.0.2-cp313-cp313t-musllinux_1_2_i686.whl", hash = "sha256:3d79d162e7be8f996986c064d1c7c817f6df3a77fe3d6859f6f9e7be4b8c213a"},
+    {file = "MarkupSafe-3.0.2-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:131a3c7689c85f5ad20f9f6fb1b866f402c445b220c19fe4308c0b147ccd2ad9"},
+    {file = "MarkupSafe-3.0.2-cp313-cp313t-win32.whl", hash = "sha256:ba8062ed2cf21c07a9e295d5b8a2a5ce678b913b45fdf68c32d95d6c1291e0b6"},
+    {file = "MarkupSafe-3.0.2-cp313-cp313t-win_amd64.whl", hash = "sha256:e444a31f8db13eb18ada366ab3cf45fd4b31e4db1236a4448f68778c1d1a5a2f"},
+    {file = "MarkupSafe-3.0.2-cp39-cp39-macosx_10_9_universal2.whl", hash = "sha256:eaa0a10b7f72326f1372a713e73c3f739b524b3af41feb43e4921cb529f5929a"},
+    {file = "MarkupSafe-3.0.2-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:48032821bbdf20f5799ff537c7ac3d1fba0ba032cfc06194faffa8cda8b560ff"},
+    {file = "MarkupSafe-3.0.2-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:1a9d3f5f0901fdec14d8d2f66ef7d035f2157240a433441719ac9a3fba440b13"},
+    {file = "MarkupSafe-3.0.2-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:88b49a3b9ff31e19998750c38e030fc7bb937398b1f78cfa599aaef92d693144"},
+    {file = "MarkupSafe-3.0.2-cp39-cp39-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:cfad01eed2c2e0c01fd0ecd2ef42c492f7f93902e39a42fc9ee1692961443a29"},
+    {file = "MarkupSafe-3.0.2-cp39-cp39-musllinux_1_2_aarch64.whl", hash = "sha256:1225beacc926f536dc82e45f8a4d68502949dc67eea90eab715dea3a21c1b5f0"},
+    {file = "MarkupSafe-3.0.2-cp39-cp39-musllinux_1_2_i686.whl", hash = "sha256:3169b1eefae027567d1ce6ee7cae382c57fe26e82775f460f0b2778beaad66c0"},
+    {file = "MarkupSafe-3.0.2-cp39-cp39-musllinux_1_2_x86_64.whl", hash = "sha256:eb7972a85c54febfb25b5c4b4f3af4dcc731994c7da0d8a0b4a6eb0640e1d178"},
+    {file = "MarkupSafe-3.0.2-cp39-cp39-win32.whl", hash = "sha256:8c4e8c3ce11e1f92f6536ff07154f9d49677ebaaafc32db9db4620bc11ed480f"},
+    {file = "MarkupSafe-3.0.2-cp39-cp39-win_amd64.whl", hash = "sha256:6e296a513ca3d94054c2c881cc913116e90fd030ad1c656b3869762b754f5f8a"},
+    {file = "markupsafe-3.0.2.tar.gz", hash = "sha256:ee55d3edf80167e48ea11a923c7386f4669df67d7994554387f84e7d8b0a2bf0"},
+]
+
+[[package]]
+name = "mpmath"
+version = "1.3.0"
+description = "Python library for arbitrary-precision floating-point arithmetic"
+optional = false
+python-versions = "*"
+files = [
+    {file = "mpmath-1.3.0-py3-none-any.whl", hash = "sha256:a0b2b9fe80bbcd81a6647ff13108738cfb482d481d826cc0e02f5b35e5c88d2c"},
+    {file = "mpmath-1.3.0.tar.gz", hash = "sha256:7a28eb2a9774d00c7bc92411c19a89209d5da7c4c9a9e227be8330a23a25b91f"},
+]
+
+[package.extras]
+develop = ["codecov", "pycodestyle", "pytest (>=4.6)", "pytest-cov", "wheel"]
+docs = ["sphinx"]
+gmpy = ["gmpy2 (>=2.1.0a4)"]
+tests = ["pytest (>=4.6)"]
+
+[[package]]
+name = "networkx"
+version = "3.4.2"
+description = "Python package for creating and manipulating graphs and networks"
+optional = false
+python-versions = ">=3.10"
+files = [
+    {file = "networkx-3.4.2-py3-none-any.whl", hash = "sha256:df5d4365b724cf81b8c6a7312509d0c22386097011ad1abe274afd5e9d3bbc5f"},
+    {file = "networkx-3.4.2.tar.gz", hash = "sha256:307c3669428c5362aab27c8a1260aa8f47c4e91d3891f48be0141738d8d053e1"},
+]
+
+[package.extras]
+default = ["matplotlib (>=3.7)", "numpy (>=1.24)", "pandas (>=2.0)", "scipy (>=1.10,!=1.11.0,!=1.11.1)"]
+developer = ["changelist (==0.5)", "mypy (>=1.1)", "pre-commit (>=3.2)", "rtoml"]
+doc = ["intersphinx-registry", "myst-nb (>=1.1)", "numpydoc (>=1.8.0)", "pillow (>=9.4)", "pydata-sphinx-theme (>=0.15)", "sphinx (>=7.3)", "sphinx-gallery (>=0.16)", "texext (>=0.6.7)"]
+example = ["cairocffi (>=1.7)", "contextily (>=1.6)", "igraph (>=0.11)", "momepy (>=0.7.2)", "osmnx (>=1.9)", "scikit-learn (>=1.5)", "seaborn (>=0.13)"]
+extra = ["lxml (>=4.6)", "pydot (>=3.0.1)", "pygraphviz (>=1.14)", "sympy (>=1.10)"]
+test = ["pytest (>=7.2)", "pytest-cov (>=4.0)"]
+
+[[package]]
+name = "numpy"
+version = "2.1.2"
+description = "Fundamental package for array computing in Python"
+optional = false
+python-versions = ">=3.10"
+files = [
+    {file = "numpy-2.1.2-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:30d53720b726ec36a7f88dc873f0eec8447fbc93d93a8f079dfac2629598d6ee"},
+    {file = "numpy-2.1.2-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:e8d3ca0a72dd8846eb6f7dfe8f19088060fcb76931ed592d29128e0219652884"},
+    {file = "numpy-2.1.2-cp310-cp310-macosx_14_0_arm64.whl", hash = "sha256:fc44e3c68ff00fd991b59092a54350e6e4911152682b4782f68070985aa9e648"},
+    {file = "numpy-2.1.2-cp310-cp310-macosx_14_0_x86_64.whl", hash = "sha256:7c1c60328bd964b53f8b835df69ae8198659e2b9302ff9ebb7de4e5a5994db3d"},
+    {file = "numpy-2.1.2-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:6cdb606a7478f9ad91c6283e238544451e3a95f30fb5467fbf715964341a8a86"},
+    {file = "numpy-2.1.2-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:d666cb72687559689e9906197e3bec7b736764df6a2e58ee265e360663e9baf7"},
+    {file = "numpy-2.1.2-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:c6eef7a2dbd0abfb0d9eaf78b73017dbfd0b54051102ff4e6a7b2980d5ac1a03"},
+    {file = "numpy-2.1.2-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:12edb90831ff481f7ef5f6bc6431a9d74dc0e5ff401559a71e5e4611d4f2d466"},
+    {file = "numpy-2.1.2-cp310-cp310-win32.whl", hash = "sha256:a65acfdb9c6ebb8368490dbafe83c03c7e277b37e6857f0caeadbbc56e12f4fb"},
+    {file = "numpy-2.1.2-cp310-cp310-win_amd64.whl", hash = "sha256:860ec6e63e2c5c2ee5e9121808145c7bf86c96cca9ad396c0bd3e0f2798ccbe2"},
+    {file = "numpy-2.1.2-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:b42a1a511c81cc78cbc4539675713bbcf9d9c3913386243ceff0e9429ca892fe"},
+    {file = "numpy-2.1.2-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:faa88bc527d0f097abdc2c663cddf37c05a1c2f113716601555249805cf573f1"},
+    {file = "numpy-2.1.2-cp311-cp311-macosx_14_0_arm64.whl", hash = "sha256:c82af4b2ddd2ee72d1fc0c6695048d457e00b3582ccde72d8a1c991b808bb20f"},
+    {file = "numpy-2.1.2-cp311-cp311-macosx_14_0_x86_64.whl", hash = "sha256:13602b3174432a35b16c4cfb5de9a12d229727c3dd47a6ce35111f2ebdf66ff4"},
+    {file = "numpy-2.1.2-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:1ebec5fd716c5a5b3d8dfcc439be82a8407b7b24b230d0ad28a81b61c2f4659a"},
+    {file = "numpy-2.1.2-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:e2b49c3c0804e8ecb05d59af8386ec2f74877f7ca8fd9c1e00be2672e4d399b1"},
+    {file = "numpy-2.1.2-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:2cbba4b30bf31ddbe97f1c7205ef976909a93a66bb1583e983adbd155ba72ac2"},
+    {file = "numpy-2.1.2-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:8e00ea6fc82e8a804433d3e9cedaa1051a1422cb6e443011590c14d2dea59146"},
+    {file = "numpy-2.1.2-cp311-cp311-win32.whl", hash = "sha256:5006b13a06e0b38d561fab5ccc37581f23c9511879be7693bd33c7cd15ca227c"},
+    {file = "numpy-2.1.2-cp311-cp311-win_amd64.whl", hash = "sha256:f1eb068ead09f4994dec71c24b2844f1e4e4e013b9629f812f292f04bd1510d9"},
+    {file = "numpy-2.1.2-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:d7bf0a4f9f15b32b5ba53147369e94296f5fffb783db5aacc1be15b4bf72f43b"},
+    {file = "numpy-2.1.2-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:b1d0fcae4f0949f215d4632be684a539859b295e2d0cb14f78ec231915d644db"},
+    {file = "numpy-2.1.2-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:f751ed0a2f250541e19dfca9f1eafa31a392c71c832b6bb9e113b10d050cb0f1"},
+    {file = "numpy-2.1.2-cp312-cp312-macosx_14_0_x86_64.whl", hash = "sha256:bd33f82e95ba7ad632bc57837ee99dba3d7e006536200c4e9124089e1bf42426"},
+    {file = "numpy-2.1.2-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:1b8cde4f11f0a975d1fd59373b32e2f5a562ade7cde4f85b7137f3de8fbb29a0"},
+    {file = "numpy-2.1.2-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6d95f286b8244b3649b477ac066c6906fbb2905f8ac19b170e2175d3d799f4df"},
+    {file = "numpy-2.1.2-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:ab4754d432e3ac42d33a269c8567413bdb541689b02d93788af4131018cbf366"},
+    {file = "numpy-2.1.2-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:e585c8ae871fd38ac50598f4763d73ec5497b0de9a0ab4ef5b69f01c6a046142"},
+    {file = "numpy-2.1.2-cp312-cp312-win32.whl", hash = "sha256:9c6c754df29ce6a89ed23afb25550d1c2d5fdb9901d9c67a16e0b16eaf7e2550"},
+    {file = "numpy-2.1.2-cp312-cp312-win_amd64.whl", hash = "sha256:456e3b11cb79ac9946c822a56346ec80275eaf2950314b249b512896c0d2505e"},
+    {file = "numpy-2.1.2-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:a84498e0d0a1174f2b3ed769b67b656aa5460c92c9554039e11f20a05650f00d"},
+    {file = "numpy-2.1.2-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:4d6ec0d4222e8ffdab1744da2560f07856421b367928026fb540e1945f2eeeaf"},
+    {file = "numpy-2.1.2-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:259ec80d54999cc34cd1eb8ded513cb053c3bf4829152a2e00de2371bd406f5e"},
+    {file = "numpy-2.1.2-cp313-cp313-macosx_14_0_x86_64.whl", hash = "sha256:675c741d4739af2dc20cd6c6a5c4b7355c728167845e3c6b0e824e4e5d36a6c3"},
+    {file = "numpy-2.1.2-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:05b2d4e667895cc55e3ff2b56077e4c8a5604361fc21a042845ea3ad67465aa8"},
+    {file = "numpy-2.1.2-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:43cca367bf94a14aca50b89e9bc2061683116cfe864e56740e083392f533ce7a"},
+    {file = "numpy-2.1.2-cp313-cp313-musllinux_1_1_x86_64.whl", hash = "sha256:76322dcdb16fccf2ac56f99048af32259dcc488d9b7e25b51e5eca5147a3fb98"},
+    {file = "numpy-2.1.2-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:32e16a03138cabe0cb28e1007ee82264296ac0983714094380b408097a418cfe"},
+    {file = "numpy-2.1.2-cp313-cp313-win32.whl", hash = "sha256:242b39d00e4944431a3cd2db2f5377e15b5785920421993770cddb89992c3f3a"},
+    {file = "numpy-2.1.2-cp313-cp313-win_amd64.whl", hash = "sha256:f2ded8d9b6f68cc26f8425eda5d3877b47343e68ca23d0d0846f4d312ecaa445"},
+    {file = "numpy-2.1.2-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:2ffef621c14ebb0188a8633348504a35c13680d6da93ab5cb86f4e54b7e922b5"},
+    {file = "numpy-2.1.2-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:ad369ed238b1959dfbade9018a740fb9392c5ac4f9b5173f420bd4f37ba1f7a0"},
+    {file = "numpy-2.1.2-cp313-cp313t-macosx_14_0_arm64.whl", hash = "sha256:d82075752f40c0ddf57e6e02673a17f6cb0f8eb3f587f63ca1eaab5594da5b17"},
+    {file = "numpy-2.1.2-cp313-cp313t-macosx_14_0_x86_64.whl", hash = "sha256:1600068c262af1ca9580a527d43dc9d959b0b1d8e56f8a05d830eea39b7c8af6"},
+    {file = "numpy-2.1.2-cp313-cp313t-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a26ae94658d3ba3781d5e103ac07a876b3e9b29db53f68ed7df432fd033358a8"},
+    {file = "numpy-2.1.2-cp313-cp313t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:13311c2db4c5f7609b462bc0f43d3c465424d25c626d95040f073e30f7570e35"},
+    {file = "numpy-2.1.2-cp313-cp313t-musllinux_1_1_x86_64.whl", hash = "sha256:2abbf905a0b568706391ec6fa15161fad0fb5d8b68d73c461b3c1bab6064dd62"},
+    {file = "numpy-2.1.2-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:ef444c57d664d35cac4e18c298c47d7b504c66b17c2ea91312e979fcfbdfb08a"},
+    {file = "numpy-2.1.2-pp310-pypy310_pp73-macosx_10_15_x86_64.whl", hash = "sha256:bdd407c40483463898b84490770199d5714dcc9dd9b792f6c6caccc523c00952"},
+    {file = "numpy-2.1.2-pp310-pypy310_pp73-macosx_14_0_x86_64.whl", hash = "sha256:da65fb46d4cbb75cb417cddf6ba5e7582eb7bb0b47db4b99c9fe5787ce5d91f5"},
+    {file = "numpy-2.1.2-pp310-pypy310_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:1c193d0b0238638e6fc5f10f1b074a6993cb13b0b431f64079a509d63d3aa8b7"},
+    {file = "numpy-2.1.2-pp310-pypy310_pp73-win_amd64.whl", hash = "sha256:a7d80b2e904faa63068ead63107189164ca443b42dd1930299e0d1cb041cec2e"},
+    {file = "numpy-2.1.2.tar.gz", hash = "sha256:13532a088217fa624c99b843eeb54640de23b3414b14aa66d023805eb731066c"},
+]
+
+[[package]]
+name = "nvidia-cublas-cu12"
+version = "12.4.5.8"
+description = "CUBLAS native runtime libraries"
+optional = false
+python-versions = ">=3"
+files = [
+    {file = "nvidia_cublas_cu12-12.4.5.8-py3-none-manylinux2014_aarch64.whl", hash = "sha256:0f8aa1706812e00b9f19dfe0cdb3999b092ccb8ca168c0db5b8ea712456fd9b3"},
+    {file = "nvidia_cublas_cu12-12.4.5.8-py3-none-manylinux2014_x86_64.whl", hash = "sha256:2fc8da60df463fdefa81e323eef2e36489e1c94335b5358bcb38360adf75ac9b"},
+    {file = "nvidia_cublas_cu12-12.4.5.8-py3-none-win_amd64.whl", hash = "sha256:5a796786da89203a0657eda402bcdcec6180254a8ac22d72213abc42069522dc"},
+]
+
+[[package]]
+name = "nvidia-cuda-cupti-cu12"
+version = "12.4.127"
+description = "CUDA profiling tools runtime libs."
+optional = false
+python-versions = ">=3"
+files = [
+    {file = "nvidia_cuda_cupti_cu12-12.4.127-py3-none-manylinux2014_aarch64.whl", hash = "sha256:79279b35cf6f91da114182a5ce1864997fd52294a87a16179ce275773799458a"},
+    {file = "nvidia_cuda_cupti_cu12-12.4.127-py3-none-manylinux2014_x86_64.whl", hash = "sha256:9dec60f5ac126f7bb551c055072b69d85392b13311fcc1bcda2202d172df30fb"},
+    {file = "nvidia_cuda_cupti_cu12-12.4.127-py3-none-win_amd64.whl", hash = "sha256:5688d203301ab051449a2b1cb6690fbe90d2b372f411521c86018b950f3d7922"},
+]
+
+[[package]]
+name = "nvidia-cuda-nvrtc-cu12"
+version = "12.4.127"
+description = "NVRTC native runtime libraries"
+optional = false
+python-versions = ">=3"
+files = [
+    {file = "nvidia_cuda_nvrtc_cu12-12.4.127-py3-none-manylinux2014_aarch64.whl", hash = "sha256:0eedf14185e04b76aa05b1fea04133e59f465b6f960c0cbf4e37c3cb6b0ea198"},
+    {file = "nvidia_cuda_nvrtc_cu12-12.4.127-py3-none-manylinux2014_x86_64.whl", hash = "sha256:a178759ebb095827bd30ef56598ec182b85547f1508941a3d560eb7ea1fbf338"},
+    {file = "nvidia_cuda_nvrtc_cu12-12.4.127-py3-none-win_amd64.whl", hash = "sha256:a961b2f1d5f17b14867c619ceb99ef6fcec12e46612711bcec78eb05068a60ec"},
+]
+
+[[package]]
+name = "nvidia-cuda-runtime-cu12"
+version = "12.4.127"
+description = "CUDA Runtime native Libraries"
+optional = false
+python-versions = ">=3"
+files = [
+    {file = "nvidia_cuda_runtime_cu12-12.4.127-py3-none-manylinux2014_aarch64.whl", hash = "sha256:961fe0e2e716a2a1d967aab7caee97512f71767f852f67432d572e36cb3a11f3"},
+    {file = "nvidia_cuda_runtime_cu12-12.4.127-py3-none-manylinux2014_x86_64.whl", hash = "sha256:64403288fa2136ee8e467cdc9c9427e0434110899d07c779f25b5c068934faa5"},
+    {file = "nvidia_cuda_runtime_cu12-12.4.127-py3-none-win_amd64.whl", hash = "sha256:09c2e35f48359752dfa822c09918211844a3d93c100a715d79b59591130c5e1e"},
+]
+
+[[package]]
+name = "nvidia-cudnn-cu12"
+version = "9.1.0.70"
+description = "cuDNN runtime libraries"
+optional = false
+python-versions = ">=3"
+files = [
+    {file = "nvidia_cudnn_cu12-9.1.0.70-py3-none-manylinux2014_x86_64.whl", hash = "sha256:165764f44ef8c61fcdfdfdbe769d687e06374059fbb388b6c89ecb0e28793a6f"},
+    {file = "nvidia_cudnn_cu12-9.1.0.70-py3-none-win_amd64.whl", hash = "sha256:6278562929433d68365a07a4a1546c237ba2849852c0d4b2262a486e805b977a"},
+]
+
+[package.dependencies]
+nvidia-cublas-cu12 = "*"
+
+[[package]]
+name = "nvidia-cufft-cu12"
+version = "11.2.1.3"
+description = "CUFFT native runtime libraries"
+optional = false
+python-versions = ">=3"
+files = [
+    {file = "nvidia_cufft_cu12-11.2.1.3-py3-none-manylinux2014_aarch64.whl", hash = "sha256:5dad8008fc7f92f5ddfa2101430917ce2ffacd86824914c82e28990ad7f00399"},
+    {file = "nvidia_cufft_cu12-11.2.1.3-py3-none-manylinux2014_x86_64.whl", hash = "sha256:f083fc24912aa410be21fa16d157fed2055dab1cc4b6934a0e03cba69eb242b9"},
+    {file = "nvidia_cufft_cu12-11.2.1.3-py3-none-win_amd64.whl", hash = "sha256:d802f4954291101186078ccbe22fc285a902136f974d369540fd4a5333d1440b"},
+]
+
+[package.dependencies]
+nvidia-nvjitlink-cu12 = "*"
+
+[[package]]
+name = "nvidia-curand-cu12"
+version = "10.3.5.147"
+description = "CURAND native runtime libraries"
+optional = false
+python-versions = ">=3"
+files = [
+    {file = "nvidia_curand_cu12-10.3.5.147-py3-none-manylinux2014_aarch64.whl", hash = "sha256:1f173f09e3e3c76ab084aba0de819c49e56614feae5c12f69883f4ae9bb5fad9"},
+    {file = "nvidia_curand_cu12-10.3.5.147-py3-none-manylinux2014_x86_64.whl", hash = "sha256:a88f583d4e0bb643c49743469964103aa59f7f708d862c3ddb0fc07f851e3b8b"},
+    {file = "nvidia_curand_cu12-10.3.5.147-py3-none-win_amd64.whl", hash = "sha256:f307cc191f96efe9e8f05a87096abc20d08845a841889ef78cb06924437f6771"},
+]
+
+[[package]]
+name = "nvidia-cusolver-cu12"
+version = "11.6.1.9"
+description = "CUDA solver native runtime libraries"
+optional = false
+python-versions = ">=3"
+files = [
+    {file = "nvidia_cusolver_cu12-11.6.1.9-py3-none-manylinux2014_aarch64.whl", hash = "sha256:d338f155f174f90724bbde3758b7ac375a70ce8e706d70b018dd3375545fc84e"},
+    {file = "nvidia_cusolver_cu12-11.6.1.9-py3-none-manylinux2014_x86_64.whl", hash = "sha256:19e33fa442bcfd085b3086c4ebf7e8debc07cfe01e11513cc6d332fd918ac260"},
+    {file = "nvidia_cusolver_cu12-11.6.1.9-py3-none-win_amd64.whl", hash = "sha256:e77314c9d7b694fcebc84f58989f3aa4fb4cb442f12ca1a9bde50f5e8f6d1b9c"},
+]
+
+[package.dependencies]
+nvidia-cublas-cu12 = "*"
+nvidia-cusparse-cu12 = "*"
+nvidia-nvjitlink-cu12 = "*"
+
+[[package]]
+name = "nvidia-cusparse-cu12"
+version = "12.3.1.170"
+description = "CUSPARSE native runtime libraries"
+optional = false
+python-versions = ">=3"
+files = [
+    {file = "nvidia_cusparse_cu12-12.3.1.170-py3-none-manylinux2014_aarch64.whl", hash = "sha256:9d32f62896231ebe0480efd8a7f702e143c98cfaa0e8a76df3386c1ba2b54df3"},
+    {file = "nvidia_cusparse_cu12-12.3.1.170-py3-none-manylinux2014_x86_64.whl", hash = "sha256:ea4f11a2904e2a8dc4b1833cc1b5181cde564edd0d5cd33e3c168eff2d1863f1"},
+    {file = "nvidia_cusparse_cu12-12.3.1.170-py3-none-win_amd64.whl", hash = "sha256:9bc90fb087bc7b4c15641521f31c0371e9a612fc2ba12c338d3ae032e6b6797f"},
+]
+
+[package.dependencies]
+nvidia-nvjitlink-cu12 = "*"
+
+[[package]]
+name = "nvidia-nccl-cu12"
+version = "2.21.5"
+description = "NVIDIA Collective Communication Library (NCCL) Runtime"
+optional = false
+python-versions = ">=3"
+files = [
+    {file = "nvidia_nccl_cu12-2.21.5-py3-none-manylinux2014_x86_64.whl", hash = "sha256:8579076d30a8c24988834445f8d633c697d42397e92ffc3f63fa26766d25e0a0"},
+]
+
+[[package]]
+name = "nvidia-nvjitlink-cu12"
+version = "12.4.127"
+description = "Nvidia JIT LTO Library"
+optional = false
+python-versions = ">=3"
+files = [
+    {file = "nvidia_nvjitlink_cu12-12.4.127-py3-none-manylinux2014_x86_64.whl", hash = "sha256:06b3b9b25bf3f8af351d664978ca26a16d2c5127dbd53c0497e28d1fb9611d57"},
+    {file = "nvidia_nvjitlink_cu12-12.4.127-py3-none-win_amd64.whl", hash = "sha256:fd9020c501d27d135f983c6d3e244b197a7ccad769e34df53a42e276b0e25fa1"},
+]
+
+[[package]]
+name = "nvidia-nvtx-cu12"
+version = "12.4.127"
+description = "NVIDIA Tools Extension"
+optional = false
+python-versions = ">=3"
+files = [
+    {file = "nvidia_nvtx_cu12-12.4.127-py3-none-manylinux2014_aarch64.whl", hash = "sha256:7959ad635db13edf4fc65c06a6e9f9e55fc2f92596db928d169c0bb031e88ef3"},
+    {file = "nvidia_nvtx_cu12-12.4.127-py3-none-manylinux2014_x86_64.whl", hash = "sha256:781e950d9b9f60d8241ccea575b32f5105a5baf4c2351cab5256a24869f12a1a"},
+    {file = "nvidia_nvtx_cu12-12.4.127-py3-none-win_amd64.whl", hash = "sha256:641dccaaa1139f3ffb0d3164b4b84f9d253397e38246a4f2f36728b48566d485"},
+]
+
+[[package]]
+name = "scipy"
+version = "1.14.1"
+description = "Fundamental algorithms for scientific computing in Python"
+optional = false
+python-versions = ">=3.10"
+files = [
+    {file = "scipy-1.14.1-cp310-cp310-macosx_10_13_x86_64.whl", hash = "sha256:b28d2ca4add7ac16ae8bb6632a3c86e4b9e4d52d3e34267f6e1b0c1f8d87e389"},
+    {file = "scipy-1.14.1-cp310-cp310-macosx_12_0_arm64.whl", hash = "sha256:d0d2821003174de06b69e58cef2316a6622b60ee613121199cb2852a873f8cf3"},
+    {file = "scipy-1.14.1-cp310-cp310-macosx_14_0_arm64.whl", hash = "sha256:8bddf15838ba768bb5f5083c1ea012d64c9a444e16192762bd858f1e126196d0"},
+    {file = "scipy-1.14.1-cp310-cp310-macosx_14_0_x86_64.whl", hash = "sha256:97c5dddd5932bd2a1a31c927ba5e1463a53b87ca96b5c9bdf5dfd6096e27efc3"},
+    {file = "scipy-1.14.1-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:2ff0a7e01e422c15739ecd64432743cf7aae2b03f3084288f399affcefe5222d"},
+    {file = "scipy-1.14.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:8e32dced201274bf96899e6491d9ba3e9a5f6b336708656466ad0522d8528f69"},
+    {file = "scipy-1.14.1-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:8426251ad1e4ad903a4514712d2fa8fdd5382c978010d1c6f5f37ef286a713ad"},
+    {file = "scipy-1.14.1-cp310-cp310-win_amd64.whl", hash = "sha256:a49f6ed96f83966f576b33a44257d869756df6cf1ef4934f59dd58b25e0327e5"},
+    {file = "scipy-1.14.1-cp311-cp311-macosx_10_13_x86_64.whl", hash = "sha256:2da0469a4ef0ecd3693761acbdc20f2fdeafb69e6819cc081308cc978153c675"},
+    {file = "scipy-1.14.1-cp311-cp311-macosx_12_0_arm64.whl", hash = "sha256:c0ee987efa6737242745f347835da2cc5bb9f1b42996a4d97d5c7ff7928cb6f2"},
+    {file = "scipy-1.14.1-cp311-cp311-macosx_14_0_arm64.whl", hash = "sha256:3a1b111fac6baec1c1d92f27e76511c9e7218f1695d61b59e05e0fe04dc59617"},
+    {file = "scipy-1.14.1-cp311-cp311-macosx_14_0_x86_64.whl", hash = "sha256:8475230e55549ab3f207bff11ebfc91c805dc3463ef62eda3ccf593254524ce8"},
+    {file = "scipy-1.14.1-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:278266012eb69f4a720827bdd2dc54b2271c97d84255b2faaa8f161a158c3b37"},
+    {file = "scipy-1.14.1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:fef8c87f8abfb884dac04e97824b61299880c43f4ce675dd2cbeadd3c9b466d2"},
+    {file = "scipy-1.14.1-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:b05d43735bb2f07d689f56f7b474788a13ed8adc484a85aa65c0fd931cf9ccd2"},
+    {file = "scipy-1.14.1-cp311-cp311-win_amd64.whl", hash = "sha256:716e389b694c4bb564b4fc0c51bc84d381735e0d39d3f26ec1af2556ec6aad94"},
+    {file = "scipy-1.14.1-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:631f07b3734d34aced009aaf6fedfd0eb3498a97e581c3b1e5f14a04164a456d"},
+    {file = "scipy-1.14.1-cp312-cp312-macosx_12_0_arm64.whl", hash = "sha256:af29a935803cc707ab2ed7791c44288a682f9c8107bc00f0eccc4f92c08d6e07"},
+    {file = "scipy-1.14.1-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:2843f2d527d9eebec9a43e6b406fb7266f3af25a751aa91d62ff416f54170bc5"},
+    {file = "scipy-1.14.1-cp312-cp312-macosx_14_0_x86_64.whl", hash = "sha256:eb58ca0abd96911932f688528977858681a59d61a7ce908ffd355957f7025cfc"},
+    {file = "scipy-1.14.1-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:30ac8812c1d2aab7131a79ba62933a2a76f582d5dbbc695192453dae67ad6310"},
+    {file = "scipy-1.14.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:8f9ea80f2e65bdaa0b7627fb00cbeb2daf163caa015e59b7516395fe3bd1e066"},
+    {file = "scipy-1.14.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:edaf02b82cd7639db00dbff629995ef185c8df4c3ffa71a5562a595765a06ce1"},
+    {file = "scipy-1.14.1-cp312-cp312-win_amd64.whl", hash = "sha256:2ff38e22128e6c03ff73b6bb0f85f897d2362f8c052e3b8ad00532198fbdae3f"},
+    {file = "scipy-1.14.1-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:1729560c906963fc8389f6aac023739ff3983e727b1a4d87696b7bf108316a79"},
+    {file = "scipy-1.14.1-cp313-cp313-macosx_12_0_arm64.whl", hash = "sha256:4079b90df244709e675cdc8b93bfd8a395d59af40b72e339c2287c91860deb8e"},
+    {file = "scipy-1.14.1-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:e0cf28db0f24a38b2a0ca33a85a54852586e43cf6fd876365c86e0657cfe7d73"},
+    {file = "scipy-1.14.1-cp313-cp313-macosx_14_0_x86_64.whl", hash = "sha256:0c2f95de3b04e26f5f3ad5bb05e74ba7f68b837133a4492414b3afd79dfe540e"},
+    {file = "scipy-1.14.1-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:b99722ea48b7ea25e8e015e8341ae74624f72e5f21fc2abd45f3a93266de4c5d"},
+    {file = "scipy-1.14.1-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:5149e3fd2d686e42144a093b206aef01932a0059c2a33ddfa67f5f035bdfe13e"},
+    {file = "scipy-1.14.1-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:e4f5a7c49323533f9103d4dacf4e4f07078f360743dec7f7596949149efeec06"},
+    {file = "scipy-1.14.1-cp313-cp313-win_amd64.whl", hash = "sha256:baff393942b550823bfce952bb62270ee17504d02a1801d7fd0719534dfb9c84"},
+    {file = "scipy-1.14.1.tar.gz", hash = "sha256:5a275584e726026a5699459aa72f828a610821006228e841b94275c4a7c08417"},
+]
+
+[package.dependencies]
+numpy = ">=1.23.5,<2.3"
+
+[package.extras]
+dev = ["cython-lint (>=0.12.2)", "doit (>=0.36.0)", "mypy (==1.10.0)", "pycodestyle", "pydevtool", "rich-click", "ruff (>=0.0.292)", "types-psutil", "typing_extensions"]
+doc = ["jupyterlite-pyodide-kernel", "jupyterlite-sphinx (>=0.13.1)", "jupytext", "matplotlib (>=3.5)", "myst-nb", "numpydoc", "pooch", "pydata-sphinx-theme (>=0.15.2)", "sphinx (>=5.0.0,<=7.3.7)", "sphinx-design (>=0.4.0)"]
+test = ["Cython", "array-api-strict (>=2.0)", "asv", "gmpy2", "hypothesis (>=6.30)", "meson", "mpmath", "ninja", "pooch", "pytest", "pytest-cov", "pytest-timeout", "pytest-xdist", "scikit-umfpack", "threadpoolctl"]
+
+[[package]]
+name = "setuptools"
+version = "75.2.0"
+description = "Easily download, build, install, upgrade, and uninstall Python packages"
+optional = false
+python-versions = ">=3.8"
+files = [
+    {file = "setuptools-75.2.0-py3-none-any.whl", hash = "sha256:a7fcb66f68b4d9e8e66b42f9876150a3371558f98fa32222ffaa5bced76406f8"},
+    {file = "setuptools-75.2.0.tar.gz", hash = "sha256:753bb6ebf1f465a1912e19ed1d41f403a79173a9acf66a42e7e6aec45c3c16ec"},
+]
+
+[package.extras]
+check = ["pytest-checkdocs (>=2.4)", "pytest-ruff (>=0.2.1)", "ruff (>=0.5.2)"]
+core = ["importlib-metadata (>=6)", "importlib-resources (>=5.10.2)", "jaraco.collections", "jaraco.functools", "jaraco.text (>=3.7)", "more-itertools", "more-itertools (>=8.8)", "packaging", "packaging (>=24)", "platformdirs (>=2.6.2)", "tomli (>=2.0.1)", "wheel (>=0.43.0)"]
+cover = ["pytest-cov"]
+doc = ["furo", "jaraco.packaging (>=9.3)", "jaraco.tidelift (>=1.4)", "pygments-github-lexers (==0.0.5)", "pyproject-hooks (!=1.1)", "rst.linker (>=1.9)", "sphinx (>=3.5)", "sphinx-favicon", "sphinx-inline-tabs", "sphinx-lint", "sphinx-notfound-page (>=1,<2)", "sphinx-reredirects", "sphinxcontrib-towncrier", "towncrier (<24.7)"]
+enabler = ["pytest-enabler (>=2.2)"]
+test = ["build[virtualenv] (>=1.0.3)", "filelock (>=3.4.0)", "ini2toml[lite] (>=0.14)", "jaraco.develop (>=7.21)", "jaraco.envs (>=2.2)", "jaraco.path (>=3.2.0)", "jaraco.test", "packaging (>=23.2)", "pip (>=19.1)", "pyproject-hooks (!=1.1)", "pytest (>=6,!=8.1.*)", "pytest-home (>=0.5)", "pytest-perf", "pytest-subprocess", "pytest-timeout", "pytest-xdist (>=3)", "tomli-w (>=1.0.0)", "virtualenv (>=13.0.0)", "wheel (>=0.44.0)"]
+type = ["importlib-metadata (>=7.0.2)", "jaraco.develop (>=7.21)", "mypy (==1.11.*)", "pytest-mypy"]
+
+[[package]]
+name = "sympy"
+version = "1.13.1"
+description = "Computer algebra system (CAS) in Python"
+optional = false
+python-versions = ">=3.8"
+files = [
+    {file = "sympy-1.13.1-py3-none-any.whl", hash = "sha256:db36cdc64bf61b9b24578b6f7bab1ecdd2452cf008f34faa33776680c26d66f8"},
+    {file = "sympy-1.13.1.tar.gz", hash = "sha256:9cebf7e04ff162015ce31c9c6c9144daa34a93bd082f54fd8f12deca4f47515f"},
+]
+
+[package.dependencies]
+mpmath = ">=1.1.0,<1.4"
+
+[package.extras]
+dev = ["hypothesis (>=6.70.0)", "pytest (>=7.1.0)"]
+
+[[package]]
+name = "torch"
+version = "2.5.0"
+description = "Tensors and Dynamic neural networks in Python with strong GPU acceleration"
+optional = false
+python-versions = ">=3.8.0"
+files = [
+    {file = "torch-2.5.0-cp310-cp310-manylinux1_x86_64.whl", hash = "sha256:7f179373a047b947dec448243f4e6598a1c960fa3bb978a9a7eecd529fbc363f"},
+    {file = "torch-2.5.0-cp310-cp310-manylinux2014_aarch64.whl", hash = "sha256:15fbc95e38d330e5b0ef1593b7bc0a19f30e5bdad76895a5cffa1a6a044235e9"},
+    {file = "torch-2.5.0-cp310-cp310-win_amd64.whl", hash = "sha256:f499212f1cffea5d587e5f06144630ed9aa9c399bba12ec8905798d833bd1404"},
+    {file = "torch-2.5.0-cp310-none-macosx_11_0_arm64.whl", hash = "sha256:c54db1fade17287aabbeed685d8e8ab3a56fea9dd8d46e71ced2da367f09a49f"},
+    {file = "torch-2.5.0-cp311-cp311-manylinux1_x86_64.whl", hash = "sha256:499a68a756d3b30d10f7e0f6214dc3767b130b797265db3b1c02e9094e2a07be"},
+    {file = "torch-2.5.0-cp311-cp311-manylinux2014_aarch64.whl", hash = "sha256:9f3df8138a1126a851440b7d5a4869bfb7c9cc43563d64fd9d96d0465b581024"},
+    {file = "torch-2.5.0-cp311-cp311-win_amd64.whl", hash = "sha256:b81da3bdb58c9de29d0e1361e52f12fcf10a89673f17a11a5c6c7da1cb1a8376"},
+    {file = "torch-2.5.0-cp311-none-macosx_11_0_arm64.whl", hash = "sha256:ba135923295d564355326dc409b6b7f5bd6edc80f764cdaef1fb0a1b23ff2f9c"},
+    {file = "torch-2.5.0-cp312-cp312-manylinux1_x86_64.whl", hash = "sha256:2dd40c885a05ef7fe29356cca81be1435a893096ceb984441d6e2c27aff8c6f4"},
+    {file = "torch-2.5.0-cp312-cp312-manylinux2014_aarch64.whl", hash = "sha256:bc52d603d87fe1da24439c0d5fdbbb14e0ae4874451d53f0120ffb1f6c192727"},
+    {file = "torch-2.5.0-cp312-cp312-win_amd64.whl", hash = "sha256:ea718746469246cc63b3353afd75698a288344adb55e29b7f814a5d3c0a7c78d"},
+    {file = "torch-2.5.0-cp312-none-macosx_11_0_arm64.whl", hash = "sha256:6de1fd253e27e7f01f05cd7c37929ae521ca23ca4620cfc7c485299941679112"},
+    {file = "torch-2.5.0-cp313-cp313-manylinux1_x86_64.whl", hash = "sha256:83dcf518685db20912b71fc49cbddcc8849438cdb0e9dcc919b02a849e2cd9e8"},
+    {file = "torch-2.5.0-cp39-cp39-manylinux1_x86_64.whl", hash = "sha256:65e0a60894435608334d68c8811e55fd8f73e5bf8ee6f9ccedb0064486a7b418"},
+    {file = "torch-2.5.0-cp39-cp39-manylinux2014_aarch64.whl", hash = "sha256:38c21ff1bd39f076d72ab06e3c88c2ea6874f2e6f235c9450816b6c8e7627094"},
+    {file = "torch-2.5.0-cp39-cp39-win_amd64.whl", hash = "sha256:ce4baeba9804da5a346e210b3b70826f5811330c343e4fe1582200359ee77fe5"},
+    {file = "torch-2.5.0-cp39-none-macosx_11_0_arm64.whl", hash = "sha256:03e53f577a96e4d41aca472da8faa40e55df89d2273664af390ce1f570e885bd"},
+]
+
+[package.dependencies]
+filelock = "*"
+fsspec = "*"
+jinja2 = "*"
+networkx = "*"
+nvidia-cublas-cu12 = {version = "12.4.5.8", markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\""}
+nvidia-cuda-cupti-cu12 = {version = "12.4.127", markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\""}
+nvidia-cuda-nvrtc-cu12 = {version = "12.4.127", markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\""}
+nvidia-cuda-runtime-cu12 = {version = "12.4.127", markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\""}
+nvidia-cudnn-cu12 = {version = "9.1.0.70", markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\""}
+nvidia-cufft-cu12 = {version = "11.2.1.3", markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\""}
+nvidia-curand-cu12 = {version = "10.3.5.147", markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\""}
+nvidia-cusolver-cu12 = {version = "11.6.1.9", markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\""}
+nvidia-cusparse-cu12 = {version = "12.3.1.170", markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\""}
+nvidia-nccl-cu12 = {version = "2.21.5", markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\""}
+nvidia-nvjitlink-cu12 = {version = "12.4.127", markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\""}
+nvidia-nvtx-cu12 = {version = "12.4.127", markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\""}
+setuptools = {version = "*", markers = "python_version >= \"3.12\""}
+sympy = {version = "1.13.1", markers = "python_version >= \"3.9\""}
+triton = {version = "3.1.0", markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\" and python_version < \"3.13\""}
+typing-extensions = ">=4.8.0"
+
+[package.extras]
+opt-einsum = ["opt-einsum (>=3.3)"]
+optree = ["optree (>=0.12.0)"]
+
+[[package]]
+name = "tqdm"
+version = "4.66.5"
+description = "Fast, Extensible Progress Meter"
+optional = false
+python-versions = ">=3.7"
+files = [
+    {file = "tqdm-4.66.5-py3-none-any.whl", hash = "sha256:90279a3770753eafc9194a0364852159802111925aa30eb3f9d85b0e805ac7cd"},
+    {file = "tqdm-4.66.5.tar.gz", hash = "sha256:e1020aef2e5096702d8a025ac7d16b1577279c9d63f8375b63083e9a5f0fcbad"},
+]
+
+[package.dependencies]
+colorama = {version = "*", markers = "platform_system == \"Windows\""}
+
+[package.extras]
+dev = ["pytest (>=6)", "pytest-cov", "pytest-timeout", "pytest-xdist"]
+notebook = ["ipywidgets (>=6)"]
+slack = ["slack-sdk"]
+telegram = ["requests"]
+
+[[package]]
+name = "triton"
+version = "3.1.0"
+description = "A language and compiler for custom Deep Learning operations"
+optional = false
+python-versions = "*"
+files = [
+    {file = "triton-3.1.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6b0dd10a925263abbe9fa37dcde67a5e9b2383fc269fdf59f5657cac38c5d1d8"},
+    {file = "triton-3.1.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:0f34f6e7885d1bf0eaaf7ba875a5f0ce6f3c13ba98f9503651c1e6dc6757ed5c"},
+    {file = "triton-3.1.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:c8182f42fd8080a7d39d666814fa36c5e30cc00ea7eeeb1a2983dbb4c99a0fdc"},
+    {file = "triton-3.1.0-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6dadaca7fc24de34e180271b5cf864c16755702e9f63a16f62df714a8099126a"},
+    {file = "triton-3.1.0-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:aafa9a20cd0d9fee523cd4504aa7131807a864cd77dcf6efe7e981f18b8c6c11"},
+]
+
+[package.dependencies]
+filelock = "*"
+
+[package.extras]
+build = ["cmake (>=3.20)", "lit"]
+tests = ["autopep8", "flake8", "isort", "llnl-hatchet", "numpy", "pytest", "scipy (>=1.7.1)"]
+tutorials = ["matplotlib", "pandas", "tabulate"]
+
+[[package]]
+name = "typing-extensions"
+version = "4.12.2"
+description = "Backported and Experimental Type Hints for Python 3.8+"
+optional = false
+python-versions = ">=3.8"
+files = [
+    {file = "typing_extensions-4.12.2-py3-none-any.whl", hash = "sha256:04e5ca0351e0f3f85c6853954072df659d0d13fac324d0072316b67d7794700d"},
+    {file = "typing_extensions-4.12.2.tar.gz", hash = "sha256:1a7ead55c7e559dd4dee8856e3a88b41225abfe1ce8df57b7c13915fe121ffb8"},
+]
+
+[metadata]
+lock-version = "2.0"
+python-versions = "^3.11"
+content-hash = "fb3198ed8518f731a442b6755e2141800cf8c817f38d3d5f502e7dad3a8da9d9"
diff --git a/pyproject.toml b/pyproject.toml
new file mode 100644
index 0000000..3d1d727
--- /dev/null
+++ b/pyproject.toml
@@ -0,0 +1,18 @@
+[tool.poetry]
+name = "liar-dice"
+version = "0.1.0"
+description = ""
+authors = ["PabloLION <36828324+PabloLION@users.noreply.github.com>"]
+readme = "README.md"
+
+[tool.poetry.dependencies]
+python = "^3.11"
+torch = "^2.5.0"
+tqdm = "^4.66.5"
+numpy = "^2.1.2"
+scipy = "^1.14.1"
+
+
+[build-system]
+requires = ["poetry-core"]
+build-backend = "poetry.core.masonry.api"
diff --git a/tests/__init__.py b/tests/__init__.py
new file mode 100644
index 0000000..e69de29