Cuequivariance Torch

import importlib
import sys

print("python:", sys.executable, sys.version.split()[0])
available, import_errors = {}, {}
for name in ("torch", "cuequivariance", "cuequivariance_torch"):
    try:
        module = importlib.import_module(name)
        available[name] = module
        print(name, getattr(module, "__version__", "unknown"), module.__file__)
    except (ImportError, OSError) as error:
        import_errors[name] = f"{type(error).__name__}: {error}"
        print(name, import_errors[name])
if "torch" in available:
    torch = available["torch"]
    print("CUDA build:", torch.version.cuda, "available:", torch.cuda.is_available())

set -e
cueq_python=python3  # Set to the selected project interpreter's path.
cueq_env=$(mktemp -d "${TMPDIR:-/tmp}/cueq-cpu.XXXXXX")
"$cueq_python" -m venv --without-pip "$cueq_env"
python3 -m pip --python "$cueq_env" install --cert "${PIP_CERT:-/etc/ssl/certs/ca-certificates.crt}" -q --no-deps -r scripts/torch-cpu-wheels.txt
python3 -m pip --python "$cueq_env" install --cert "${PIP_CERT:-/etc/ssl/certs/ca-certificates.crt}" -q -r scripts/requirements.txt
cueq_python="$cueq_env/bin/python"
"$cueq_python" -c 'import sys, torch, cuequivariance_torch; assert torch.version.cuda is None; print("CPU setup complete:", sys.executable, torch.__version__, cuequivariance_torch.__version__)'

set -e
cueq_run=$(mktemp -d "${TMPDIR:-/tmp}/cueq-check.XXXXXX")
printf 'run directory: %s\n' "$cueq_run"
if "$cueq_python" -u check.py >"$cueq_run/output.txt" 2>&1; then
    cueq_status=0
else
    cueq_status=$?
fi
printf '%s\n' "$cueq_status" >"$cueq_run/exit-code.txt"
tail -n 30 "$cueq_run/output.txt"
printf 'exit code: %s\n' "$cueq_status"
test "$cueq_status" -eq 0

import numpy as np
import torch
import cuequivariance as cue
import cuequivariance_torch as cuet

torch.manual_seed(0)
dtype = torch.float64
print("torch:", torch.__version__, "CUDA build:", torch.version.cuda,
      "CUDA available:", torch.cuda.is_available(), "device: cpu", "dtype:", dtype,
      "method: naive")

irreps = cue.Irreps("SO3", "0 + 1")
sh = cuet.SphericalHarmonics([0, 1], normalize=False, method="naive")
tp = cuet.FullyConnectedTensorProduct(
    irreps, irreps, irreps, layout=cue.mul_ir, method="naive", dtype=dtype,
)
features = torch.randn(2, irreps.dim, dtype=dtype, requires_grad=True)
vectors = torch.tensor([[0.3, -0.7, 1.1], [1.2, 0.4, -0.2]],
                       dtype=dtype, requires_grad=True)
result = tp(features, sh(vectors))
assert result.device.type == "cpu" and result.shape == (2, irreps.dim)
result.square().sum().backward()
gradient_norms = {}
for name, tensor in [("features", features), ("vectors", vectors), *tp.named_parameters()]:
    assert tensor.grad is not None and torch.isfinite(tensor.grad).all()
    assert tensor.grad.abs().sum() > 0
    gradient_norms[name] = tensor.grad.norm().item()
print("CPU tensor product:", tuple(result.shape), "gradient norms:", gradient_norms)

ls = [0, 1, 2]
sh = cuet.SphericalHarmonics(ls, normalize=False, method="naive")
vectors = torch.tensor([[0.3, -0.7, 1.1], [1.2, 0.4, -0.2]],
                       dtype=dtype, requires_grad=True)
values = sh(vectors)
descriptor = cue.descriptors.spherical_harmonics(cue.SO3(1), ls)
reference = descriptor(vectors.detach().numpy())[0]
np.testing.assert_allclose(values.detach().numpy(), reference, atol=1e-12)
assert values.shape == (2, sum(2 * l + 1 for l in ls))
gradcheck_ok = torch.autograd.gradcheck(sh, (vectors,))
assert gradcheck_ok
print("Spherical harmonics:", tuple(values.shape), "normalize: False",
      "max error:", np.max(np.abs(values.detach().numpy() - reference)),
      "gradcheck:", gradcheck_ok)

irreps_in = cue.Irreps("SO3", "3x0 + 2x1")
irreps_out = cue.Irreps("SO3", "2x0 + 4x1")
options = dict(internal_weights=False, shared_weights=True, method="naive", dtype=dtype)
linear = cuet.Linear(irreps_in, irreps_out, layout=cue.mul_ir, **options)
linear_ir = cuet.Linear(irreps_in, irreps_out, layout=cue.ir_mul, **options)
to_ir = cuet.TransposeIrrepsLayout(
    irreps_in, source=cue.mul_ir, target=cue.ir_mul, use_fallback=True
)
to_mul = cuet.TransposeIrrepsLayout(
    irreps_out, source=cue.ir_mul, target=cue.mul_ir, use_fallback=True
)
x = torch.randn(2, irreps_in.dim, dtype=dtype, requires_grad=True)
w = torch.randn(1, linear.weight_numel, dtype=dtype, requires_grad=True)
y = linear(x, w)
y_ir = to_mul(linear_ir(to_ir(x), w))
torch.testing.assert_close(y, y_ir)
rotation = np.array([0.2, -0.3, 0.1])
din, dout = [torch.from_numpy(cue.IrrepsAndLayout(r, cue.mul_ir).exp_map(rotation, []))
             for r in (irreps_in, irreps_out)]
rotated, expected_rotated = linear(x @ din.T, w), y @ dout.T
torch.testing.assert_close(rotated, expected_rotated)
grads = torch.autograd.grad(y.square().sum(), (x, w))
for g in grads:
    assert torch.isfinite(g).all() and g.abs().sum() > 0
print("Linear:", tuple(y.shape), "shared external weights:", tuple(w.shape),
      "layout error:", (y - y_ir).detach().abs().max().item(),
      "rotation error:", (rotated - expected_rotated).detach().abs().max().item(),
      "gradient norms:", [g.norm().item() for g in grads])

a, b = cue.Irreps("SO3", "3x0 + 3x1"), cue.Irreps("SO3", "0 + 1")
keep = [cue.SO3(0), cue.SO3(1)]
tp = cuet.ChannelWiseTensorProduct(
    a, b, keep, layout=cue.ir_mul, shared_weights=False,
    internal_weights=False, method="naive", dtype=dtype,
)
x1, x2, w = [torch.randn(2, n, dtype=dtype, requires_grad=True)
             for n in (a.dim, b.dim, tp.weight_numel)]
y = tp(x1, x2, weight=w)
descriptor = cue.descriptors.channelwise_tensor_product(a, b, keep)
reference = descriptor(*(t.detach().numpy() for t in (w, x1, x2)))[0]
np.testing.assert_allclose(y.detach().numpy(), reference, atol=1e-12)
assert y.shape == (2, tp.irreps_out.dim)
grads = torch.autograd.grad(y.square().sum(), (x1, x2, w))
for g in grads:
    assert torch.isfinite(g).all() and g.abs().sum() > 0
print("Channelwise:", tuple(y.shape), "per-item weights:", tuple(w.shape),
      "filter selects types; actual ir_mul output order:", tp.irreps_out,
      "max error:", np.max(np.abs(y.detach().numpy() - reference)),
      "gradient norms:", [g.norm().item() for g in grads])

descriptor = cue.descriptors.linear(cue.Irreps("SO3", "3x0"), cue.Irreps("SO3", "2x0"))
sp = cuet.SegmentedPolynomial(descriptor.polynomial, method="naive")
nodes = torch.randn(5, 3, dtype=dtype, requires_grad=True)
w = torch.randn(1, descriptor.inputs[0].dim, dtype=dtype, requires_grad=True)
senders = torch.tensor([0, 2, 1, 0], dtype=torch.int64)
receivers = torch.tensor([1, 0, 1, 2], dtype=torch.int64)
[actual] = sp([w, nodes], input_indices={1: senders},
              output_indices={0: receivers}, output_shapes={0: nodes.new_empty(4, 1)})
[edges] = sp([w, nodes[senders]])
reference = nodes.new_zeros(4, descriptor.outputs[0].dim).index_add_(0, receivers, edges)
torch.testing.assert_close(actual, reference)
assert torch.count_nonzero(actual[3]) == 0
grads = torch.autograd.grad(actual.square().sum(), (nodes, w))
ref_grads = torch.autograd.grad(reference.square().sum(), (nodes, w))
for g, ref in zip(grads, ref_grads):
    torch.testing.assert_close(g, ref)
print("Gather/scatter:", tuple(actual.shape),
      "input/output index dtypes:", senders.dtype, receivers.dtype,
      "forward error:", (actual - reference).detach().abs().max().item(),
      "gradient errors:", [(g - r).abs().max().item() for g, r in zip(grads, ref_grads)],
      "isolated destination:", actual[3].detach().tolist())

class SimpleGNN(torch.nn.Module):
    def __init__(self, irreps_in, irreps_out):
        super().__init__()
        self.tp = cuet.ChannelWiseTensorProduct(
            irreps_in, cue.Irreps("SO3", "0 + 1 + 2"), layout=cue.mul_ir,
            shared_weights=False, internal_weights=False, method="naive",
        )
        self.radial = torch.nn.Sequential(
            torch.nn.Linear(1, 16), torch.nn.SiLU(),
            torch.nn.Linear(16, self.tp.weight_numel),
        )
        self.sh = cuet.SphericalHarmonics([0, 1, 2], normalize=True, method="naive")
        self.linear = cuet.Linear(
            self.tp.irreps_out, irreps_out, layout=cue.mul_ir, method="naive"
        )

    def forward(self, node_features, edge_vectors, senders, receivers, num_dst):
        weights = self.radial(edge_vectors.norm(dim=-1, keepdim=True))
        messages = self.tp(
            node_features, self.sh(edge_vectors), weight=weights,
            indices_1=senders, indices_out=receivers, size_out=num_dst,
        )
        return self.linear(messages)