Bootstrap .venv and verify a TPU runtime receipt

# Step 1 — Define the closed-form sum-of-squares oracle and pmap receipt builder: Each local device computes the sum of squares of its 8-element row...
# Import hashlib for this computation.
import hashlib
import json
import platform
import jax
import jax.numpy as jnp
import numpy as np


# Function `closed_form_sum_of_squares(num_devices, width)` implementing this stage's computation:
def closed_form_sum_of_squares(num_devices, width=8):
    # Evaluate `num_devices) * int(width` and convert the result into Python scalar/collection `n`.
    n = int(num_devices) * int(width)
    # Return `float((n - 1) * n * (2 * n - 1) // 6)` to the caller.
    return float((n - 1) * n * (2 * n - 1) // 6)


# Function `build_runtime_receipt()` implementing this stage's computation:
def build_runtime_receipt():
    # Query the active JAX devices into `devices`.
    devices = jax.local_devices()
    # Run `len` to compute `n_dev`.
    n_dev = len(devices)
    # Construct and reshape `x` into the target tensor dimensions.
    x = jnp.arange(n_dev * 8, dtype=jnp.float32).reshape(n_dev, 8)
    # Reduce across the target axis to summarize `per_device`.
    per_device = jax.pmap(lambda row: jnp.sum(row * row))(x)
    # Synchronize host execution until asynchronous device computation completes.
    observed_total = float(jax.block_until_ready(jnp.sum(per_device)))
    # Run `closed_form_sum_of_squares` to compute `expected_total`.
    expected_total = closed_form_sum_of_squares(n_dev, 8)
    # Check which hardware backend (`cpu`, `gpu`, or `tpu`) JAX selected for `payload`.
    payload = {
        "python": platform.python_version(),
        "jax": jax.__version__,
        "backend": jax.default_backend(),
        "local_device_count": n_dev,
        "per_device_sums": [float(v) for v in np.asarray(per_device)],
        "observed_total": observed_total,
        "expected_total": expected_total,
        "verified": bool(np.isclose(observed_total, expected_total)),
    }
    # Read or serialize artifact data on disk (`payload['receipt_sha256']`).
    payload["receipt_sha256"] = hashlib.sha256(
        json.dumps(payload, sort_keys=True, separators=(",", ":")).encode()
    ).hexdigest()[:16]
    # Return `payload` to the caller.
    return payload

# Step 2 — Verify the active runtime receipt and tabulate multi-chip expectations: Knowing the exact expected sum for 1, 4, and 8 devices lets you...
slice_expectations = [
    {"slice": "1-device (CPU / v5e-1)", "devices": 1, "expected_sum": closed_form_sum_of_squares(1)},
    {"slice": "4-device (v5litepod-4 / v6e-4)", "devices": 4, "expected_sum": closed_form_sum_of_squares(4)},
    {"slice": "8-device (v5litepod-8 / v6e-8)", "devices": 8, "expected_sum": closed_form_sum_of_squares(8)},
]
# Run `build_runtime_receipt` to compute `receipt`.
receipt = build_runtime_receipt()
# Assert invariant `receipt["verified"] is True` holds
assert receipt["verified"] is True
# Print the observed values to compare against the expected result.
print("Runtime receipt:", json.dumps(receipt, sort_keys=True))
# Print diagnostic summary of the computed outputs.
print("Expected pmap sums by slice:", {s["slice"]: s["expected_sum"] for s in slice_expectations})

receipt = build_runtime_receipt()
assert receipt["verified"] is True