Sign skills

.github/evals/reference_jax_challenges.py

@@ -0,0 +1,73 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Maintainer reference only. Never copy this file into an agent environment."""
+
+import cuequivariance_jax.ir_dict as ir_dict
+import jax
+import jax.numpy as jnp
+
+import cuequivariance as cue
+import cuequivariance_jax as cuex
+
+jax.config.update("jax_enable_x64", True)
+
+
+def to_blocks(irreps, flat):
+    data = cuex.RepArray(cue.IrrepsAndLayout(irreps, cue.ir_mul), flat).regroup()
+    return ir_dict.flat_to_dict(data.irreps, data.array, layout="ir_mul")
+
+
+def from_blocks(irreps, blocks):
+    offsets, parts = {}, []
+    for mul, ir in irreps:
+        start = offsets.get(ir, 0)
+        block = blocks[ir][..., start : start + mul, :]
+        parts.append(block.swapaxes(-1, -2).reshape(block.shape[:-2] + (mul * ir.dim,)))
+        offsets[ir] = start + mul
+    return jnp.concatenate(parts, axis=-1)
+
+
+INPUT = cue.Irreps("SO3", "2x0 + 2x1")
+OUTPUT = cue.Irreps("SO3", "0 + 1")
+PRODUCT = cue.descriptors.channelwise_tensor_product(
+    INPUT, cue.Irreps("SO3", "0 + 1 + 2"), [cue.SO3(0), cue.SO3(1)]
+)
+PROJECTION = cue.descriptors.linear(PRODUCT.outputs[0].irreps, OUTPUT)
+HARMONICS = cue.descriptors.spherical_harmonics(cue.SO3(1), [0, 1, 2])
+
+
+def transpose(flat, irreps, *, source):
+    parts, offset = [], 0
+    for mul, ir in irreps:
+        suffix = (mul, ir.dim) if source == "mul_ir" else (ir.dim, mul)
+        block = flat[..., offset : offset + mul * ir.dim]
+        parts.append(
+            block.reshape(flat.shape[:-1] + suffix)
+            .swapaxes(-1, -2)
+            .reshape(flat.shape[:-1] + (mul * ir.dim,))
+        )
+        offset += mul * ir.dim
+    return jnp.concatenate(parts, axis=-1)
+
+
+def run(descriptor, operands, rows):
+    return cuex.segmented_polynomial(
+        descriptor.polynomial,
+        operands,
+        [jax.ShapeDtypeStruct((rows, descriptor.outputs[0].dim), jnp.float64)],
+        method="naive",
+    )[0]
+
+
+def apply(params, features, vectors, senders, receivers, num_dst):
+    lengths = jnp.linalg.norm(vectors, axis=-1, keepdims=True)
+    weights = jnp.tanh(lengths @ params["radial_in"]) @ params["radial_out"]
+    sh = run(HARMONICS, [vectors / lengths], vectors.shape[0])
+    inputs = transpose(features, INPUT, source="mul_ir")
+    messages = run(PRODUCT, [weights, inputs[senders], sh], vectors.shape[0])
+    summed = (
+        jnp.zeros((num_dst, messages.shape[-1]), messages.dtype)
+        .at[receivers]
+        .add(messages)
+    )
+    output = run(PROJECTION, [params["projection"], summed], num_dst)
+    return transpose(output, OUTPUT, source="ir_mul")

.github/evals/reference_torch_challenges.py

@@ -0,0 +1,73 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Maintainer reference only. Never copy this file into an agent environment."""
+
+import torch
+
+import cuequivariance as cue
+import cuequivariance_torch as cuet
+
+
+def to_blocks(irreps, flat):
+    parts, offset = {}, 0
+    for mul, ir in irreps:
+        block = flat[..., offset : offset + mul * ir.dim]
+        block = block.reshape(flat.shape[:-1] + (ir.dim, mul)).transpose(-1, -2)
+        parts.setdefault(ir, []).append(block)
+        offset += mul * ir.dim
+    return {ir: torch.cat(values, dim=-2) for ir, values in parts.items()}
+
+
+def from_blocks(irreps, blocks):
+    offsets, parts = {}, []
+    for mul, ir in irreps:
+        start = offsets.get(ir, 0)
+        block = blocks[ir][..., start : start + mul, :]
+        parts.append(
+            block.transpose(-1, -2).reshape(block.shape[:-2] + (mul * ir.dim,))
+        )
+        offsets[ir] = start + mul
+    return torch.cat(parts, dim=-1)
+
+
+class GraphLayer(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.product = cuet.ChannelWiseTensorProduct(
+            cue.Irreps("SO3", "2x0 + 2x1"),
+            cue.Irreps("SO3", "0 + 1 + 2"),
+            [cue.SO3(0), cue.SO3(1)],
+            shared_weights=False,
+            internal_weights=False,
+            layout=cue.mul_ir,
+            method="naive",
+            dtype=torch.float64,
+        )
+        self.readout = cuet.Linear(
+            self.product.irreps_out,
+            cue.Irreps("SO3", "0 + 1"),
+            internal_weights=False,
+            layout=cue.mul_ir,
+            method="naive",
+            dtype=torch.float64,
+        )
+        self.sh = cuet.SphericalHarmonics([0, 1, 2], normalize=True, method="naive")
+        self.radial_in = torch.nn.Parameter(torch.zeros(1, 4, dtype=torch.float64))
+        self.radial_out = torch.nn.Parameter(
+            torch.zeros(4, self.product.weight_numel, dtype=torch.float64)
+        )
+        self.projection = torch.nn.Parameter(
+            torch.zeros(self.readout.weight_numel, dtype=torch.float64)
+        )
+
+    def forward(self, features, vectors, senders, receivers, num_dst):
+        lengths = torch.linalg.vector_norm(vectors, dim=-1, keepdim=True)
+        weights = torch.tanh(lengths @ self.radial_in) @ self.radial_out
+        messages = self.product(
+            features,
+            self.sh(vectors),
+            weight=weights,
+            indices_1=senders,
+            indices_out=receivers,
+            size_out=num_dst,
+        )
+        return self.readout(messages, self.projection[None, :])

.github/scripts/check_skill_packaging.py

@@ -0,0 +1,115 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Check skill resources and CLI through the same build path as publish.yml."""
+
+import argparse
+import os
+import shutil
+import subprocess
+import sys
+import tarfile
+import tempfile
+import zipfile
+from pathlib import Path
+
+ROOT = Path(__file__).resolve().parents[2]
+PACKAGES = ("cuequivariance", "cuequivariance_jax", "cuequivariance_torch")
+
+
+def check_cli(package, import_root, expected, cwd):
+    env = {**os.environ, "PYTHONPATH": str(import_root)}
+    # Verify import provenance too: an editable install must not hide a missing
+    # resource or module in the wheel being checked.
+    subprocess.run(
+        [
+            sys.executable,
+            "-c",
+            "import importlib, pathlib, sys; "
+            "m = importlib.import_module(sys.argv[1]); "
+            "assert pathlib.Path(m.__file__).is_relative_to(sys.argv[2]), m.__file__",
+            package,
+            str(import_root),
+        ],
+        cwd=cwd,
+        env=env,
+        check=True,
+    )
+    actual = subprocess.check_output(
+        [sys.executable, "-m", package, "skill"], cwd=cwd, env=env
+    )
+    assert actual == expected + b"\n", (
+        f"{package}: CLI output differs from canonical skill"
+    )
+
+
+def build(project, kind, destination):
+    subprocess.run(
+        [
+            sys.executable,
+            "-m",
+            "hatchling",
+            "build",
+            "-t",
+            kind,
+            "-d",
+            str(destination),
+        ],
+        cwd=project,
+        check=True,
+    )
+
+
+def check_package(package, temp):
+    canonical = ROOT / "skills" / package.replace("_", "-") / "SKILL.md"
+    expected = canonical.read_bytes()
+    assert not canonical.is_symlink(), f"{canonical}: nv-carps needs a real file"
+    assert (ROOT / package / package / "SKILL.md").read_bytes() == expected
+    check_cli(package, ROOT / package, expected, temp)
+
+    # Publishing already dereferences LICENSE, README, VERSION and SKILL links.
+    # Copy with dereferencing to exercise that behavior without modifying the
+    # developer's checkout. Direct sdists from the linked tree are not releases.
+    project = temp / "source" / package
+    shutil.copytree(
+        ROOT / package,
+        project,
+        symlinks=False,
+        ignore=shutil.ignore_patterns(
+            "__pycache__", ".pytest_cache", "dist", "*.egg-info"
+        ),
+    )
+    wheel_dir, sdist_dir = temp / "wheel", temp / "sdist"
+    build(project, "wheel", wheel_dir)
+    build(project, "sdist", sdist_dir)
+    archive = next(sdist_dir.glob("*.tar.gz"))
+    with tarfile.open(archive) as source:
+        members = [
+            m for m in source.getmembers() if m.name.endswith(f"/{package}/SKILL.md")
+        ]
+        assert len(members) == 1 and members[0].isfile(), (
+            "sdist must contain a regular SKILL.md"
+        )
+        assert source.extractfile(members[0]).read() == expected
+        source.extractall(temp / "unpacked", filter="data")
+    rebuilt = temp / "rebuilt"
+    build(next((temp / "unpacked").iterdir()), "wheel", rebuilt)
+    for label, directory in (("wheel", wheel_dir), ("rebuilt", rebuilt)):
+        wheel = next(directory.glob("*.whl"))
+        import_root = temp / f"import-{label}"
+        with zipfile.ZipFile(wheel) as contents:
+            assert contents.read(f"{package}/SKILL.md") == expected
+            contents.extractall(import_root)
+        check_cli(package, import_root, expected, temp)
+    print(f"{package}: source CLI, wheel, sdist, and rebuilt wheel passed")
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--package", choices=PACKAGES)
+    args = parser.parse_args()
+    for package in (args.package,) if args.package else PACKAGES:
+        with tempfile.TemporaryDirectory(prefix=f"{package}-skill-packaging-") as temp:
+            check_package(package, Path(temp))
+
+
+if __name__ == "__main__":
+    main()