diff --git a/.gitignore b/.gitignore deleted file mode 100644 index 5481b47..0000000 --- a/.gitignore +++ /dev/null @@ -1,2 +0,0 @@ -__pycache__/ -competition/.venv-competition/ diff --git a/Justfile b/Justfile deleted file mode 100644 index f8ef3b5..0000000 --- a/Justfile +++ /dev/null @@ -1,2 +0,0 @@ -vendor: - pip download -r competition/requirements-vendor.txt -d competition/wheelhouse/ --python-version 3.12 --platform manylinux --no-deps diff --git a/competition-submission.zip b/competition-submission.zip deleted file mode 100644 index 5903f80..0000000 Binary files a/competition-submission.zip and /dev/null differ diff --git a/competition-submission.zip.bak b/competition-submission.zip.bak deleted file mode 100644 index e22ec3e..0000000 Binary files a/competition-submission.zip.bak and /dev/null differ diff --git a/competition/build.sh b/competition/build.sh deleted file mode 100755 index bcd1014..0000000 --- a/competition/build.sh +++ /dev/null @@ -1,7 +0,0 @@ -#!/usr/bin/env bash -# One-time intake step: install the vendored wheels offline. -# The sandbox has no network, so everything must come from wheelhouse/. -set -euo pipefail -cd "$(dirname "$0")" -python -m pip install --no-index --find-links=wheelhouse -r requirements-vendor.txt -echo "[build] vendored wheels installed" >&2 diff --git a/competition/competition.py b/competition/competition.py deleted file mode 100644 index 507b50f..0000000 --- a/competition/competition.py +++ /dev/null @@ -1,336 +0,0 @@ -"""Competition stdio agent: direct policy sampling over a trained equinox checkpoint. - -Self-contained competition bot. It speaks the wire protocol in -`generals-bots/competition/protocol.py` (handshake, then one observation frame -per turn, one action line per turn, EOF on stdin = game over) and, on each -turn, samples an action from the policy network exactly as the model is used -during training/evaluation — no tree search. - -Unlike `scripts/mcts_agent.py`, this module does **not** import the -`general_bots_training` training package. It inlines the policy-value network -and the observation encoding so the script can be dropped into the competition -sandbox with just the engine (`generals`), `equinox`, `jax`, and `numpy` -available. - -This is the search-free counterpart to `scripts/competition_puct.py`: it shares -the same network and observation encoding, but selects an action by sampling -from the masked policy (as in `PolicyValueNetwork.__call__`) instead of running -particle PUCT. - -Run it directly through the bundled matchup driver, e.g.: - - PYTHONPATH=src:generals-bots .venv/bin/python \ - generals-bots/competition/matchup.py \ - scripts/competition.py --checkpoint ppo_model.eqx \ - generals-bots/competition/agents/expander_python/run.sh \ - --mode competition - -or wrap it in a `run.sh` that passes the checkpoint path. -""" - -import argparse -import os -import sys -import time - -import equinox as eqx -import jax -import jax.numpy as jnp -import jax.random as jrandom -import numpy as np -from generals.core import game -from generals.core.action import compute_valid_move_mask -from generals.core.game import GameState -from generals.core.observation import Observation -from generals.modifiers import build_castles, deathtouch - -os.environ.setdefault("JAX_PLATFORMS", "cpu") - -PASS_ACTION = np.array([1, 0, 0, 0, 0], dtype=np.int32) - - -# --------------------------------------------------------------------------- -# Policy-value network (inlined from general_bots_training.network) -# --------------------------------------------------------------------------- - - -def obs_to_tensor(obs: Observation) -> jnp.ndarray: - """Encode an Observation into a (C, H, W) float32 tensor for the network. - - Armies and army-counts are log-normalized; land counts are divided by the - number of cells; scalar values are broadcast to spatial planes so a plain - conv stack can consume them. Grid-size agnostic. - """ - H, W = obs.armies.shape - armies = jnp.log1p(obs.armies.astype(jnp.float32)) / jnp.log(50.0) - - def bcast(scalar): - return jnp.broadcast_to(scalar.astype(jnp.float32), (H, W)) - - own_land = bcast(obs.owned_land_count / (H * W)) - own_army = bcast(jnp.log1p(obs.owned_army_count.astype(jnp.float32)) / jnp.log(50.0)) - opp_land = bcast(obs.opponent_land_count / (H * W)) - opp_army = bcast(jnp.log1p(obs.opponent_army_count.astype(jnp.float32)) / jnp.log(50.0)) - timestep = bcast(jnp.log1p(obs.timestep.astype(jnp.float32)) / jnp.log(1201.0)) - - return jnp.stack( - [ - armies, - obs.generals.astype(jnp.float32), - obs.castles.astype(jnp.float32), - obs.mountains.astype(jnp.float32), - obs.neutral_cells.astype(jnp.float32), - obs.owned_cells.astype(jnp.float32), - obs.opponent_cells.astype(jnp.float32), - obs.fog_cells.astype(jnp.float32), - obs.structures_in_fog.astype(jnp.float32), - own_land, - own_army, - opp_land, - opp_army, - timestep, - ], - axis=0, - ) - - -class PolicyValueNetwork(eqx.Module): - """Conv policy-value network. - - Action layout: 4 full-move directions and 4 half-move (split) directions - per source cell, plus one global pass action. Invalid moves are masked to - -1e9; pass is always available. - - The value head uses global average pooling over the spatial grid, so the - network works for any board size without reshaping linear layers. - """ - - conv1: eqx.nn.Conv2d - conv2: eqx.nn.Conv2d - conv3: eqx.nn.Conv2d - conv4: eqx.nn.Conv2d - policy_conv: eqx.nn.Conv2d - value_conv: eqx.nn.Conv2d - value_linear1: eqx.nn.Linear - value_linear2: eqx.nn.Linear - - def __init__(self, key, in_channels: int = 14, channels=(32, 32, 32, 16)): - keys = jrandom.split(key, 8) - self.conv1 = eqx.nn.Conv2d(in_channels, channels[0], kernel_size=3, padding=1, key=keys[0]) - self.conv2 = eqx.nn.Conv2d(channels[0], channels[1], kernel_size=3, padding=1, key=keys[1]) - self.conv3 = eqx.nn.Conv2d(channels[1], channels[2], kernel_size=3, padding=1, key=keys[2]) - self.conv4 = eqx.nn.Conv2d(channels[2], channels[3], kernel_size=3, padding=1, key=keys[3]) - # 9 = 4 dirs (full) + 4 dirs (half) + 1 pass - self.policy_conv = eqx.nn.Conv2d(channels[3], 9, kernel_size=1, key=keys[4]) - self.value_conv = eqx.nn.Conv2d(channels[3], 4, kernel_size=1, key=keys[5]) - self.value_linear1 = eqx.nn.Linear(4, 64, key=keys[6]) - self.value_linear2 = eqx.nn.Linear(64, 1, key=keys[7]) - - def _features(self, obs): - x = jax.nn.relu(self.conv1(obs)) - x = jax.nn.relu(self.conv2(x)) - x = jax.nn.relu(self.conv3(x)) - x = jax.nn.relu(self.conv4(x)) - return x - - def _value_from_features(self, feat): - v = jax.nn.relu(self.value_conv(feat)) # (4, H, W) - v = v.mean(axis=(1, 2)) # (4,) global average pool - v = jax.nn.relu(self.value_linear1(v)) - return self.value_linear2(v)[0] - - def policy_value(self, obs, mask): - """Return deterministic masked policy logits and the critic value.""" - features = self._features(obs) - return self._policy_logits(features, mask), self._value_from_features(features) - - def _policy_logits(self, feat, mask): - """Return 8*H*W move logits plus one global pass logit.""" - logits = self.policy_conv(feat) # (9, H, W) - mask_t = jnp.transpose(mask, (2, 0, 1)) # (4, H, W) - penalty = (1.0 - mask_t) * -1e9 - move_penalty = jnp.concatenate([penalty, penalty], axis=0) - move_logits = (logits[:8] + move_penalty).reshape(-1) - pass_logit = jnp.mean(logits[8])[None] - return jnp.concatenate([move_logits, pass_logit]) - - -# --------------------------------------------------------------------------- -# Action selection (matches training/evaluation rollout) -# --------------------------------------------------------------------------- - - -@jax.jit -def competition_move_mask(obs: Observation) -> jnp.ndarray: - """Legal move mask, identical to the one used during training/evaluation.""" - return compute_valid_move_mask(obs.armies, obs.owned_cells, obs.mountains) - - -def _decode_policy_index(index: int, height: int, width: int) -> np.ndarray: - cells = height * width - if index == 8 * cells: - return PASS_ACTION.copy() - encoded_direction, position = divmod(index, cells) - row, col = divmod(position, width) - split = int(encoded_direction >= 4) - direction = encoded_direction - 4 if split else encoded_direction - return np.array([0, row, col, direction, split], dtype=np.int32) - - -def observation_from_wire( - timestep: int, - owned_land_count: int, - owned_army_count: int, - opponent_land_count: int, - opponent_army_count: int, - type_grid, - owner_grid, - army_grid, -) -> Observation: - """Convert a competition wire frame into the network's Observation type.""" - types = jnp.asarray(type_grid, dtype=jnp.int32) - owners = jnp.asarray(owner_grid, dtype=jnp.int32) - armies = jnp.asarray(army_grid, dtype=jnp.int32) - visible = (types != 0) & (types != 5) - return Observation( - armies=armies, - generals=types == 4, - castles=types == 3, - mountains=types == 2, - neutral_cells=visible & (owners == 0) & (types != 2), - owned_cells=owners == 1, - opponent_cells=owners == 2, - fog_cells=types == 0, - structures_in_fog=types == 5, - owned_land_count=jnp.int32(owned_land_count), - owned_army_count=jnp.int32(owned_army_count), - opponent_land_count=jnp.int32(opponent_land_count), - opponent_army_count=jnp.int32(opponent_army_count), - timestep=jnp.int32(timestep), - ) - - -@jax.jit -def competition_step(state: GameState, actions: jnp.ndarray): - state, actions = build_castles.apply_build_actions(state, actions) - return deathtouch.step(state, actions, turn=800) - - -class PolicyAgent: - """Direct policy agent that samples actions as in training/evaluation. - - On each turn it computes the masked policy logits over all legal moves and - samples one action with `jrandom.categorical`, exactly mirroring - `PolicyValueNetwork.__call__` used during rollout collection. - """ - - def __init__( - self, - network: PolicyValueNetwork, - seed: int = 0, - ): - self.network = network - self.key = jrandom.PRNGKey(seed) - self._infer = eqx.filter_jit( - lambda network, tensor, mask: network.policy_value(tensor, mask) - ) - - def warmup(self, height: int, width: int) -> None: - """Compile shape-dependent network and competition transition kernels.""" - grid = jnp.zeros((height, width), dtype=jnp.int32) - grid = grid.at[0, 0].set(1).at[height - 1, width - 1].set(2) - state = game.create_initial_state(grid) - state = state._replace( - armies=state.armies.at[0, 0].set(100).at[height - 1, width - 1].set(100), - time=jnp.int32(100), - ) - obs = game.get_observation(state, 0) - self.policy_value(obs) - actions = jnp.stack([jnp.asarray(PASS_ACTION), jnp.asarray(PASS_ACTION)]) - warmed_state, _ = competition_step(state, actions) - jax.block_until_ready(warmed_state) - - def policy_value(self, obs: Observation): - mask = competition_move_mask(obs) - logits, value = self._infer(self.network, obs_to_tensor(obs), mask) - return np.asarray(logits), float(value), mask - - def select(self, obs: Observation) -> tuple[np.ndarray, dict[str, float]]: - """Sample one action from the masked policy and report statistics.""" - started_at = time.perf_counter() - logits, value, _ = self.policy_value(obs) - self.key, sample_key = jrandom.split(self.key) - index = int(jrandom.categorical(sample_key, jnp.asarray(logits))) - height, width = obs.armies.shape - action = _decode_policy_index(index, height, width) - elapsed_ms = (time.perf_counter() - started_at) * 1000 - return action, { - "value": value, - "elapsed_ms": elapsed_ms, - } - - -# --------------------------------------------------------------------------- -# Stdio driver -# --------------------------------------------------------------------------- - - -def _read_grid(stream, height: int): - return [[int(value) for value in stream.readline().split()] for _ in range(height)] - - -def parse_args(): - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("checkpoint", help="Equinox policy checkpoint") - parser.add_argument("--seed", type=int, default=0) - return parser.parse_args() - - -def main(): - args = parse_args() - handshake = sys.stdin.readline() - if not handshake: - return - player_index, height, width = (int(value) for value in handshake.split()) - - network = PolicyValueNetwork(jrandom.PRNGKey(args.seed)) - network = eqx.tree_deserialise_leaves(args.checkpoint, network) - agent = PolicyAgent(network, seed=args.seed) - agent.warmup(height, width) - print(f"[competition] warmup complete for {height}x{width}", file=sys.stderr, flush=True) - - while True: - scalar_line = sys.stdin.readline() - if not scalar_line: - return - timestep, own_land, own_army, opponent_land, opponent_army = ( - int(value) for value in scalar_line.split() - ) - type_grid = _read_grid(sys.stdin, height) - owner_grid = _read_grid(sys.stdin, height) - army_grid = _read_grid(sys.stdin, height) - observation = observation_from_wire( - timestep, - own_land, - own_army, - opponent_land, - opponent_army, - type_grid, - owner_grid, - army_grid, - ) - - started_at = time.perf_counter() - action, stats = agent.select(observation) - elapsed_ms = (time.perf_counter() - started_at) * 1000 - print( - f"[competition] turn={timestep} value={stats['value']:+.3f} " - f"elapsed_ms={elapsed_ms:.1f}", - file=sys.stderr, - flush=True, - ) - print(" ".join(str(int(value)) for value in action), flush=True) - - -if __name__ == "__main__": - main() diff --git a/competition/model.eqx b/competition/model.eqx deleted file mode 100644 index 549c7ad..0000000 Binary files a/competition/model.eqx and /dev/null differ diff --git a/competition/requirements-vendor.txt b/competition/requirements-vendor.txt deleted file mode 100644 index 5df83a1..0000000 --- a/competition/requirements-vendor.txt +++ /dev/null @@ -1,3 +0,0 @@ -equinox==0.13.8 -jaxtyping==0.3.11 -wadler-lindig==0.1.7 diff --git a/competition/run.sh b/competition/run.sh deleted file mode 100755 index 6dc19c2..0000000 --- a/competition/run.sh +++ /dev/null @@ -1,8 +0,0 @@ -#!/usr/bin/env bash -set -euo pipefail -# Competition entrypoint. The sandbox runs `run.sh` from the submission root; -# cd to our own directory so relative paths work regardless of the cwd. -# `python` is the sandbox's CPython 3.12 with jax/numpy pre-installed; equinox -# and its deps are installed by build.sh from wheelhouse/. -cd "$(dirname "$0")" -exec python -u competition.py model.eqx diff --git a/competition/wheelhouse/equinox-0.13.8-py3-none-any.whl b/competition/wheelhouse/equinox-0.13.8-py3-none-any.whl deleted file mode 100644 index 8b4a9ca..0000000 Binary files a/competition/wheelhouse/equinox-0.13.8-py3-none-any.whl and /dev/null differ diff --git a/competition/wheelhouse/jaxtyping-0.3.11-py3-none-any.whl b/competition/wheelhouse/jaxtyping-0.3.11-py3-none-any.whl deleted file mode 100644 index dadbbde..0000000 Binary files a/competition/wheelhouse/jaxtyping-0.3.11-py3-none-any.whl and /dev/null differ diff --git a/competition/wheelhouse/wadler_lindig-0.1.7-py3-none-any.whl b/competition/wheelhouse/wadler_lindig-0.1.7-py3-none-any.whl deleted file mode 100644 index 2e24523..0000000 Binary files a/competition/wheelhouse/wadler_lindig-0.1.7-py3-none-any.whl and /dev/null differ diff --git a/ppo_model.eqx b/ppo_model.eqx index 549c7ad..ff41fc0 100644 Binary files a/ppo_model.eqx and b/ppo_model.eqx differ diff --git a/ppo_model.eqx.3 b/ppo_model.eqx.3 deleted file mode 100644 index f30bc9d..0000000 Binary files a/ppo_model.eqx.3 and /dev/null differ diff --git a/ppo_model.eqx.4 b/ppo_model.eqx.4 deleted file mode 100644 index 40891b5..0000000 Binary files a/ppo_model.eqx.4 and /dev/null differ diff --git a/ppo_model.eqx.5 b/ppo_model.eqx.5 deleted file mode 100644 index 7af97e3..0000000 Binary files a/ppo_model.eqx.5 and /dev/null differ diff --git a/ppo_model.eqx.6 b/ppo_model.eqx.6 deleted file mode 100644 index 4bcf7cb..0000000 Binary files a/ppo_model.eqx.6 and /dev/null differ diff --git a/ppo_model.eqx.7 b/ppo_model.eqx.7 deleted file mode 100644 index 1e5d698..0000000 Binary files a/ppo_model.eqx.7 and /dev/null differ diff --git a/ppo_model_h100.eqx b/ppo_model_h100.eqx deleted file mode 100644 index 33f3255..0000000 Binary files a/ppo_model_h100.eqx and /dev/null differ diff --git a/ppo_model_h100.eqx.1 b/ppo_model_h100.eqx.1 deleted file mode 100644 index 441aaf7..0000000 Binary files a/ppo_model_h100.eqx.1 and /dev/null differ diff --git a/ppo_model_h100.eqx.2 b/ppo_model_h100.eqx.2 deleted file mode 100644 index 3509540..0000000 Binary files a/ppo_model_h100.eqx.2 and /dev/null differ diff --git a/ppo_model_h100.eqx.3 b/ppo_model_h100.eqx.3 deleted file mode 100644 index 597ac33..0000000 Binary files a/ppo_model_h100.eqx.3 and /dev/null differ diff --git a/ppo_model_h100.eqx.4 b/ppo_model_h100.eqx.4 deleted file mode 100644 index 597ac33..0000000 Binary files a/ppo_model_h100.eqx.4 and /dev/null differ diff --git a/pyproject.toml b/pyproject.toml index 14d8040..3a9c92d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -5,20 +5,12 @@ requires-python = ">=3.13" dependencies = [ "equinox>=0.13.8", "generals-bots", - "jax==0.11.0", + "jax[cuda]>=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" } diff --git a/requirements-competition.txt b/requirements-competition.txt deleted file mode 100644 index bf814de..0000000 --- a/requirements-competition.txt +++ /dev/null @@ -1,11 +0,0 @@ -numpy==2.4.6 -scipy==1.18.0 -pandas==3.0.5 -scikit-learn==1.9.0 -# On Linux, the default PyPI wheel bundles CUDA. For a CPU-only build, uncomment: -jax==0.11.0 -numba==0.66.0 -networkx==3.6.1 -safetensors==0.8.0 -gymnasium==1.3.0 -# torch==2.13.0 diff --git a/requirements-vendor.txt b/requirements-vendor.txt deleted file mode 100644 index 5df83a1..0000000 --- a/requirements-vendor.txt +++ /dev/null @@ -1,3 +0,0 @@ -equinox==0.13.8 -jaxtyping==0.3.11 -wadler-lindig==0.1.7 diff --git a/requirements.txt b/requirements.txt deleted file mode 100644 index 2b12fbc..0000000 --- a/requirements.txt +++ /dev/null @@ -1,81 +0,0 @@ -# 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 diff --git a/scripts/competition.py b/scripts/competition.py deleted file mode 100644 index 2b1e67b..0000000 --- a/scripts/competition.py +++ /dev/null @@ -1,342 +0,0 @@ -"""Competition stdio agent: direct policy sampling over a trained equinox checkpoint. - -Self-contained competition bot. It speaks the wire protocol in -`generals-bots/competition/protocol.py` (handshake, then one observation frame -per turn, one action line per turn, EOF on stdin = game over) and, on each -turn, samples an action from the policy network exactly as the model is used -during training/evaluation — no tree search. - -Unlike `scripts/mcts_agent.py`, this module does **not** import the -`general_bots_training` training package. It inlines the policy-value network -and the observation encoding so the script can be dropped into the competition -sandbox with just the engine (`generals`), `equinox`, `jax`, and `numpy` -available. - -This is the search-free counterpart to `scripts/competition_puct.py`: it shares -the same network and observation encoding, but selects an action by sampling -from the masked policy (as in `PolicyValueNetwork.__call__`) instead of running -particle PUCT. - -Run it directly through the bundled matchup driver, e.g.: - - PYTHONPATH=src:generals-bots .venv/bin/python \ - generals-bots/competition/matchup.py \ - scripts/competition.py --checkpoint ppo_model.eqx \ - generals-bots/competition/agents/expander_python/run.sh \ - --mode competition - -or wrap it in a `run.sh` that passes the checkpoint path. -""" - -import argparse -import os -import sys -import time - -import equinox as eqx -import jax -import jax.numpy as jnp -import jax.random as jrandom -import numpy as np -from generals.core import game -from generals.core.action import compute_valid_move_mask -from generals.core.game import GameState -from generals.core.observation import Observation -from generals.modifiers import build_castles, deathtouch - -os.environ.setdefault("JAX_PLATFORMS", "cpu") - -PASS_ACTION = np.array([1, 0, 0, 0, 0], dtype=np.int32) - - -# --------------------------------------------------------------------------- -# Policy-value network (inlined from general_bots_training.network) -# --------------------------------------------------------------------------- - - -def obs_to_tensor(obs: Observation) -> jnp.ndarray: - """Encode an Observation into a (C, H, W) float32 tensor for the network. - - Armies and army-counts are log-normalized; land counts are divided by the - number of cells; scalar values are broadcast to spatial planes so a plain - conv stack can consume them. Grid-size agnostic. - """ - H, W = obs.armies.shape - armies = jnp.log1p(obs.armies.astype(jnp.float32)) / jnp.log(50.0) - - def bcast(scalar): - return jnp.broadcast_to(scalar.astype(jnp.float32), (H, W)) - - own_land = bcast(obs.owned_land_count / (H * W)) - own_army = bcast(jnp.log1p(obs.owned_army_count.astype(jnp.float32)) / jnp.log(50.0)) - opp_land = bcast(obs.opponent_land_count / (H * W)) - opp_army = bcast(jnp.log1p(obs.opponent_army_count.astype(jnp.float32)) / jnp.log(50.0)) - timestep = bcast(jnp.log1p(obs.timestep.astype(jnp.float32)) / jnp.log(1201.0)) - - return jnp.stack( - [ - armies, - obs.generals.astype(jnp.float32), - obs.castles.astype(jnp.float32), - obs.mountains.astype(jnp.float32), - obs.neutral_cells.astype(jnp.float32), - obs.owned_cells.astype(jnp.float32), - obs.opponent_cells.astype(jnp.float32), - obs.fog_cells.astype(jnp.float32), - obs.structures_in_fog.astype(jnp.float32), - own_land, - own_army, - opp_land, - opp_army, - timestep, - ], - axis=0, - ) - - -class PolicyValueNetwork(eqx.Module): - """Conv policy-value network. - - Action layout: 4 full-move directions and 4 half-move (split) directions - per source cell, plus one global pass action. Invalid moves are masked to - -1e9; pass is always available. - - The value head uses global average pooling over the spatial grid, so the - network works for any board size without reshaping linear layers. - """ - - conv1: eqx.nn.Conv2d - conv2: eqx.nn.Conv2d - conv3: eqx.nn.Conv2d - conv4: eqx.nn.Conv2d - policy_conv: eqx.nn.Conv2d - value_conv: eqx.nn.Conv2d - value_linear1: eqx.nn.Linear - value_linear2: eqx.nn.Linear - - def __init__(self, key, in_channels: int = 14, channels=(32, 32, 32, 16)): - keys = jrandom.split(key, 8) - self.conv1 = eqx.nn.Conv2d(in_channels, channels[0], kernel_size=3, padding=1, key=keys[0]) - self.conv2 = eqx.nn.Conv2d(channels[0], channels[1], kernel_size=3, padding=1, key=keys[1]) - self.conv3 = eqx.nn.Conv2d(channels[1], channels[2], kernel_size=3, padding=1, key=keys[2]) - self.conv4 = eqx.nn.Conv2d(channels[2], channels[3], kernel_size=3, padding=1, key=keys[3]) - # 9 = 4 dirs (full) + 4 dirs (half) + 1 pass - self.policy_conv = eqx.nn.Conv2d(channels[3], 9, kernel_size=1, key=keys[4]) - self.value_conv = eqx.nn.Conv2d(channels[3], 4, kernel_size=1, key=keys[5]) - self.value_linear1 = eqx.nn.Linear(4, 64, key=keys[6]) - self.value_linear2 = eqx.nn.Linear(64, 1, key=keys[7]) - - def _features(self, obs): - x = jax.nn.relu(self.conv1(obs)) - x = jax.nn.relu(self.conv2(x)) - x = jax.nn.relu(self.conv3(x)) - x = jax.nn.relu(self.conv4(x)) - return x - - def _value_from_features(self, feat): - v = jax.nn.relu(self.value_conv(feat)) # (4, H, W) - v = v.mean(axis=(1, 2)) # (4,) global average pool - v = jax.nn.relu(self.value_linear1(v)) - return self.value_linear2(v)[0] - - def policy_value(self, obs, mask): - """Return deterministic masked policy logits and the critic value.""" - features = self._features(obs) - return self._policy_logits(features, mask), self._value_from_features(features) - - def _policy_logits(self, feat, mask): - """Return 8*H*W move logits plus one global pass logit.""" - logits = self.policy_conv(feat) # (9, H, W) - mask_t = jnp.transpose(mask, (2, 0, 1)) # (4, H, W) - penalty = (1.0 - mask_t) * -1e9 - move_penalty = jnp.concatenate([penalty, penalty], axis=0) - move_logits = (logits[:8] + move_penalty).reshape(-1) - pass_logit = jnp.mean(logits[8])[None] - return jnp.concatenate([move_logits, pass_logit]) - - -# --------------------------------------------------------------------------- -# Action selection (matches training/evaluation rollout) -# --------------------------------------------------------------------------- - - -@jax.jit -def competition_move_mask(obs: Observation) -> jnp.ndarray: - """Legal move mask, identical to the one used during training/evaluation.""" - return compute_valid_move_mask(obs.armies, obs.owned_cells, obs.mountains) - - -def _decode_policy_index(index: int, height: int, width: int) -> np.ndarray: - cells = height * width - if index == 8 * cells: - return PASS_ACTION.copy() - encoded_direction, position = divmod(index, cells) - row, col = divmod(position, width) - split = int(encoded_direction >= 4) - direction = encoded_direction - 4 if split else encoded_direction - return np.array([0, row, col, direction, split], dtype=np.int32) - - -def observation_from_wire( - timestep: int, - owned_land_count: int, - owned_army_count: int, - opponent_land_count: int, - opponent_army_count: int, - type_grid, - owner_grid, - army_grid, -) -> Observation: - """Convert a competition wire frame into the network's Observation type.""" - types = jnp.asarray(type_grid, dtype=jnp.int32) - owners = jnp.asarray(owner_grid, dtype=jnp.int32) - armies = jnp.asarray(army_grid, dtype=jnp.int32) - visible = (types != 0) & (types != 5) - return Observation( - armies=armies, - generals=types == 4, - castles=types == 3, - mountains=types == 2, - neutral_cells=visible & (owners == 0) & (types != 2), - owned_cells=owners == 1, - opponent_cells=owners == 2, - fog_cells=types == 0, - structures_in_fog=types == 5, - owned_land_count=jnp.int32(owned_land_count), - owned_army_count=jnp.int32(owned_army_count), - opponent_land_count=jnp.int32(opponent_land_count), - opponent_army_count=jnp.int32(opponent_army_count), - timestep=jnp.int32(timestep), - ) - - -@jax.jit -def competition_step(state: GameState, actions: jnp.ndarray): - state, actions = build_castles.apply_build_actions(state, actions) - return deathtouch.step(state, actions, turn=800) - - -class PolicyAgent: - """Direct policy agent that samples actions as in training/evaluation. - - On each turn it computes the masked policy logits over all legal moves and - samples one action with `jrandom.categorical`, exactly mirroring - `PolicyValueNetwork.__call__` used during rollout collection. - """ - - def __init__( - self, - network: PolicyValueNetwork, - seed: int = 0, - ): - self.network = network - self.key = jrandom.PRNGKey(seed) - self._infer = eqx.filter_jit( - lambda network, tensor, mask: network.policy_value(tensor, mask) - ) - - def warmup(self, height: int, width: int) -> None: - """Compile shape-dependent network and competition transition kernels.""" - grid = jnp.zeros((height, width), dtype=jnp.int32) - grid = grid.at[0, 0].set(1).at[height - 1, width - 1].set(2) - state = game.create_initial_state(grid) - state = state._replace( - armies=state.armies.at[0, 0].set(100).at[height - 1, width - 1].set(100), - time=jnp.int32(100), - ) - obs = game.get_observation(state, 0) - self.policy_value(obs) - actions = jnp.stack([jnp.asarray(PASS_ACTION), jnp.asarray(PASS_ACTION)]) - warmed_state, _ = competition_step(state, actions) - jax.block_until_ready(warmed_state) - - def policy_value(self, obs: Observation): - mask = competition_move_mask(obs) - logits, value = self._infer(self.network, obs_to_tensor(obs), mask) - return np.asarray(logits), float(value), mask - - def select(self, obs: Observation) -> tuple[np.ndarray, dict[str, float]]: - """Sample one action from the masked policy and report statistics.""" - started_at = time.perf_counter() - logits, value, _ = self.policy_value(obs) - self.key, sample_key = jrandom.split(self.key) - index = int(jrandom.categorical(sample_key, jnp.asarray(logits))) - height, width = obs.armies.shape - action = _decode_policy_index(index, height, width) - elapsed_ms = (time.perf_counter() - started_at) * 1000 - return action, { - "value": value, - "elapsed_ms": elapsed_ms, - } - - -# --------------------------------------------------------------------------- -# Stdio driver -# --------------------------------------------------------------------------- - - -def _read_grid(stream, height: int): - return [[int(value) for value in stream.readline().split()] for _ in range(height)] - - -def parse_args(): - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("checkpoint", help="Equinox policy checkpoint") - # Accepted for compatibility with the original run.sh; the direct policy - # agent does no search, so these are ignored. - parser.add_argument("--time-budget-ms", type=float, default=125.0) - parser.add_argument("--max-simulations", type=int, default=128) - parser.add_argument("--rollout-depth", type=int, default=2) - parser.add_argument("--top-k", type=int, default=20) - parser.add_argument("--seed", type=int, default=0) - return parser.parse_args() - - -def main(): - args = parse_args() - handshake = sys.stdin.readline() - if not handshake: - return - player_index, height, width = (int(value) for value in handshake.split()) - - network = PolicyValueNetwork(jrandom.PRNGKey(args.seed)) - network = eqx.tree_deserialise_leaves(args.checkpoint, network) - agent = PolicyAgent(network, seed=args.seed) - agent.warmup(height, width) - print(f"[competition] warmup complete for {height}x{width}", file=sys.stderr, flush=True) - - while True: - scalar_line = sys.stdin.readline() - if not scalar_line: - return - timestep, own_land, own_army, opponent_land, opponent_army = ( - int(value) for value in scalar_line.split() - ) - type_grid = _read_grid(sys.stdin, height) - owner_grid = _read_grid(sys.stdin, height) - army_grid = _read_grid(sys.stdin, height) - observation = observation_from_wire( - timestep, - own_land, - own_army, - opponent_land, - opponent_army, - type_grid, - owner_grid, - army_grid, - ) - - started_at = time.perf_counter() - action, stats = agent.select(observation) - elapsed_ms = (time.perf_counter() - started_at) * 1000 - print( - f"[competition] turn={timestep} value={stats['value']:+.3f} " - f"elapsed_ms={elapsed_ms:.1f}", - file=sys.stderr, - flush=True, - ) - print(" ".join(str(int(value)) for value in action), flush=True) - - -if __name__ == "__main__": - main() diff --git a/scripts/competition_puct.py b/scripts/competition_puct.py deleted file mode 100644 index e398c97..0000000 --- a/scripts/competition_puct.py +++ /dev/null @@ -1,552 +0,0 @@ -"""Competition stdio agent: particle PUCT over a trained equinox checkpoint. - -Self-contained competition bot. It speaks the wire protocol in -`generals-bots/competition/protocol.py` (handshake, then one observation frame -per turn, one action line per turn, EOF on stdin = game over) and runs -deadline-bounded root PUCT using only the perspective-relative wire -observation. - -Unlike `scripts/mcts_agent.py`, this module does **not** import the -`general_bots_training` training package. It inlines the policy-value network, -the observation encoding, and the particle search so the script can be dropped -into the competition sandbox with just the engine (`generals`), `equinox`, -`jax`, and `numpy` available. - -Run it directly through the bundled matchup driver, e.g.: - - PYTHONPATH=src:generals-bots .venv/bin/python \ - generals-bots/competition/matchup.py \ - scripts/competition.py --checkpoint ppo_model.eqx \ - generals-bots/competition/agents/expander_python/run.sh \ - --mode competition - -or wrap it in a `run.sh` that passes the checkpoint path. -""" - -import argparse -import os -import sys -import time -from dataclasses import dataclass - -import equinox as eqx -import jax -import jax.numpy as jnp -import jax.random as jrandom -import numpy as np -from generals.core import game -from generals.core.action import compute_valid_move_mask -from generals.core.game import GameState -from generals.core.observation import Observation -from generals.modifiers import build_castles, deathtouch - -os.environ.setdefault("JAX_PLATFORMS", "cpu") - -PASS_ACTION = np.array([1, 0, 0, 0, 0], dtype=np.int32) - -# Search/agent defaults, previously exposed as CLI flags. Kept as constants so -# the script can be launched by a bare `run.sh` with no extra arguments. -TIME_BUDGET_MS = 125.0 -MAX_SIMULATIONS = 128 -ROLLOUT_DEPTH = 2 -TOP_K = 20 -SEED = 0 - - -# --------------------------------------------------------------------------- -# Policy-value network (inlined from general_bots_training.network) -# --------------------------------------------------------------------------- - - -def obs_to_tensor(obs: Observation) -> jnp.ndarray: - """Encode an Observation into a (C, H, W) float32 tensor for the network. - - Armies and army-counts are log-normalized; land counts are divided by the - number of cells; scalar values are broadcast to spatial planes so a plain - conv stack can consume them. Grid-size agnostic. - """ - H, W = obs.armies.shape - armies = jnp.log1p(obs.armies.astype(jnp.float32)) / jnp.log(50.0) - - def bcast(scalar): - return jnp.broadcast_to(scalar.astype(jnp.float32), (H, W)) - - own_land = bcast(obs.owned_land_count / (H * W)) - own_army = bcast(jnp.log1p(obs.owned_army_count.astype(jnp.float32)) / jnp.log(50.0)) - opp_land = bcast(obs.opponent_land_count / (H * W)) - opp_army = bcast(jnp.log1p(obs.opponent_army_count.astype(jnp.float32)) / jnp.log(50.0)) - timestep = bcast(jnp.log1p(obs.timestep.astype(jnp.float32)) / jnp.log(1201.0)) - - return jnp.stack( - [ - armies, - obs.generals.astype(jnp.float32), - obs.castles.astype(jnp.float32), - obs.mountains.astype(jnp.float32), - obs.neutral_cells.astype(jnp.float32), - obs.owned_cells.astype(jnp.float32), - obs.opponent_cells.astype(jnp.float32), - obs.fog_cells.astype(jnp.float32), - obs.structures_in_fog.astype(jnp.float32), - own_land, - own_army, - opp_land, - opp_army, - timestep, - ], - axis=0, - ) - - -class PolicyValueNetwork(eqx.Module): - """Conv policy-value network. - - Action layout: 4 full-move directions and 4 half-move (split) directions - per source cell, plus one global pass action. Invalid moves are masked to - -1e9; pass is always available. - - The value head uses global average pooling over the spatial grid, so the - network works for any board size without reshaping linear layers. - """ - - conv1: eqx.nn.Conv2d - conv2: eqx.nn.Conv2d - conv3: eqx.nn.Conv2d - conv4: eqx.nn.Conv2d - policy_conv: eqx.nn.Conv2d - value_conv: eqx.nn.Conv2d - value_linear1: eqx.nn.Linear - value_linear2: eqx.nn.Linear - - def __init__(self, key, in_channels: int = 14, channels=(32, 32, 32, 16)): - keys = jrandom.split(key, 8) - self.conv1 = eqx.nn.Conv2d(in_channels, channels[0], kernel_size=3, padding=1, key=keys[0]) - self.conv2 = eqx.nn.Conv2d(channels[0], channels[1], kernel_size=3, padding=1, key=keys[1]) - self.conv3 = eqx.nn.Conv2d(channels[1], channels[2], kernel_size=3, padding=1, key=keys[2]) - self.conv4 = eqx.nn.Conv2d(channels[2], channels[3], kernel_size=3, padding=1, key=keys[3]) - # 9 = 4 dirs (full) + 4 dirs (half) + 1 pass - self.policy_conv = eqx.nn.Conv2d(channels[3], 9, kernel_size=1, key=keys[4]) - self.value_conv = eqx.nn.Conv2d(channels[3], 4, kernel_size=1, key=keys[5]) - self.value_linear1 = eqx.nn.Linear(4, 64, key=keys[6]) - self.value_linear2 = eqx.nn.Linear(64, 1, key=keys[7]) - - def _features(self, obs): - x = jax.nn.relu(self.conv1(obs)) - x = jax.nn.relu(self.conv2(x)) - x = jax.nn.relu(self.conv3(x)) - x = jax.nn.relu(self.conv4(x)) - return x - - def _value_from_features(self, feat): - v = jax.nn.relu(self.value_conv(feat)) # (4, H, W) - v = v.mean(axis=(1, 2)) # (4,) global average pool - v = jax.nn.relu(self.value_linear1(v)) - return self.value_linear2(v)[0] - - def policy_value(self, obs, mask): - """Return deterministic masked policy logits and the critic value.""" - features = self._features(obs) - return self._policy_logits(features, mask), self._value_from_features(features) - - def _policy_logits(self, feat, mask): - """Return 8*H*W move logits plus one global pass logit.""" - logits = self.policy_conv(feat) # (9, H, W) - mask_t = jnp.transpose(mask, (2, 0, 1)) # (4, H, W) - penalty = (1.0 - mask_t) * -1e9 - move_penalty = jnp.concatenate([penalty, penalty], axis=0) - move_logits = (logits[:8] + move_penalty).reshape(-1) - pass_logit = jnp.mean(logits[8])[None] - return jnp.concatenate([move_logits, pass_logit]) - - -# --------------------------------------------------------------------------- -# Particle PUCT search (inlined from general_bots_training.mcts) -# --------------------------------------------------------------------------- - - -@dataclass(frozen=True) -class MCTSConfig: - time_budget_ms: float = 100.0 - max_simulations: int = 128 - rollout_depth: int = 2 - top_k: int = 20 - opponent_top_k: int = 12 - max_build_actions: int = 3 - c_puct: float = 1.5 - value_scale: float = 5.0 - - -@jax.jit -def competition_move_mask(obs: Observation) -> jnp.ndarray: - """Legal move mask that conservatively blocks unresolved fog structures.""" - blocked = obs.mountains | obs.structures_in_fog - return compute_valid_move_mask(obs.armies, obs.owned_cells, blocked) - - -@jax.jit -def build_cost_grid(obs: Observation) -> jnp.ndarray: - """Exact own-castle build costs derivable from a fog observation.""" - structures = ((obs.castles | obs.generals) & obs.owned_cells).astype(jnp.int32) - height, width = structures.shape - radius = 6 - padded = jnp.pad(structures, radius) - costs = jnp.full((height, width), 35, dtype=jnp.int32) - for row_offset in range(-radius, radius + 1): - for col_offset in range(-radius, radius + 1): - surcharge = 14 - 2 * (abs(row_offset) + abs(col_offset)) - if surcharge > 0: - shifted = padded[ - radius + row_offset : radius + row_offset + height, - radius + col_offset : radius + col_offset + width, - ] - costs = costs + surcharge * shifted - return costs - - -@jax.jit -def valid_build_mask(obs: Observation) -> jnp.ndarray: - costs = build_cost_grid(obs) - plain_owned = obs.owned_cells & ~obs.generals & ~obs.castles - return plain_owned & (obs.armies >= costs) - - -def _decode_policy_index(index: int, height: int, width: int) -> np.ndarray: - cells = height * width - if index == 8 * cells: - return PASS_ACTION.copy() - encoded_direction, position = divmod(index, cells) - row, col = divmod(position, width) - split = int(encoded_direction >= 4) - direction = encoded_direction - 4 if split else encoded_direction - return np.array([0, row, col, direction, split], dtype=np.int32) - - -def observation_from_wire( - timestep: int, - owned_land_count: int, - owned_army_count: int, - opponent_land_count: int, - opponent_army_count: int, - type_grid, - owner_grid, - army_grid, -) -> Observation: - """Convert a competition wire frame into the network's Observation type.""" - types = jnp.asarray(type_grid, dtype=jnp.int32) - owners = jnp.asarray(owner_grid, dtype=jnp.int32) - armies = jnp.asarray(army_grid, dtype=jnp.int32) - visible = (types != 0) & (types != 5) - return Observation( - armies=armies, - generals=types == 4, - castles=types == 3, - mountains=types == 2, - neutral_cells=visible & (owners == 0) & (types != 2), - owned_cells=owners == 1, - opponent_cells=owners == 2, - fog_cells=types == 0, - structures_in_fog=types == 5, - owned_land_count=jnp.int32(owned_land_count), - owned_army_count=jnp.int32(owned_army_count), - opponent_land_count=jnp.int32(opponent_land_count), - opponent_army_count=jnp.int32(opponent_army_count), - timestep=jnp.int32(timestep), - ) - - -def sample_determinization( - obs: Observation, player_index: int, rng: np.random.Generator -) -> GameState: - """Sample a conservative hidden state consistent with visible cells/totals.""" - armies = np.asarray(obs.armies, dtype=np.int32).copy() - owned = np.asarray(obs.owned_cells, dtype=bool) - visible_opponent = np.asarray(obs.opponent_cells, dtype=bool) - mountains = np.asarray(obs.mountains | obs.structures_in_fog, dtype=bool) - castles = np.asarray(obs.castles, dtype=bool) - generals = np.asarray(obs.generals, dtype=bool).copy() - fog = np.asarray(obs.fog_cells, dtype=bool) & ~mountains - - opponent = visible_opponent.copy() - hidden_candidates = np.argwhere(fog & ~owned) - missing_land = max(0, int(obs.opponent_land_count) - int(opponent.sum())) - if len(hidden_candidates): - chosen = hidden_candidates[ - rng.choice( - len(hidden_candidates), - size=min(missing_land, len(hidden_candidates)), - replace=False, - ) - ] - opponent[chosen[:, 0], chosen[:, 1]] = True - - visible_enemy_general = generals & opponent - if not visible_enemy_general.any(): - candidates = np.argwhere((opponent | fog) & ~owned & ~mountains) - if len(candidates): - row, col = candidates[rng.integers(len(candidates))] - generals[row, col] = True - opponent[row, col] = True - - hidden_opponent = opponent & ~visible_opponent - remaining_army = max(0, int(obs.opponent_army_count) - int(armies[visible_opponent].sum())) - hidden_cells = np.argwhere(hidden_opponent) - if len(hidden_cells): - base = min(remaining_army, len(hidden_cells)) - armies[hidden_opponent] = 0 - armies[hidden_cells[:base, 0], hidden_cells[:base, 1]] = 1 - remaining_army -= base - if remaining_army: - allocations = rng.multinomial( - remaining_army, np.full(len(hidden_cells), 1 / len(hidden_cells)) - ) - armies[hidden_cells[:, 0], hidden_cells[:, 1]] += allocations.astype(np.int32) - - ownership_relative = np.stack([owned, opponent]) - if player_index == 1: - ownership = ownership_relative[::-1] - else: - ownership = ownership_relative - passable = ~mountains - neutral = passable & ~ownership[0] & ~ownership[1] - - general_positions = [] - for absolute_player in range(2): - positions = np.argwhere(generals & ownership[absolute_player]) - if len(positions): - general_positions.append(positions[0]) - else: - fallback = np.argwhere(ownership[absolute_player] & passable) - general_positions.append(fallback[0] if len(fallback) else np.array([0, 0])) - - return GameState( - armies=jnp.asarray(armies), - ownership=jnp.asarray(ownership), - ownership_neutral=jnp.asarray(neutral), - generals=jnp.asarray(generals), - castles=jnp.asarray(castles), - mountains=jnp.asarray(mountains), - passable=jnp.asarray(passable), - general_positions=jnp.asarray(general_positions, dtype=jnp.int32), - time=jnp.asarray(obs.timestep, dtype=jnp.int32), - winner=jnp.int32(-1), - pool_idx=jnp.int32(0), - ) - - -@jax.jit -def competition_step(state: GameState, actions: jnp.ndarray): - state, actions = build_castles.apply_build_actions(state, actions) - return deathtouch.step(state, actions, turn=800) - - -class ParticlePUCT: - """Deadline-bounded root PUCT with policy-guided simultaneous rollouts.""" - - def __init__( - self, - network: PolicyValueNetwork, - player_index: int, - config: MCTSConfig = MCTSConfig(), - seed: int = 0, - ): - self.network = network - self.player_index = player_index - self.config = config - self.rng = np.random.default_rng(seed) - self._infer = eqx.filter_jit( - lambda network, tensor, mask: network.policy_value(tensor, mask) - ) - - def warmup(self, height: int, width: int) -> None: - """Compile shape-dependent network and competition transition kernels.""" - grid = jnp.zeros((height, width), dtype=jnp.int32) - grid = grid.at[0, 0].set(1).at[height - 1, width - 1].set(2) - state = game.create_initial_state(grid) - state = state._replace( - armies=state.armies.at[0, 0].set(100).at[height - 1, width - 1].set(100), - time=jnp.int32(100), - ) - obs = game.get_observation(state, self.player_index) - self.policy_value(obs) - actions = jnp.stack([jnp.asarray(PASS_ACTION), jnp.asarray(PASS_ACTION)]) - warmed_state, _ = competition_step(state, actions) - jax.block_until_ready(warmed_state) - candidates, _ = self.candidates(obs) - self._simulate(obs, candidates[0]) - self._simulate(obs, candidates[min(1, len(candidates) - 1)]) - - def policy_value(self, obs: Observation): - mask = competition_move_mask(obs) - logits, value = self._infer(self.network, obs_to_tensor(obs), mask) - return np.asarray(logits), float(value), mask - - def candidates(self, obs: Observation, top_k: int | None = None): - logits, _, _ = self.policy_value(obs) - height, width = obs.armies.shape - valid_indices = np.flatnonzero(logits > -1e8) - count = min(top_k or self.config.top_k, len(valid_indices)) - ranked = np.argsort(logits[valid_indices])[::-1][:count] - indices = valid_indices[ranked] - selected_logits = logits[indices] - selected_logits = selected_logits - selected_logits.max() - priors = np.exp(selected_logits) - actions = [_decode_policy_index(int(index), height, width) for index in indices] - - build_mask = np.asarray(valid_build_mask(obs)) - build_positions = np.argwhere(build_mask) - if len(build_positions): - costs = np.asarray(build_cost_grid(obs)) - armies = np.asarray(obs.armies) - scores = np.array([armies[r, c] - costs[r, c] for r, c in build_positions]) - order = np.argsort(scores)[::-1][: self.config.max_build_actions] - build_prior = max(float(priors.sum()) * 0.05, 1e-3) - for position_index in order: - row, col = build_positions[position_index] - actions.append(np.array([2, row, col, 0, 0], dtype=np.int32)) - priors = np.append(priors, build_prior) - - priors = priors / priors.sum() - return actions, priors - - def _sample_policy_action(self, obs: Observation, top_k: int) -> np.ndarray: - actions, priors = self.candidates(obs, top_k) - return actions[int(self.rng.choice(len(actions), p=priors))] - - def _leaf_value(self, state: GameState, info) -> float: - if bool(info.is_done): - winner = int(info.winner) - return 0.0 if winner < 0 else (1.0 if winner == self.player_index else -1.0) - obs = game.get_observation(state, self.player_index) - _, value, _ = self.policy_value(obs) - return float(np.tanh(value / self.config.value_scale)) - - def _simulate(self, root_obs: Observation, root_action: np.ndarray) -> float: - state = sample_determinization(root_obs, self.player_index, self.rng) - info = game.get_info(state) - our_action = root_action - for _ in range(self.config.rollout_depth): - opponent_index = 1 - self.player_index - opponent_obs = game.get_observation(state, opponent_index) - opponent_action = self._sample_policy_action(opponent_obs, self.config.opponent_top_k) - joint_actions = [None, None] - joint_actions[self.player_index] = our_action - joint_actions[opponent_index] = opponent_action - state, info = competition_step(state, jnp.asarray(joint_actions, dtype=jnp.int32)) - jax.block_until_ready(state) - if bool(info.is_done): - break - our_obs = game.get_observation(state, self.player_index) - our_action = self._sample_policy_action(our_obs, self.config.top_k) - return self._leaf_value(state, info) - - def search(self, obs: Observation) -> tuple[np.ndarray, dict[str, float]]: - started_at = time.perf_counter() - deadline = started_at + self.config.time_budget_ms / 1000.0 - actions, priors = self.candidates(obs) - visits = np.zeros(len(actions), dtype=np.int32) - value_sums = np.zeros(len(actions), dtype=np.float64) - simulations = 0 - estimated_simulation_seconds = 0.0 - - while ( - simulations < self.config.max_simulations - and time.perf_counter() + estimated_simulation_seconds < deadline - ): - total_visits = max(1, int(visits.sum())) - q_values = np.divide( - value_sums, - visits, - out=np.zeros_like(value_sums), - where=visits > 0, - ) - scores = q_values + self.config.c_puct * priors * np.sqrt(total_visits) / (1 + visits) - action_index = int(np.argmax(scores)) - simulation_started_at = time.perf_counter() - value = self._simulate(obs, actions[action_index]) - simulation_elapsed = time.perf_counter() - simulation_started_at - estimated_simulation_seconds = max(estimated_simulation_seconds, simulation_elapsed) - visits[action_index] += 1 - value_sums[action_index] += value - simulations += 1 - - selected = int(np.argmax(visits)) if simulations else int(np.argmax(priors)) - return actions[selected], { - "simulations": float(simulations), - "selected_visits": float(visits[selected]), - "elapsed_budget_ms": self.config.time_budget_ms, - } - - -# --------------------------------------------------------------------------- -# Stdio driver -# --------------------------------------------------------------------------- - - -def _read_grid(stream, height: int): - return [[int(value) for value in stream.readline().split()] for _ in range(height)] - - -def parse_args(): - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("checkpoint", help="Equinox policy checkpoint") - return parser.parse_args() - - -def main(): - args = parse_args() - handshake = sys.stdin.readline() - if not handshake: - return - player_index, height, width = (int(value) for value in handshake.split()) - - network = PolicyValueNetwork(jrandom.PRNGKey(SEED)) - network = eqx.tree_deserialise_leaves(args.checkpoint, network) - search = ParticlePUCT( - network, - player_index, - MCTSConfig( - time_budget_ms=TIME_BUDGET_MS, - max_simulations=MAX_SIMULATIONS, - rollout_depth=ROLLOUT_DEPTH, - top_k=TOP_K, - ), - seed=SEED, - ) - search.warmup(height, width) - print(f"[competition] warmup complete for {height}x{width}", file=sys.stderr, flush=True) - - while True: - scalar_line = sys.stdin.readline() - if not scalar_line: - return - timestep, own_land, own_army, opponent_land, opponent_army = ( - int(value) for value in scalar_line.split() - ) - type_grid = _read_grid(sys.stdin, height) - owner_grid = _read_grid(sys.stdin, height) - army_grid = _read_grid(sys.stdin, height) - observation = observation_from_wire( - timestep, - own_land, - own_army, - opponent_land, - opponent_army, - type_grid, - owner_grid, - army_grid, - ) - - started_at = time.perf_counter() - action, stats = search.search(observation) - elapsed_ms = (time.perf_counter() - started_at) * 1000 - print( - f"[competition] turn={timestep} simulations={int(stats['simulations'])} " - f"elapsed_ms={elapsed_ms:.1f}", - file=sys.stderr, - flush=True, - ) - print(" ".join(str(int(value)) for value in action), flush=True) - - -if __name__ == "__main__": - main() diff --git a/scripts/evaluate.py b/scripts/evaluate.py deleted file mode 100644 index b9c7b91..0000000 --- a/scripts/evaluate.py +++ /dev/null @@ -1,123 +0,0 @@ -"""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=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=model:ppo_model.eqx`. Games are run in a single vmapped batch of -`num_games` parallel envs. -""" - -from dataclasses import dataclass - -import jax -import jax.numpy as jnp -import jax.random as jrandom -from generals.core import game -from generals.core.env import GeneralsEnv -from omegaconf import DictConfig, OmegaConf - -from general_bots_training.opponents import make_opponent - - -@dataclass -class EvalConfig: - grid_dims: tuple[int, int] = (21, 21) - truncation: int = 1200 - num_games: int = 100 - agent0: str = "random" - agent1: str = "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 - - -def main(config: DictConfig): - key = jrandom.PRNGKey(config.seed) - key, env_key = jrandom.split(key, 2) - - agent0 = make_opponent(config.agent0) - agent1 = make_opponent(config.agent1) - label0, label1 = config.agent0, 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()) diff --git a/scripts/run.sh b/scripts/run.sh deleted file mode 100644 index 704f9a4..0000000 --- a/scripts/run.sh +++ /dev/null @@ -1,4 +0,0 @@ -#!/usr/bin/env bash -set -euo pipefail -cd "$(dirname "$0")" -exec python -u competition.py ../ppo_model.eqx diff --git a/src/general_bots_training/__init__.py b/src/general_bots_training/__init__.py index 5c24ebc..899572c 100644 --- a/src/general_bots_training/__init__.py +++ b/src/general_bots_training/__init__.py @@ -1,6 +1,6 @@ from .mcts import MCTSConfig, ParticlePUCT from .network import PolicyValueNetwork, obs_to_tensor -from .opponents import ModelOpponent, Opponent, SelfPlayOpponent, StaticOpponent, make_opponent +from .opponents import Opponent, SelfPlayOpponent, StaticOpponent, make_opponent from .ppo import compute_advantages_and_returns, compute_gae, make_train_epoch, ppo_loss from .rollout import make_collect_rollout, make_rollout_step @@ -17,7 +17,6 @@ __all__ = [ "make_rollout_step", "make_opponent", "StaticOpponent", - "ModelOpponent", "SelfPlayOpponent", "Opponent", ] diff --git a/src/general_bots_training/__pycache__/__init__.cpython-313.pyc b/src/general_bots_training/__pycache__/__init__.cpython-313.pyc new file mode 100644 index 0000000..5b3e23f Binary files /dev/null and b/src/general_bots_training/__pycache__/__init__.cpython-313.pyc differ diff --git a/src/general_bots_training/__pycache__/mcts.cpython-313.pyc b/src/general_bots_training/__pycache__/mcts.cpython-313.pyc new file mode 100644 index 0000000..d969103 Binary files /dev/null and b/src/general_bots_training/__pycache__/mcts.cpython-313.pyc differ diff --git a/src/general_bots_training/__pycache__/network.cpython-313.pyc b/src/general_bots_training/__pycache__/network.cpython-313.pyc new file mode 100644 index 0000000..4cb8cab Binary files /dev/null and b/src/general_bots_training/__pycache__/network.cpython-313.pyc differ diff --git a/src/general_bots_training/__pycache__/opponents.cpython-313.pyc b/src/general_bots_training/__pycache__/opponents.cpython-313.pyc new file mode 100644 index 0000000..3616d0d Binary files /dev/null and b/src/general_bots_training/__pycache__/opponents.cpython-313.pyc differ diff --git a/src/general_bots_training/__pycache__/ppo.cpython-313.pyc b/src/general_bots_training/__pycache__/ppo.cpython-313.pyc new file mode 100644 index 0000000..56eef53 Binary files /dev/null and b/src/general_bots_training/__pycache__/ppo.cpython-313.pyc differ diff --git a/src/general_bots_training/__pycache__/rollout.cpython-313.pyc b/src/general_bots_training/__pycache__/rollout.cpython-313.pyc new file mode 100644 index 0000000..ad76e46 Binary files /dev/null and b/src/general_bots_training/__pycache__/rollout.cpython-313.pyc differ diff --git a/src/general_bots_training/__pycache__/train.cpython-313.pyc b/src/general_bots_training/__pycache__/train.cpython-313.pyc new file mode 100644 index 0000000..871538a Binary files /dev/null and b/src/general_bots_training/__pycache__/train.cpython-313.pyc differ diff --git a/src/general_bots_training/opponents.py b/src/general_bots_training/opponents.py index a77e136..ee5abe5 100644 --- a/src/general_bots_training/opponents.py +++ b/src/general_bots_training/opponents.py @@ -4,15 +4,10 @@ from typing import Protocol from dataclasses import dataclass -import equinox as eqx import jax.numpy as jnp -import jax.random as jrandom from generals.agents import Agent, ExpanderAgent, HunterAgent, RandomAgent -from generals.core.action import compute_valid_move_mask from generals.core.observation import Observation -from .network import PolicyValueNetwork, obs_to_tensor - class StaticOpponent(Protocol): """Opponent policy that carries no per-environment mutable state.""" @@ -25,21 +20,6 @@ class SelfPlayOpponent: """Marker selecting the current training network as player 1.""" -class ModelOpponent: - """Static opponent that plays a fixed policy loaded from an equinox checkpoint.""" - - def __init__(self, network: PolicyValueNetwork): - self._network = network - - def act(self, observation: Observation, key: jnp.ndarray) -> jnp.ndarray: - 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 - - Opponent = StaticOpponent | SelfPlayOpponent @@ -50,36 +30,15 @@ OPPONENT_TYPES: dict[str, type[Agent]] = { } -def make_opponent(spec: str) -> Opponent: - """Create an opponent strategy from a single configuration string. - - The spec is either a bare opponent name ("random", "expander", "hunter", - "self_play") or a ``type:checkpoint`` pair for a model opponent, e.g. - ``"model:ppo_model.eqx"``. - """ - name, sep, checkpoint = spec.partition(":") +def make_opponent(name: str) -> Opponent: + """Create an opponent strategy by configuration name.""" normalized_name = name.lower().replace("-", "_") if normalized_name == "self_play": return SelfPlayOpponent() - if normalized_name == "model": - if not sep: - raise ValueError( - "opponent 'model' requires a checkpoint path, e.g. 'model:ppo_model.eqx'" - ) - # The key only seeds the network structure; deserialization overwrites - # every leaf, so a fixed seed is sufficient. - network = PolicyValueNetwork(jrandom.PRNGKey(0), in_channels=14) - network = eqx.tree_deserialise_leaves(checkpoint, network) - return ModelOpponent(network) - - if sep: - choices = ", ".join([*sorted(OPPONENT_TYPES), "model", "self_play"]) - raise ValueError(f"opponent {spec!r} does not take a checkpoint; choose one of: {choices}") - try: opponent_type = OPPONENT_TYPES[normalized_name] except KeyError as error: - choices = ", ".join([*sorted(OPPONENT_TYPES), "model", "self_play"]) - raise ValueError(f"unknown opponent {spec!r}; choose one of: {choices}") from error + choices = ", ".join([*sorted(OPPONENT_TYPES), "self_play"]) + raise ValueError(f"unknown opponent {name!r}; choose one of: {choices}") from error return opponent_type() diff --git a/tests/__pycache__/test_checkpoint.cpython-313-pytest-9.1.1.pyc b/tests/__pycache__/test_checkpoint.cpython-313-pytest-9.1.1.pyc new file mode 100644 index 0000000..3aa7e31 Binary files /dev/null and b/tests/__pycache__/test_checkpoint.cpython-313-pytest-9.1.1.pyc differ diff --git a/tests/__pycache__/test_mcts.cpython-313-pytest-9.1.1.pyc b/tests/__pycache__/test_mcts.cpython-313-pytest-9.1.1.pyc new file mode 100644 index 0000000..ec155ac Binary files /dev/null and b/tests/__pycache__/test_mcts.cpython-313-pytest-9.1.1.pyc differ diff --git a/tests/__pycache__/test_opponents.cpython-313-pytest-9.1.1.pyc b/tests/__pycache__/test_opponents.cpython-313-pytest-9.1.1.pyc new file mode 100644 index 0000000..a5150c5 Binary files /dev/null and b/tests/__pycache__/test_opponents.cpython-313-pytest-9.1.1.pyc differ diff --git a/tests/__pycache__/test_self_play.cpython-313-pytest-9.1.1.pyc b/tests/__pycache__/test_self_play.cpython-313-pytest-9.1.1.pyc new file mode 100644 index 0000000..8aeda74 Binary files /dev/null and b/tests/__pycache__/test_self_play.cpython-313-pytest-9.1.1.pyc differ diff --git a/tests/__pycache__/test_train_config.cpython-313-pytest-9.1.1.pyc b/tests/__pycache__/test_train_config.cpython-313-pytest-9.1.1.pyc new file mode 100644 index 0000000..75beb00 Binary files /dev/null and b/tests/__pycache__/test_train_config.cpython-313-pytest-9.1.1.pyc differ diff --git a/tests/__pycache__/test_training.cpython-313-pytest-9.1.1.pyc b/tests/__pycache__/test_training.cpython-313-pytest-9.1.1.pyc new file mode 100644 index 0000000..0705288 Binary files /dev/null and b/tests/__pycache__/test_training.cpython-313-pytest-9.1.1.pyc differ diff --git a/tests/test_opponents.py b/tests/test_opponents.py index 362870d..e553630 100644 --- a/tests/test_opponents.py +++ b/tests/test_opponents.py @@ -1,7 +1,7 @@ import pytest from generals.agents import ExpanderAgent, HunterAgent, RandomAgent -from general_bots_training.opponents import ModelOpponent, SelfPlayOpponent, make_opponent +from general_bots_training.opponents import SelfPlayOpponent, make_opponent @pytest.mark.parametrize( @@ -28,28 +28,3 @@ def test_make_opponent_supports_self_play_aliases(): def test_make_opponent_rejects_unknown_name(): with pytest.raises(ValueError, match="unknown opponent"): make_opponent("turtle") - - -def test_make_opponent_model_requires_checkpoint(): - with pytest.raises(ValueError, match="requires a checkpoint path"): - make_opponent("model") - - -def test_make_opponent_model_loads_checkpoint(tmp_path): - import equinox as eqx - import jax.random as jrandom - - from general_bots_training.network import PolicyValueNetwork - - network = PolicyValueNetwork(jrandom.PRNGKey(0), in_channels=14) - checkpoint_path = tmp_path / "model.eqx" - eqx.tree_serialise_leaves(checkpoint_path, network) - - opponent = make_opponent(f"model:{checkpoint_path}") - - assert isinstance(opponent, ModelOpponent) - - -def test_make_opponent_rejects_checkpoint_on_non_model(): - with pytest.raises(ValueError, match="does not take a checkpoint"): - make_opponent("random:some.eqx") diff --git a/uv.lock b/uv.lock index 6ea418c..6ba0efd 100644 --- a/uv.lock +++ b/uv.lock @@ -132,20 +132,12 @@ source = { editable = "." } dependencies = [ { name = "equinox" }, { name = "generals-bots" }, - { name = "jax" }, + { name = "jax", extra = ["cuda"] }, { 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" }, @@ -156,14 +148,11 @@ dev = [ requires-dist = [ { name = "equinox", specifier = ">=0.13.8" }, { name = "generals-bots", git = "https://github.com/strakam/generals-bots.git" }, - { 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 = "jax", extras = ["cuda"], 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 = [ @@ -264,10 +253,6 @@ 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" @@ -310,48 +295,6 @@ 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" @@ -482,18 +425,6 @@ 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" @@ -506,15 +437,6 @@ 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" @@ -524,24 +446,6 @@ 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" @@ -551,20 +455,6 @@ 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" @@ -574,15 +464,6 @@ 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" @@ -592,15 +473,6 @@ 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" @@ -622,30 +494,6 @@ 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" @@ -658,20 +506,6 @@ 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" @@ -686,18 +520,6 @@ 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" @@ -719,24 +541,6 @@ 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" @@ -758,28 +562,6 @@ 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"