added evaluation script

pyproject.toml

@@ -5,12 +5,20 @@ requires-python = ">=3.13"
 dependencies = [
     "equinox>=0.13.8",
     "generals-bots",
-    "jax[cuda]>=0.11.0",
+    "jax==0.11.0",
     "jaxtyping>=0.3.11",
     "omegaconf>=2.3.1",
     "optax>=0.2.8",
 ]
 
+[project.optional-dependencies]
+cuda = [
+    "jax[cuda]==0.11.0"
+]
+cuda13 = [
+    "jax[cuda13]==0.11.0"
+]
+
 [tool.uv.sources]
 generals-bots = { git = "https://github.com/strakam/generals-bots.git" }
 

requirements.txt

@@ -0,0 +1,81 @@
+# This file was autogenerated by uv via the following command:
+#    uv pip compile pyproject.toml -o requirements.txt
+absl-py==2.5.0
+    # via optax
+antlr4-python3-runtime==4.9.3
+    # via omegaconf
+bidict==0.23.1
+    # via python-socketio
+certifi==2026.7.22
+    # via requests
+charset-normalizer==3.4.9
+    # via requests
+equinox==0.13.8
+    # via general-bots-training (pyproject.toml)
+generals-bots @ git+https://github.com/strakam/generals-bots.git@9e3b9d13cca51caa1bb07db48bb85c9e90ce0462
+    # via general-bots-training (pyproject.toml)
+h11==0.16.0
+    # via wsproto
+idna==3.18
+    # via requests
+jax==0.11.0
+    # via
+    #   general-bots-training (pyproject.toml)
+    #   equinox
+    #   generals-bots
+    #   optax
+jaxlib==0.11.0
+    # via
+    #   generals-bots
+    #   jax
+    #   optax
+jaxtyping==0.3.11
+    # via
+    #   general-bots-training (pyproject.toml)
+    #   equinox
+ml-dtypes==0.5.4
+    # via
+    #   jax
+    #   jaxlib
+numpy==2.5.1
+    # via
+    #   generals-bots
+    #   jax
+    #   jaxlib
+    #   ml-dtypes
+    #   optax
+    #   scipy
+omegaconf==2.3.1
+    # via general-bots-training (pyproject.toml)
+opt-einsum==3.4.0
+    # via jax
+optax==0.2.8
+    # via general-bots-training (pyproject.toml)
+pygame==2.6.1
+    # via generals-bots
+python-engineio==4.13.4
+    # via python-socketio
+python-socketio==5.16.4
+    # via generals-bots
+pyyaml==6.0.3
+    # via omegaconf
+requests==2.34.2
+    # via python-socketio
+scipy==1.18.0
+    # via
+    #   jax
+    #   jaxlib
+simple-websocket==1.1.0
+    # via python-engineio
+typing-extensions==4.16.0
+    # via equinox
+urllib3==2.7.0
+    # via requests
+wadler-lindig==0.1.7
+    # via
+    #   equinox
+    #   jaxtyping
+websocket-client==1.9.0
+    # via python-socketio
+wsproto==1.3.2
+    # via simple-websocket

scripts/evaluate.py

@@ -0,0 +1,173 @@
+"""Evaluate Generals.io agents by playing them against each other.
+
+Run with `.venv/bin/python scripts/evaluate.py` or `uv run python scripts/evaluate.py`.
+Override defaults with OmegaConf arguments such as `num_games=200 agent0.kind=hunter`.
+
+Each agent is either one of the bundled JAX agents (`random`, `expander`,
+`hunter`) or a trained policy loaded from an equinox checkpoint via
+`agent0.kind=model agent0.checkpoint=ppo_model.eqx`. Games are run in a single
+vmapped batch of `num_games` parallel envs.
+"""
+
+from dataclasses import dataclass, field
+
+import equinox as eqx
+import jax
+import jax.numpy as jnp
+import jax.random as jrandom
+from generals.core import game
+from generals.core.action import compute_valid_move_mask
+from generals.core.env import GeneralsEnv
+from generals.core.observation import Observation
+from omegaconf import DictConfig, OmegaConf
+
+from general_bots_training.network import PolicyValueNetwork, obs_to_tensor
+from general_bots_training.opponents import OPPONENT_TYPES, StaticOpponent
+
+
+@dataclass
+class AgentConfig:
+    """One side of the matchup.
+
+    kind: "random" | "expander" | "hunter" | "model"
+    checkpoint: path to an equinox checkpoint (only used when kind == "model")
+    """
+
+    kind: str = "random"
+    checkpoint: str | None = None
+
+
+@dataclass
+class EvalConfig:
+    grid_dims: tuple[int, int] = (21, 21)
+    truncation: int = 500
+    num_games: int = 100
+    agent0: AgentConfig = field(default_factory=lambda: AgentConfig(kind="random"))
+    agent1: AgentConfig = field(default_factory=lambda: AgentConfig(kind="expander"))
+    seed: int = 0
+
+
+def load_config(args: list[str] | None = None) -> DictConfig:
+    """Merge CLI overrides into the typed default evaluation configuration."""
+    defaults = OmegaConf.structured(EvalConfig)
+    return OmegaConf.merge(defaults, OmegaConf.from_cli(args))  # type: ignore
+
+
+class NetworkAgent:
+    """Wrap a PolicyValueNetwork in the stateless StaticOpponent interface."""
+
+    def __init__(self, network):
+        self._network = network
+
+    def act(self, observation: Observation, key):
+        obs_arr = obs_to_tensor(observation)
+        mask = compute_valid_move_mask(
+            observation.armies, observation.owned_cells, observation.mountains
+        )
+        action, _, _, _ = self._network(obs_arr, mask, key, None)
+        return action
+
+
+def make_agent(cfg: AgentConfig, key) -> StaticOpponent:
+    """Build an agent from its config. `key` seeds the network when needed."""
+    kind = cfg.kind.lower().replace("-", "_")
+    if kind == "model":
+        if cfg.checkpoint is None:
+            raise ValueError("agent kind 'model' requires a checkpoint path")
+        network = PolicyValueNetwork(key, in_channels=14)
+        network = eqx.tree_deserialise_leaves(cfg.checkpoint, network)
+        return NetworkAgent(network)
+    if kind in OPPONENT_TYPES:
+        return OPPONENT_TYPES[kind]()
+    choices = ", ".join([*sorted(OPPONENT_TYPES), "model"])
+    raise ValueError(f"unknown agent kind {cfg.kind!r}; choose one of: {choices}")
+
+
+def agent_label(cfg: AgentConfig) -> str:
+    return cfg.checkpoint if cfg.kind == "model" and cfg.checkpoint else cfg.kind
+
+
+def main(config: DictConfig):
+    key = jrandom.PRNGKey(config.seed)
+    key, agent0_key, agent1_key, env_key = jrandom.split(key, 4)
+
+    agent0 = make_agent(config.agent0, agent0_key)
+    agent1 = make_agent(config.agent1, agent1_key)
+    label0, label1 = agent_label(config.agent0), agent_label(config.agent1)
+
+    env = GeneralsEnv(grid_dims=tuple(config.grid_dims), truncation=config.truncation)
+    pool, _ = env.reset(env_key)
+    get_obs = game.get_full_observation if env.perfect_info else game.get_observation
+
+    n = config.num_games
+    key, init_key = jrandom.split(key)
+    init_keys = jrandom.split(init_key, n)
+    states = jax.vmap(env.init_state)(init_keys)
+    states = states._replace(pool_idx=jnp.arange(n, dtype=states.pool_idx.dtype) % env.pool_size)
+
+    step_env = jax.vmap(env.step, in_axes=(0, 0, None))
+
+    def run_episode(states, key):
+        """Play one full game per env; return final winner per env.
+
+        The env auto-resets on done and overwrites `state.winner` with -1, so we
+        capture the winner from `timestep.info` on the first done step per env and
+        latch it for the remainder of the scan.
+        """
+
+        def body(carry, _):
+            states, latched_winner, key = carry
+            key, p0_key, p1_key = jrandom.split(key, 3)
+            obs_p0 = jax.vmap(lambda s: get_obs(s, 0))(states)
+            obs_p1 = jax.vmap(lambda s: get_obs(s, 1))(states)
+            keys_p0 = jrandom.split(p0_key, n)
+            keys_p1 = jrandom.split(p1_key, n)
+            actions_p0 = jax.vmap(agent0.act)(obs_p0, keys_p0)
+            actions_p1 = jax.vmap(agent1.act)(obs_p1, keys_p1)
+            actions = jnp.stack([actions_p0, actions_p1], axis=1)
+            timesteps, states = step_env(states, actions, pool)
+            done = timesteps.terminated | timesteps.truncated
+            # On the first done, latch the winner (truncated games stay at -1 = draw).
+            new_winner = jnp.where(
+                done & (latched_winner < 0), timesteps.info.winner, latched_winner
+            )
+            return (states, new_winner, key), done
+
+        # Run a fixed number of steps equal to the truncation length, which
+        # guarantees every env has terminated or truncated at least once.
+        (states, latched_winner, key), _ = jax.lax.scan(
+            body, (states, jnp.full(n, -1, dtype=jnp.int32), key), None, length=config.truncation
+        )
+        return latched_winner
+
+    print("Generals.io evaluation")
+    print(OmegaConf.to_yaml(config, resolve=True).rstrip())
+    print(f"device: {jax.devices()[0]}")
+    print(f"agent0: {label0}")
+    print(f"agent1: {label1}")
+    print(f"games:  {n}")
+    print()
+    print("warming up (jit compile)...")
+    winners = run_episode(states, key)
+    jax.block_until_ready(winners)
+    print("warming up done\n")
+
+    print("playing...")
+    winners = run_episode(states, key)
+    jax.block_until_ready(winners)
+
+    wins0 = int(jnp.sum(winners == 0))
+    wins1 = int(jnp.sum(winners == 1))
+    draws = int(jnp.sum(winners < 0))
+
+    print("\nsummary")
+    print("-" * 40)
+    print(f"agent0 ({label0}): {wins0:4d} wins ({wins0 / n * 100:.1f}%)")
+    print(f"agent1 ({label1}): {wins1:4d} wins ({wins1 / n * 100:.1f}%)")
+    print(f"draws:              {draws:4d}      ({draws / n * 100:.1f}%)")
+    print("-" * 40)
+    print(f"total games: {n}")
+
+
+if __name__ == "__main__":
+    main(load_config())

uv.lock

@@ -132,12 +132,20 @@ source = { editable = "." }
 dependencies = [
     { name = "equinox" },
     { name = "generals-bots" },
-    { name = "jax", extra = ["cuda"] },
+    { name = "jax" },
     { name = "jaxtyping" },
     { name = "omegaconf" },
     { name = "optax" },
 ]
 
+[package.optional-dependencies]
+cuda = [
+    { name = "jax", extra = ["cuda"] },
+]
+cuda13 = [
+    { name = "jax", extra = ["cuda13"] },
+]
+
 [package.dev-dependencies]
 dev = [
     { name = "ipython" },
@@ -148,11 +156,14 @@ dev = [
 requires-dist = [
     { name = "equinox", specifier = ">=0.13.8" },
     { name = "generals-bots", git = "https://github.com/strakam/generals-bots.git" },
-    { name = "jax", extras = ["cuda"], specifier = ">=0.11.0" },
+    { name = "jax", specifier = "==0.11.0" },
+    { name = "jax", extras = ["cuda"], marker = "extra == 'cuda'", specifier = "==0.11.0" },
+    { name = "jax", extras = ["cuda13"], marker = "extra == 'cuda13'", specifier = "==0.11.0" },
     { name = "jaxtyping", specifier = ">=0.3.11" },
     { name = "omegaconf", specifier = ">=2.3.1" },
     { name = "optax", specifier = ">=0.2.8" },
 ]
+provides-extras = ["cuda", "cuda13"]
 
 [package.metadata.requires-dev]
 dev = [
@@ -253,6 +264,10 @@ cuda = [
     { name = "jax-cuda12-plugin", extra = ["with-cuda"] },
     { name = "jaxlib" },
 ]
+cuda13 = [
+    { name = "jax-cuda13-plugin", extra = ["with-cuda"] },
+    { name = "jaxlib" },
+]
 
 [[package]]
 name = "jax-cuda12-pjrt"
@@ -295,6 +310,48 @@ with-cuda = [
     { name = "nvidia-nvshmem-cu12", marker = "sys_platform == 'linux'" },
 ]
 
+[[package]]
+name = "jax-cuda13-pjrt"
+version = "0.11.0"
+source = { registry = "https://pypi.org/simple" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/dd/cd/8fa8034e2f7bc50a78c2319f576aedf06c633b94b3e5d187fb9becfc1379/jax_cuda13_pjrt-0.11.0-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:e72c7b6b266b3fbe85792b2b4569b6c87447e77c454770d5afcc081a0b87d140", size = 120632519, upload-time = "2026-07-16T18:58:54.921Z" },
+    { url = "https://files.pythonhosted.org/packages/6e/0e/140944fc5cc1e0d5891b78209fb70a2f1ce951a656292f598904ac7efe5d/jax_cuda13_pjrt-0.11.0-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:814092a06a4da93a0b01df9bcdf5f8e0900d6c0a4c3ca7e7ec4ec967de7e96d9", size = 125815692, upload-time = "2026-07-16T18:58:59.251Z" },
+]
+
+[[package]]
+name = "jax-cuda13-plugin"
+version = "0.11.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "jax-cuda13-pjrt" },
+]
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/19/59/e744e3362e8139e497b1a0c1bb27666ff7b1a4a4e20007b5105afc27deaf/jax_cuda13_plugin-0.11.0-cp313-cp313-manylinux_2_27_aarch64.whl", hash = "sha256:7d11c242dc94cd24c44ce0b67f49d5771040a5e83b594b31b9ad42e88fb8821b", size = 7547142, upload-time = "2026-07-16T18:59:06.604Z" },
+    { url = "https://files.pythonhosted.org/packages/05/c2/f5eaf303b8497d3c1b8445935a8e11181591ce4a99054699d57e2cfe8f42/jax_cuda13_plugin-0.11.0-cp313-cp313-manylinux_2_27_x86_64.whl", hash = "sha256:5fda3b7b12e50b53a0d8f9cd8c6b6efe160de9b1054dd3c9652f1504d34fc9d8", size = 7571108, upload-time = "2026-07-16T18:59:08.06Z" },
+    { url = "https://files.pythonhosted.org/packages/06/14/74da5a18490a7d24bdc21dedf90b20eab4e4658879f333bea9772f36291b/jax_cuda13_plugin-0.11.0-cp314-cp314-manylinux_2_27_aarch64.whl", hash = "sha256:d9147b3db90f140bd94afb29589a7eaeb7fea1b11159420394c60d466ea9a7c1", size = 7547841, upload-time = "2026-07-16T18:59:10.03Z" },
+    { url = "https://files.pythonhosted.org/packages/42/35/9131f3182b4ef67dce4845c697c8cab79bc0df8ffb15e7806bd099cae4e3/jax_cuda13_plugin-0.11.0-cp314-cp314-manylinux_2_27_x86_64.whl", hash = "sha256:09be21c598788f0b714f27cc94e3ff8a169b7854b6db883287d6cbc6d4d8d911", size = 7571369, upload-time = "2026-07-16T18:59:11.811Z" },
+    { url = "https://files.pythonhosted.org/packages/29/67/7f057e2533f14fcdfa236deb2272e202549fd0a1b461c65d8f9c93fc2c6c/jax_cuda13_plugin-0.11.0-cp314-cp314t-manylinux_2_27_aarch64.whl", hash = "sha256:87d847502f03e7ce37e6ef570245fca1323af9a7c0b86ffb615f794e24525d70", size = 7561528, upload-time = "2026-07-16T18:59:13.288Z" },
+    { url = "https://files.pythonhosted.org/packages/3f/27/1ada0a898c34d009478814fb5a1c2c73cb59036d68563992a1ee0dfcadde/jax_cuda13_plugin-0.11.0-cp314-cp314t-manylinux_2_27_x86_64.whl", hash = "sha256:8f9e14c1b7ef52cccf62594fecd0cde2bb150a0f933e831b0b66b2cb27d91000", size = 7580472, upload-time = "2026-07-16T18:59:14.734Z" },
+]
+
+[package.optional-dependencies]
+with-cuda = [
+    { name = "nvidia-cublas", marker = "sys_platform == 'linux'" },
+    { name = "nvidia-cuda-cupti", marker = "sys_platform == 'linux'" },
+    { name = "nvidia-cuda-nvcc", marker = "sys_platform == 'linux'" },
+    { name = "nvidia-cuda-nvrtc", marker = "sys_platform == 'linux'" },
+    { name = "nvidia-cuda-runtime", marker = "sys_platform == 'linux'" },
+    { name = "nvidia-cudnn-cu13", marker = "sys_platform == 'linux'" },
+    { name = "nvidia-cufft", marker = "sys_platform == 'linux'" },
+    { name = "nvidia-cusolver", marker = "sys_platform == 'linux'" },
+    { name = "nvidia-cusparse", marker = "sys_platform == 'linux'" },
+    { name = "nvidia-nccl-cu13", marker = "sys_platform == 'linux'" },
+    { name = "nvidia-nvjitlink", marker = "sys_platform == 'linux'" },
+    { name = "nvidia-nvshmem-cu13", marker = "sys_platform == 'linux'" },
+    { name = "nvidia-nvvm" },
+]
+
 [[package]]
 name = "jaxlib"
 version = "0.11.0"
@@ -425,6 +482,18 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/a1/5a/4d2b1601df3602dba7a14f3348ba9bfe94a18adb428e693df6154c293831/numpy-2.5.1-cp314-cp314t-win_arm64.whl", hash = "sha256:5a6db61f9aaa57e369905c67d852045d3c4f7126405b29d09b19dec118e9c9cb", size = 10697674, upload-time = "2026-07-04T17:07:58.506Z" },
 ]
 
+[[package]]
+name = "nvidia-cublas"
+version = "13.6.1.10"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "nvidia-cuda-nvrtc" },
+]
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/42/51/a174a0e79793e528ef4645a9d554e3aa48fd8da3f9c9cf523c0176bb121b/nvidia_cublas-13.6.1.10-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:e05a431062a17cb9b02e2f37e67817b637051ce8fad57b388482c594396ddbb4", size = 518610411, upload-time = "2026-07-30T18:18:32.969Z" },
+    { url = "https://files.pythonhosted.org/packages/dd/e9/288a93d8234b8f8ea0ccac7b8cc5d7aa4c663d9676c64f4a2eb3a1f9ffc4/nvidia_cublas-13.6.1.10-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:feb2ed8a1e211bc5774413efc0f1a08c4d5269b56f68b4ac6fe5408e57f7dc1c", size = 410475904, upload-time = "2026-07-30T18:19:26.388Z" },
+]
+
 [[package]]
 name = "nvidia-cublas-cu12"
 version = "12.9.2.10"
@@ -437,6 +506,15 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/cb/c0/0a517bfe63ccd3b92eb254d264e28fca3c7cab75d07daea315250fb1bf73/nvidia_cublas_cu12-12.9.2.10-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:e4f53a8ca8c5d6e8c492d0d0a3d565ecb59a751b19cfdaa4f6da0ab2104c1702", size = 581240110, upload-time = "2026-04-08T18:52:31.532Z" },
 ]
 
+[[package]]
+name = "nvidia-cuda-cccl"
+version = "13.3.3.4.1"
+source = { registry = "https://pypi.org/simple" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/96/bd/572971ffc14bd36676c821fc15d991b08fe6179cb09368250147475f954d/nvidia_cuda_cccl-13.3.3.4.1-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:067d19b4b3c9d0f2ebec9f29a311b2863db96bf98e058bbc331597d51ce818cf", size = 3454030, upload-time = "2026-06-29T16:41:49.092Z" },
+    { url = "https://files.pythonhosted.org/packages/f8/ab/049726d90147865a3ea53bae6cb7c35b98bf1fdf96cdb967101329625f83/nvidia_cuda_cccl-13.3.3.4.1-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:cc0adc188d570b09f4d606c7dc05a42aa3d8aa082e0d60f7bbfc5b6435f627c6", size = 3454034, upload-time = "2026-06-29T16:42:07.435Z" },
+]
+
 [[package]]
 name = "nvidia-cuda-cccl-cu12"
 version = "12.9.27"
@@ -446,6 +524,24 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/18/2a/d4cd8506d2044e082f8cd921be57392e6a9b5ccd3ffdf050362430a3d5d5/nvidia_cuda_cccl_cu12-12.9.27-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:37869e17ce2e1ecec6eddf1927cca0f8c34e64fd848d40453df559091e2d7117", size = 3152243, upload-time = "2025-05-01T19:32:13.955Z" },
 ]
 
+[[package]]
+name = "nvidia-cuda-crt"
+version = "13.3.73"
+source = { registry = "https://pypi.org/simple" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/fa/41/2089e411507d66458d67208bdd1bc562d492bb6458c3d2aea4603072a219/nvidia_cuda_crt-13.3.73-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:60aacc0b5e1e8b40c62abe4d1ab16440add91b99bd2f17f62dd091586b73d166", size = 157353, upload-time = "2026-06-29T16:42:38.163Z" },
+    { url = "https://files.pythonhosted.org/packages/7e/ce/16d76f4b5b3f7460f5ebd17516685495c149c66651ffdd381f90e4d4e65c/nvidia_cuda_crt-13.3.73-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:df14a17ae1c5c3171265411212246654d780f89344ea85344466c6b955247543", size = 157352, upload-time = "2026-06-29T16:43:09.209Z" },
+]
+
+[[package]]
+name = "nvidia-cuda-cupti"
+version = "13.3.75"
+source = { registry = "https://pypi.org/simple" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/d1/7b/d332e362c2c48929f913a5930a596c707c30632ae3248dd16be35f18fc11/nvidia_cuda_cupti-13.3.75-py3-none-manylinux_2_25_aarch64.whl", hash = "sha256:79f5a8c55378a998427d6974b1d85ec91ef684035d016b3fb3b85006914f94f0", size = 11824795, upload-time = "2026-06-29T16:46:00.894Z" },
+    { url = "https://files.pythonhosted.org/packages/d1/68/ee003f64d79bd3cb95e243cfdc906bb0f1ae73b1aa2857baf1d7834dd1ee/nvidia_cuda_cupti-13.3.75-py3-none-manylinux_2_25_x86_64.whl", hash = "sha256:a969cc446c44db612e417bd0af722cfad9009c027dabd4ed255262c0e9a9a3d2", size = 12044533, upload-time = "2026-06-29T16:46:20.662Z" },
+]
+
 [[package]]
 name = "nvidia-cuda-cupti-cu12"
 version = "12.9.79"
@@ -455,6 +551,20 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/c1/2e/b84e32197e33f39907b455b83395a017e697c07a449a2b15fd07fc1c9981/nvidia_cuda_cupti_cu12-12.9.79-py3-none-manylinux_2_25_x86_64.whl", hash = "sha256:096bcf334f13e1984ba36685ad4c1d6347db214de03dbb6eebb237b41d9d934f", size = 10814997, upload-time = "2025-06-05T20:01:10.168Z" },
 ]
 
+[[package]]
+name = "nvidia-cuda-nvcc"
+version = "13.3.73"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "nvidia-cuda-crt" },
+    { name = "nvidia-cuda-runtime" },
+    { name = "nvidia-nvvm" },
+]
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/5c/14/9f5cdc994d5431e2f08f62ffe34509e7feabd1f2e18517e2d7720c6ff0fd/nvidia_cuda_nvcc-13.3.73-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:70f250825355d2c3aa6c7a972a0ec00f020bad66d2679e527eb4336301c904aa", size = 39515578, upload-time = "2026-06-29T16:47:40.318Z" },
+    { url = "https://files.pythonhosted.org/packages/83/19/e46ef3597ba47a9f8a91ab24533db42a600b659fc418dbe4af0b630bcb41/nvidia_cuda_nvcc-13.3.73-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:f483af83166c4fa356a21606076d553b0b4ceaebbd9912e537545080db695bdd", size = 44942138, upload-time = "2026-06-29T16:48:13.615Z" },
+]
+
 [[package]]
 name = "nvidia-cuda-nvcc-cu12"
 version = "12.9.86"
@@ -464,6 +574,15 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/d6/5c/8cc072436787104bbbcbde1f76ab4a0d89e68f7cebc758dd2ad7913a43d0/nvidia_cuda_nvcc_cu12-12.9.86-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:44e1eca4d08926193a558d2434b1bf83d57b4d5743e0c431c0c83d51da1df62b", size = 39411138, upload-time = "2025-06-05T20:01:43.182Z" },
 ]
 
+[[package]]
+name = "nvidia-cuda-nvrtc"
+version = "13.3.33"
+source = { registry = "https://pypi.org/simple" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/8b/2c/86916c8a34dcdb0c3ddd1c0e30545041bd781184e437b9cb76fcda70560b/nvidia_cuda_nvrtc-13.3.33-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:82530788b8c6164a54d3fd9ae8bcca8893d397c4aeb998861982a03bbe41e204", size = 51110910, upload-time = "2026-05-26T16:38:16.116Z" },
+    { url = "https://files.pythonhosted.org/packages/e7/b6/60a3641111d39ebfcfcd8b8bfd0290d7623c4b8b5f90952c2d84776f8ca4/nvidia_cuda_nvrtc-13.3.33-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:7b05ecda494c6dabc44231a608b060a71008a730d9dfda932cc508e6d29159e0", size = 49260054, upload-time = "2026-05-26T16:37:51.177Z" },
+]
+
 [[package]]
 name = "nvidia-cuda-nvrtc-cu12"
 version = "12.9.86"
@@ -473,6 +592,15 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/64/eb/c2295044b8f3b3b08860e2f6a912b702fc92568a167259df5dddb78f325e/nvidia_cuda_nvrtc_cu12-12.9.86-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:096d4de6bda726415dfaf3198d4f5c522b8e70139c97feef5cd2ca6d4cd9cead", size = 44528905, upload-time = "2025-06-05T20:02:29.754Z" },
 ]
 
+[[package]]
+name = "nvidia-cuda-runtime"
+version = "13.3.29"
+source = { registry = "https://pypi.org/simple" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/5f/e5/c1a221c8e6fecd071b80ea44c20fc253ae24f56e15e3f77cfbc3fb76e724/nvidia_cuda_runtime-13.3.29-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:73291e19c9dd919c140c91bda2f80b0eca487da5ee30a086ef7bc4918ecb90ea", size = 2356574, upload-time = "2026-05-26T16:29:56.333Z" },
+    { url = "https://files.pythonhosted.org/packages/97/be/5699b6e642b372f7d24c59c2f41383e2696825e20bab85f7399c7c6a56f7/nvidia_cuda_runtime-13.3.29-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:e04420616e72f563167a7733272992d7e6df6dc5cb54b2f94f9f1520ea9e30c1", size = 2339786, upload-time = "2026-05-26T16:30:21.584Z" },
+]
+
 [[package]]
 name = "nvidia-cuda-runtime-cu12"
 version = "12.9.79"
@@ -494,6 +622,30 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/10/13/b8887c869cf2471339a24b60d3c28e761facbb534935f572b61423371abb/nvidia_cudnn_cu12-9.24.0.43-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:f424192dd85e7d29f44be18df2dae4c80d32c67a29c0d42f5c283c40cfdf871c", size = 799083985, upload-time = "2026-07-02T16:25:37.467Z" },
 ]
 
+[[package]]
+name = "nvidia-cudnn-cu13"
+version = "9.24.0.43"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "nvidia-cublas" },
+]
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/ca/30/7c257e3d5cb4fecb147b93895c66e29c93f8e76d74b45bb418ff0587c4ec/nvidia_cudnn_cu13-9.24.0.43-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:a6812a554a1ff0413e9c52b84c26c050380649ab9615f9c16bded368ce9f421f", size = 650976863, upload-time = "2026-07-02T16:23:39.248Z" },
+    { url = "https://files.pythonhosted.org/packages/5c/ba/791cffd048fe5b044e620df55267e3e95c0e6e07d50b41e377c03dfc910f/nvidia_cudnn_cu13-9.24.0.43-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:71f181cd810e90f9b6023b01186fe82d13d65f0ec098581ee201d39fad769e4b", size = 553099438, upload-time = "2026-07-02T16:27:42.58Z" },
+]
+
+[[package]]
+name = "nvidia-cufft"
+version = "12.3.0.29"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "nvidia-nvjitlink" },
+]
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/f9/bc/543e6dfbcb659ed07a94a72d8925c8b5688dabc1f616ab2f6575483a3cb6/nvidia_cufft-12.3.0.29-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:3c268ec73fdbfd22e42913dcef6e49e0c305ab09cd9848268a841cbe9f33f748", size = 184700143, upload-time = "2026-05-26T16:45:25.443Z" },
+    { url = "https://files.pythonhosted.org/packages/e7/00/fab4a29fa1d7eb43bc6b94de4e86312c5e425d5582e58b9641300b9dffc7/nvidia_cufft-12.3.0.29-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:edb25c0626bd202ee5acc035b5dd361a3b89ed3b75a81a52df72c89150cb57c2", size = 184730406, upload-time = "2026-05-26T16:46:18.193Z" },
+]
+
 [[package]]
 name = "nvidia-cufft-cu12"
 version = "11.4.1.4"
@@ -506,6 +658,20 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/95/f4/61e6996dd20481ee834f57a8e9dca28b1869366a135e0d42e2aa8493bdd4/nvidia_cufft_cu12-11.4.1.4-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:c67884f2a7d276b4b80eb56a79322a95df592ae5e765cf1243693365ccab4e28", size = 200877592, upload-time = "2025-06-05T20:05:45.862Z" },
 ]
 
+[[package]]
+name = "nvidia-cusolver"
+version = "12.2.6.9"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "nvidia-cublas" },
+    { name = "nvidia-cusparse" },
+    { name = "nvidia-nvjitlink" },
+]
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/13/bd/61dcd7917f64c9e81a35f939f1bf4ab6535adf712192e915fb472e5bbf6c/nvidia_cusolver-12.2.6.9-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:658ce9e6ed8a321033419fe6e55faf0f245a15aa23669eda03d7f2dcce436d73", size = 267275679, upload-time = "2026-06-29T16:59:54.924Z" },
+    { url = "https://files.pythonhosted.org/packages/1c/3a/23e46a0919ba5ecf1a864400542a6c78d18c76893a0a77e1856e7033af31/nvidia_cusolver-12.2.6.9-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:e32cb40c8a4d3df475b556caaf6af5808ef064df7e291a049528a0d28017b674", size = 244667799, upload-time = "2026-06-29T17:00:42.543Z" },
+]
+
 [[package]]
 name = "nvidia-cusolver-cu12"
 version = "11.7.5.82"
@@ -520,6 +686,18 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/33/40/79b0c64d44d6c166c0964ec1d803d067f4a145cca23e23925fd351d0e642/nvidia_cusolver_cu12-11.7.5.82-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:15da72d1340d29b5b3cf3fd100e3cd53421dde36002eda6ed93811af63c40d88", size = 338117415, upload-time = "2025-06-05T20:07:16.809Z" },
 ]
 
+[[package]]
+name = "nvidia-cusparse"
+version = "12.8.2.51"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "nvidia-nvjitlink" },
+]
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/ac/bc/925058f9343d93426864cbb113a3f5cf6db97a7d11642aaf4159b6f0c57a/nvidia_cusparse-12.8.2.51-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:00469fcf62c4d464a1225abd9b20864ecff35e3fbc9fb992572e83d358927755", size = 172868683, upload-time = "2026-06-29T17:01:23.244Z" },
+    { url = "https://files.pythonhosted.org/packages/a6/dc/752e401a5d6fc5f1948cf190add5afc27e440e5ca08bc3fab66826c16213/nvidia_cusparse-12.8.2.51-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:65cbcc4e37a34fca4ee7df2fd57da103593842cda1bbb4a144664ecfe59873a5", size = 154538196, upload-time = "2026-06-29T17:02:06.078Z" },
+]
+
 [[package]]
 name = "nvidia-cusparse-cu12"
 version = "12.5.10.65"
@@ -541,6 +719,24 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/50/32/e7ffa9c324ae260e5dbb4af2cd557bf7a8d155c8ac7b79a785fe1796fb92/nvidia_nccl_cu12-2.30.7-py3-none-manylinux_2_18_x86_64.whl", hash = "sha256:8ce1b8213f61f2bfac132e6df890af6450b77cbd140c6ce4e98cb0c2d8e678c9", size = 303361239, upload-time = "2026-06-09T03:24:53.816Z" },
 ]
 
+[[package]]
+name = "nvidia-nccl-cu13"
+version = "2.30.7"
+source = { registry = "https://pypi.org/simple" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/d1/21/a73174c6157101bdf1ffc22b517f76ff0082613989dd9bc8f43e8034caac/nvidia_nccl_cu13-2.30.7-py3-none-manylinux_2_18_aarch64.whl", hash = "sha256:ca786ffa5a647c75d4d1f5cc72a6c4f537947e2ba8823d7c8aaf768e7a7b9f77", size = 215983881, upload-time = "2026-06-09T03:23:15.633Z" },
+    { url = "https://files.pythonhosted.org/packages/3f/34/c500f90c7ae641b8e0f98965b36b8a7ac79cc8b296e8d251fe3eb592ee54/nvidia_nccl_cu13-2.30.7-py3-none-manylinux_2_18_x86_64.whl", hash = "sha256:cefa7fdb9710efd0f39c5f1be1d61ff6fc9a996c451265bd7fbdcf9455ed4b50", size = 215965170, upload-time = "2026-06-09T03:23:39.73Z" },
+]
+
+[[package]]
+name = "nvidia-nvjitlink"
+version = "13.3.33"
+source = { registry = "https://pypi.org/simple" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/f0/ee/580ca6f29dcab0221db8706badca1bbbb084f1975c4d4e83329c3a7e31f0/nvidia_nvjitlink-13.3.33-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:26a6de7fb4c8fdaa7703d3dad720d6d427ddfea5c48a528fd97c11733ad830e5", size = 40742423, upload-time = "2026-05-26T16:54:51.613Z" },
+    { url = "https://files.pythonhosted.org/packages/69/30/45414e35ff2eee7db3da037e5707037ccf9d2b5218ffbdb055ea4d5aa98a/nvidia_nvjitlink-13.3.33-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:ce48b37dfeb3cb1eae4cf85adacb47d7a6539ea2272870c9a3628ce275c2037e", size = 39168635, upload-time = "2026-05-26T16:54:13.906Z" },
+]
+
 [[package]]
 name = "nvidia-nvjitlink-cu12"
 version = "12.9.86"
@@ -562,6 +758,28 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/63/03/cde130e3ff706784f9efd47204f299b15113ea08203e98a556603dac245d/nvidia_nvshmem_cu12-3.7.2-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:02b6d3482d90ea8ba214bc400362297f573257763990a98c9846ec5c786a7227", size = 230171132, upload-time = "2026-07-17T19:25:07.092Z" },
 ]
 
+[[package]]
+name = "nvidia-nvshmem-cu13"
+version = "3.7.2"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "nvidia-cuda-cccl" },
+]
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/81/11/1767ef120299972822a43aceffbc81ac813aad2089d817e36e4e12aa4c9c/nvidia_nvshmem_cu13-3.7.2-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:0e748308b9e0a8a8d81859353f3e6c98fbe85a85f652a915a2ba60149a0bd124", size = 135174619, upload-time = "2026-07-17T19:24:15.404Z" },
+    { url = "https://files.pythonhosted.org/packages/4f/4a/037dfbbafc96ad1828ef62326f27970cae2c0eb271cf5b8e6922f32a7427/nvidia_nvshmem_cu13-3.7.2-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:b9a4f843c62b5dded8ad33fd6d220e007dcd39f06cee549292f5ba14ce4c235c", size = 135397797, upload-time = "2026-07-17T19:25:26.4Z" },
+]
+
+[[package]]
+name = "nvidia-nvvm"
+version = "13.3.73"
+source = { registry = "https://pypi.org/simple" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/f3/e7/ff646aa6015c7e6d12aad234e68925c87b6681d8d18c3ac40535994a3b0d/nvidia_nvvm-13.3.73-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:0e28e0858a3475e11ac67d35301cd5bf82666a1c0dc4ec4e80ceaf3a5fd1dea8", size = 69250424, upload-time = "2026-06-29T17:08:07.453Z" },
+    { url = "https://files.pythonhosted.org/packages/2f/05/35754a7105563fd9b496e5ee8e1acd986aef8258760c3cbccf419aee861a/nvidia_nvvm-13.3.73-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:e2bcdd5783b5481445f1f0e7170cb836cc0d72999839ba850bbba6dc97b76bb8", size = 66984478, upload-time = "2026-06-29T17:07:43.765Z" },
+    { url = "https://files.pythonhosted.org/packages/ad/6b/d5756f485012b920475cbc01457c1b9a7d0485bfb04b92598c5e1ef3e9ab/nvidia_nvvm-13.3.73-py3-none-win_amd64.whl", hash = "sha256:b5c91dfa59ee4cee90b2dfb19c6203f31c914b9c9b5ca10726c2da7cf8ed401d", size = 59981103, upload-time = "2026-06-29T17:21:43.334Z" },
+]
+
 [[package]]
 name = "omegaconf"
 version = "2.3.1"