Jax Development

#!/usr/bin/env python3
"""Honest JAX benchmarking starter."""

import time

import jax
import jax.numpy as jnp


def workload(x):
    return jnp.tanh(x @ x.T).sum()


def main():
    x = jax.random.normal(jax.random.key(0), (2048, 256))
    workload_jit = jax.jit(workload)

    # Warm-up / compile
    workload_jit(x).block_until_ready()

    times = []
    for _ in range(10):
        t0 = time.perf_counter()
        y = workload_jit(x)
        y.block_until_ready()
        times.append((time.perf_counter() - t0) * 1e3)

    print("steady-state times_ms:", times)
    print("mean_ms:", sum(times) / len(times))


if __name__ == "__main__":
    main()

#!/usr/bin/env python3
"""Runtime checks that survive jit via checkify."""

from jax.experimental import checkify
import jax
import jax.numpy as jnp


def f(x, i):
    checkify.check(i >= 0, "index must be non-negative: {i}", i=i)
    checkify.check(i < x.shape[0], "index out of bounds: {i}", i=i)
    y = x[i]
    checkify.check(jnp.isfinite(y).all(), "non-finite output")
    return y


def main():
    x = jnp.arange(8, dtype=jnp.float32)
    checked = checkify.checkify(f, errors=checkify.user_checks | checkify.index_checks)
    err, y = jax.jit(checked)(x, 3)
    err.throw()
    print("y:", y)


if __name__ == "__main__":
    main()

#!/usr/bin/env python3
"""Custom VJP starter."""

import jax
import jax.numpy as jnp


@jax.custom_vjp
def clipped_log1p(x):
    return jnp.log1p(jnp.maximum(x, -0.999999))


def clipped_log1p_fwd(x):
    y = clipped_log1p(x)
    return y, x


def clipped_log1p_bwd(residual, g):
    (x,) = (residual,)
    grad = 1.0 / (1.0 + jnp.maximum(x, -0.999999))
    return (g * grad,)


clipped_log1p.defvjp(clipped_log1p_fwd, clipped_log1p_bwd)


def main():
    x = jnp.array([0.0, 1.0, 2.0], dtype=jnp.float32)
    print(clipped_log1p(x))
    print(jax.grad(lambda t: jnp.sum(clipped_log1p(t)))(x))


if __name__ == "__main__":
    main()

#!/usr/bin/env python3
"""Export / serialisation starter."""

import numpy as np

import jax
import jax.numpy as jnp
from jax import export


def f(x):
    return jnp.sin(x) + 2.0 * x


def main():
    sig = jax.ShapeDtypeStruct((4,), np.float32)
    exported = export.export(jax.jit(f))(sig)
    print("platforms:", exported.platforms)
    print("in_avals:", exported.in_avals)
    print("out_avals:", exported.out_avals)


if __name__ == "__main__":
    main()

#!/usr/bin/env python3
"""Minimal reproducible example template for JAX bugs."""

import jax
import jax.numpy as jnp


def f(x):
    # Replace with the smallest function that still reproduces the problem.
    return x + 1


def main():
    x = jnp.arange(4, dtype=jnp.float32)
    print("eager:", f(x))
    print("jit:", jax.jit(f)(x))


if __name__ == "__main__":
    main()

#!/usr/bin/env python3
"""Minimal Pallas skeleton.

This is only a starting point for real kernel work after higher-level JAX
optimisation has been exhausted.
"""

import jax
import jax.numpy as jnp
from jax.experimental import pallas as pl


def add_one_kernel(x_ref, y_ref):
    idx = pl.program_id(0)
    y_ref[idx] = x_ref[idx] + 1.0


def main():
    x = jnp.arange(16, dtype=jnp.float32)
    y = pl.pallas_call(
        add_one_kernel,
        out_shape=jax.ShapeDtypeStruct(x.shape, x.dtype),
        grid=(x.shape[0],),
    )(x)
    print(y)


if __name__ == "__main__":
    main()

#!/usr/bin/env python3
"""Trace and memory-profile starter."""

import jax
import jax.numpy as jnp
import jax.profiler


@jax.jit
def step(x):
    return jnp.tanh(x @ x.T)


def main():
    x = jax.random.normal(jax.random.key(0), (2048, 256))

    with jax.profiler.trace("/tmp/jax-trace"):
        for _ in range(5):
            x = step(x[:, :256])
            x.block_until_ready()

    # Optional: capture device memory profile for OOM / leak work.
    jax.profiler.save_device_memory_profile("/tmp/jax-memory.prof")
    print("Wrote /tmp/jax-trace and /tmp/jax-memory.prof")


if __name__ == "__main__":
    main()

#!/usr/bin/env python3
"""Carry-state loop starter using lax.scan."""

import jax
import jax.numpy as jnp


def step(carry, x_t):
    state = carry
    new_state = 0.95 * state + x_t
    y_t = jnp.tanh(new_state)
    return new_state, y_t


@jax.jit
def run(init_state, xs):
    final_state, ys = jax.lax.scan(step, init_state, xs)
    return final_state, ys


def main():
    xs = jnp.linspace(0.0, 1.0, 10)
    final_state, ys = run(jnp.array(0.0), xs)
    print("final_state:", final_state)
    print("ys:", ys)


if __name__ == "__main__":
    main()

#!/usr/bin/env python3
"""Starter for manual SPMD code with jax.shard_map."""

from __future__ import annotations

import numpy as np

import jax
import jax.numpy as jnp
from jax.sharding import Mesh, PartitionSpec as P


def main():
    devices = np.array(jax.devices())
    if devices.size < 2:
        raise SystemExit("This template expects at least 2 devices.")

    mesh = Mesh(devices.reshape((devices.size,)), ("data",))

    @jax.shard_map(mesh=mesh, in_specs=P("data"), out_specs=P("data"))
    def per_shard_scale(x):
        # `x` is the local shard view. Add collectives here if needed.
        return 2.0 * x

    x = jnp.arange(devices.size * 4, dtype=jnp.float32).reshape(devices.size, 4)
    y = per_shard_scale(x)
    print(y)


if __name__ == "__main__":
    main()

#!/usr/bin/env python3
"""Starter for global-view sharding with Mesh and NamedSharding."""

from __future__ import annotations

import numpy as np

import jax
import jax.numpy as jnp
from jax.sharding import Mesh, NamedSharding, PartitionSpec as P


def main():
    devices = np.array(jax.devices())
    if devices.size < 2:
        raise SystemExit("This template expects at least 2 devices.")

    mesh = Mesh(devices.reshape((devices.size,)), ("data",))
    data_sharding = NamedSharding(mesh, P("data", None))
    repl_sharding = NamedSharding(mesh, P())

    x = jax.device_put(jnp.arange(devices.size * 8, dtype=jnp.float32).reshape(devices.size, 8), data_sharding)
    w = jax.device_put(jnp.eye(8, dtype=jnp.float32), repl_sharding)

    @jax.jit
    def f(x, w):
        return x @ w

    y = f(x, w)
    print("input sharding:", x.sharding)
    print("output sharding:", y.sharding)


if __name__ == "__main__":
    main()

#!/usr/bin/env python3
"""Compiled training-step starter with explicit PRNG plumbing."""

from __future__ import annotations

import functools

import jax
import jax.numpy as jnp


def init_params(key, in_dim: int, out_dim: int):
    k1, k2 = jax.random.split(key)
    w = 0.01 * jax.random.normal(k1, (in_dim, out_dim))
    b = jnp.zeros((out_dim,), dtype=w.dtype)
    return {"w": w, "b": b}, k2


def model(params, x):
    return x @ params["w"] + params["b"]


def loss_fn(params, batch, key):
    x, y = batch
    key, dropout_key = jax.random.split(key)
    logits = model(params, x)
    keep = jax.random.bernoulli(dropout_key, p=0.9, shape=logits.shape)
    logits = jnp.where(keep, logits / 0.9, 0.0)
    loss = jnp.mean((logits - y) ** 2)
    metrics = {"loss": loss}
    return (loss, metrics), key


@functools.partial(jax.jit, donate_argnums=(0,))
def train_step(params, batch, key, lr=1e-2):
    def wrapped_loss(p):
        (loss, metrics), new_key = loss_fn(p, batch, key)
        return loss, (metrics, new_key)

    (loss, (metrics, new_key)), grads = jax.value_and_grad(wrapped_loss, has_aux=True)(params)
    new_params = jax.tree.map(lambda p, g: p - lr * g, params, grads)
    sq_norm = sum((jnp.sum(g * g) for g in jax.tree.leaves(grads)), start=jnp.array(0.0, dtype=loss.dtype))
    metrics = {**metrics, "grad_norm": jnp.sqrt(sq_norm)}
    return new_params, new_key, metrics


def main():
    key = jax.random.key(0)
    params, key = init_params(key, in_dim=8, out_dim=4)
    x = jax.random.normal(key, (16, 8))
    y = jax.random.normal(key, (16, 4))
    params, key, metrics = train_step(params, (x, y), key)
    print(jax.device_get(metrics))


if __name__ == "__main__":
    main()

{
  "skill_name": "jax-development",
  "evals": [
    {
      "id": "control-flow-cond",
      "prompt": "Refactor evals/files/bad_control_flow.py into idiomatic JAX that works under jit, and explain why the original version fails.",
      "expected_output": "A corrected version that replaces data-dependent Python control flow with a JAX control-flow primitive or array-wise selection, plus a clear explanation of tracer/concretisation semantics.",
      "files": [
        "evals/files/bad_control_flow.py"
      ],
      "assertions": [
        "The solution replaces Python data-dependent control flow inside jitted code with `lax.cond`, `lax.switch`, or `jnp.where`.",
        "The response explains that Python tried to use a traced value as a concrete boolean or equivalent control-flow guard.",
        "The final code is runnable and remains pure."
      ]
    },
    {
      "id": "prng-threading",
      "prompt": "Fix the PRNG handling in evals/files/prng_reuse.py. I want explicit key threading and a short note on how to verify we are not reusing keys.",
      "expected_output": "A rewritten version that accepts and returns keys explicitly, splits keys correctly, and explains why reusing a key gives repeated random behaviour.",
      "files": [
        "evals/files/prng_reuse.py"
      ],
      "assertions": [
        "The solution removes the hidden module-level key from the hot path or clearly avoids using it.",
        "The rewritten code uses `split` or equivalent explicit key management.",
        "The response explains how to verify the fix, for example by repeated calls or key-reuse debugging."
      ]
    },
    {
      "id": "benchmark-honesty",
      "prompt": "Review evals/files/naive_benchmark.py and show me how to benchmark it properly on JAX.",
      "expected_output": "An answer that distinguishes compile time from steady-state execution, uses blocking, and avoids timing only dispatch.",
      "files": [
        "evals/files/naive_benchmark.py"
      ],
      "assertions": [
        "The response mentions warm-up or first-call compilation separately from repeated-call timing.",
        "The timing method uses `.block_until_ready()` or `jax.block_until_ready(...)`.",
        "The answer warns about asynchronous dispatch or host timing pitfalls."
      ]
    },
    {
      "id": "pmap-migration",
      "prompt": "Look at evals/files/legacy_pmap.py and outline a sensible migration path to newer sharding APIs without changing semantics prematurely.",
      "expected_output": "A measured migration plan that preserves semantics first, explains the current `pmap` role, and suggests modern sharding APIs deliberately.",
      "files": [
        "evals/files/legacy_pmap.py"
      ],
      "assertions": [
        "The answer does not recommend a blind rewrite without discussing semantics.",
        "The response recognises `pmap` as legacy-oriented relative to newer sharding APIs.",
        "The migration plan mentions meshes, named sharding, shard_map, or equivalent modern concepts."
      ]
    },
    {
      "id": "host-roundtrip",
      "prompt": "Diagnose evals/files/host_roundtrip.py. It works sometimes, but I suspect there's an accidental host round-trip inside jit.",
      "expected_output": "A diagnosis that identifies the host conversion inside the compiled region and rewrites the code to stay in JAX space until the edge.",
      "files": [
        "evals/files/host_roundtrip.py"
      ],
      "assertions": [
        "The response identifies the NumPy conversion inside `jit` as a problem.",
        "The rewritten code returns a JAX value from the jitted function and moves host conversion outside if needed.",
        "The explanation mentions performance or tracer implications of host conversion."
      ]
    },
    {
      "id": "closure-recompile",
      "prompt": "Review evals/files/closure_recompile.py and tell me why it is likely to trigger repeated tracing or compilation. Then rewrite it cleanly.",
      "expected_output": "A diagnosis of jitted function creation inside a loop or changing closure capture, plus a cleaner design with hoisted compilation.",
      "files": [
        "evals/files/closure_recompile.py"
      ],
      "assertions": [
        "The answer points out that `jax.jit` is being created inside a loop or repeatedly around a closure.",
        "The rewritten version hoists `jit` or otherwise stabilises the compilation boundary.",
        "The response mentions compile-cache churn or repeated tracing."
      ]
    },
    {
      "id": "scan-rewrite",
      "prompt": "Rewrite evals/files/python_loop_blowup.py in a way that is more JAX-friendly for compile time and explain why.",
      "expected_output": "A version using `lax.scan` or `lax.fori_loop` with an explanation that large Python loops under jit can inflate the traced program.",
      "files": [
        "evals/files/python_loop_blowup.py"
      ],
      "assertions": [
        "The answer replaces the large Python loop with a structured JAX loop primitive.",
        "The response explains why the original loop could blow up compile size or trace time.",
        "The rewritten code preserves the intended accumulation semantics."
      ]
    },
    {
      "id": "dynamic-mask",
      "prompt": "Fix evals/files/dynamic_mask.py for use under jit. I still want something that preserves the semantics as much as possible.",
      "expected_output": "An answer that explains why boolean masking creates a dynamic-size result and proposes a fixed-shape alternative such as `where`, padding, or a mask/value pair.",
      "files": [
        "evals/files/dynamic_mask.py"
      ],
      "assertions": [
        "The response explains the dynamic-shape issue with boolean masking under jit.",
        "The answer proposes a fixed-shape rewrite or an explicit alternative representation.",
        "The answer does not pretend that `x[x > 0]` is directly safe inside jit."
      ]
    }
  ]
}

import jax
import jax.numpy as jnp

@jax.jit
def clip_or_negate(x):
    if x.mean() > 0:
        return jnp.clip(x, 0.0, 1.0)
    return -x

import jax
import jax.numpy as jnp

def run(xs, scale):
    outputs = []
    for x in xs:
        f = jax.jit(lambda y: scale * y + 1.0)
        outputs.append(f(x))
    return outputs

import jax
import jax.numpy as jnp

@jax.jit
def positives_only(x):
    return x[x > 0]

import numpy as np
import jax
import jax.numpy as jnp

@jax.jit
def f(x):
    y = jnp.sin(x)
    return np.asarray(y).sum()

import jax
import jax.numpy as jnp

@jax.pmap(axis_name="data")
def step(x):
    y = x - jax.lax.pmean(x, "data")
    return y / (1e-6 + jnp.linalg.norm(y, axis=-1, keepdims=True))

import jax
import jax.numpy as jnp

def matmul_step(x, w):
    return jnp.tanh(x @ w)

def benchmark_once():
    x = jax.random.normal(jax.random.key(0), (2048, 512))
    w = jax.random.normal(jax.random.key(1), (512, 512))
    fast = jax.jit(matmul_step)
    return fast(x, w)  # caller times this without blocking

import jax
import jax.numpy as jnp

GLOBAL_KEY = jax.random.key(0)

@jax.jit
def sample_pair():
    a = jax.random.normal(GLOBAL_KEY, (4,))
    b = jax.random.uniform(GLOBAL_KEY, (4,))
    return a, b

import jax
import jax.numpy as jnp

@jax.jit
def accumulate(x):
    acc = jnp.zeros_like(x[0])
    for i in range(1000):
        acc = acc + jnp.sin(x[i])
    return acc

[
  {
    "query": "Can you turn this NumPy finite-difference update into JAX and batch it over 64 parameter sets with vmap?",
    "should_trigger": true
  },
  {
    "query": "why does jit throw ConcretizationTypeError when I branch on an array value?",
    "should_trigger": true
  },
  {
    "query": "I think my dropout mask repeats every step in JAX. Please inspect the key plumbing.",
    "should_trigger": true
  },
  {
    "query": "Benchmark this JAX training step properly on GPU and tell me whether compile or runtime dominates.",
    "should_trigger": true
  },
  {
    "query": "Migrate this pmap loop to the newer sharding style without breaking semantics.",
    "should_trigger": true
  },
  {
    "query": "Show me where custom_vjp and batching rules live in the JAX repo.",
    "should_trigger": true
  },
  {
    "query": "Export this jitted function as StableHLO and explain the shape assumptions.",
    "should_trigger": true
  },
  {
    "query": "Why does jax.debug.print show values in a weird order inside my compiled function?",
    "should_trigger": true
  },
  {
    "query": "I need to update formulas in an Excel workbook and fix the conditional formatting.",
    "should_trigger": false
  },
  {
    "query": "Optimise this pure NumPy code with Numba; JAX is not allowed here.",
    "should_trigger": false
  },
  {
    "query": "Write a CUDA kernel for a custom NCCL all-reduce without using any JAX code.",
    "should_trigger": false
  },
  {
    "query": "Can you review my PyTorch Lightning callback setup for checkpoint naming?",
    "should_trigger": false
  }
]

[
  {
    "query": "Port this PyTorch forward pass to functional JAX and make the recurrent loop compile efficiently.",
    "should_trigger": true
  },
  {
    "query": "I need help understanding Mesh, PartitionSpec, and NamedSharding for a 2x4 device setup.",
    "should_trigger": true
  },
  {
    "query": "This JAX program gets slower every time shapes change. Can you check for recompiles?",
    "should_trigger": true
  },
  {
    "query": "Please write a tiny Pallas kernel skeleton for a hotspot after ordinary jit did not help.",
    "should_trigger": true
  },
  {
    "query": "Could you clean up my pandas DataFrame and make a chart from the CSV?",
    "should_trigger": false
  },
  {
    "query": "Recommend a Linux driver version for CUDA generally; there is no JAX involved.",
    "should_trigger": false
  },
  {
    "query": "Explain autograd in general, not JAX specifically.",
    "should_trigger": false
  },
  {
    "query": "Help me refactor this plain Python state machine with lots of dictionaries and strings.",
    "should_trigger": false
  }
]

\
#!/usr/bin/env python3
"""Benchmark a Python callable with optional JAX JIT and proper blocking."""

from __future__ import annotations

import argparse
import importlib
import importlib.util
import json
import statistics
import sys
import time
from pathlib import Path
from typing import Any


def load_module(module_name: str | None, file_path: str | None) -> Any:
    if bool(module_name) == bool(file_path):
        raise ValueError("Exactly one of --module or --file is required.")
    if module_name:
        return importlib.import_module(module_name)

    path = Path(file_path or "")
    if not path.exists():
        raise FileNotFoundError(f"Module file not found: {path}")

    spec = importlib.util.spec_from_file_location(path.stem, path)
    if spec is None or spec.loader is None:
        raise ImportError(f"Could not load module spec from: {path}")
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


def resolve_attr(obj: Any, dotted: str) -> Any:
    current = obj
    for part in dotted.split("."):
        current = getattr(current, part)
    return current


def load_json_arg(raw: str | None, file_path: str | None, default: Any) -> Any:
    if raw is not None and file_path is not None:
        raise ValueError("Choose either the inline JSON form or the file form, not both.")
    if raw is not None:
        return json.loads(raw)
    if file_path is not None:
        return json.loads(Path(file_path).read_text(encoding="utf-8"))
    return default


def numeric_tree(value: Any) -> bool:
    if isinstance(value, (int, float, bool)):
        return True
    if isinstance(value, list):
        return all(numeric_tree(v) for v in value)
    return False


def maybe_import_jax():
    try:
        jax = importlib.import_module("jax")
        jnp = importlib.import_module("jax.numpy")
        return jax, jnp
    except Exception:
        return None, None


def tree_arrayify(value: Any, jnp_module: Any) -> Any:
    if isinstance(value, dict):
        return {k: tree_arrayify(v, jnp_module) for k, v in value.items()}
    if isinstance(value, list):
        if numeric_tree(value):
            return jnp_module.array(value)
        return [tree_arrayify(v, jnp_module) for v in value]
    return value


def tree_device_put(value: Any, jax_module: Any) -> Any:
    if isinstance(value, dict):
        return {k: tree_device_put(v, jax_module) for k, v in value.items()}
    if isinstance(value, list):
        return [tree_device_put(v, jax_module) for v in value]
    try:
        return jax_module.device_put(value)
    except Exception:
        return value


def parse_int_tuple(raw: str | None) -> tuple[int, ...] | None:
    if raw is None or raw == "":
        return None
    return tuple(int(part.strip()) for part in raw.split(",") if part.strip())


def block_until_ready(value: Any, jax_module: Any | None) -> Any:
    if hasattr(value, "block_until_ready"):
        return value.block_until_ready()
    if jax_module is not None:
        try:
            return jax_module.block_until_ready(value)
        except Exception:
            pass
    return value


def run_once(fn: Any, args: list[Any], kwargs: dict[str, Any], jax_module: Any | None) -> Any:
    out = fn(*args, **kwargs)
    block_until_ready(out, jax_module)
    return out


def summary(times_ms: list[float]) -> dict[str, float]:
    return {
        "mean_ms": statistics.mean(times_ms),
        "median_ms": statistics.median(times_ms),
        "min_ms": min(times_ms),
        "max_ms": max(times_ms),
        "stdev_ms": statistics.pstdev(times_ms) if len(times_ms) > 1 else 0.0,
    }


def benchmark(fn: Any, args: list[Any], kwargs: dict[str, Any], repeat: int, jax_module: Any | None) -> list[float]:
    times_ms = []
    for _ in range(repeat):
        t0 = time.perf_counter()
        run_once(fn, args, kwargs, jax_module)
        times_ms.append((time.perf_counter() - t0) * 1e3)
    return times_ms


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Benchmark a Python callable with optional JAX JIT and proper blocking.",
        formatter_class=argparse.RawTextHelpFormatter,
        epilog="""Exit codes:
  0 success
  2 operational error

Examples:
  python3 scripts/jax_benchmark_harness.py --file evals/files/naive_benchmark.py --function matmul_step \\
      --args-json '[[[1.0, 2.0], [3.0, 4.0]], [[1.0], [2.0]]]' --arrayify --jit --compare-eager

  python3 scripts/jax_benchmark_harness.py --module mypkg.train --function step \\
      --args-file args.json --kwargs-file kwargs.json --jit --repeat 20

  python3 scripts/jax_benchmark_harness.py --file train.py --function step \\
      --args-file args.json --jit --static-argnums 2 --donate-argnums 0,1
""",
    )
    source = parser.add_mutually_exclusive_group(required=True)
    source.add_argument("--module", help="Import path for the module containing the callable.")
    source.add_argument("--file", help="Path to a Python file containing the callable.")
    parser.add_argument("--function", required=True, help="Callable name or dotted attribute path.")
    parser.add_argument("--args-json", help="JSON list of positional arguments.")
    parser.add_argument("--args-file", help="Path to a JSON file containing positional arguments.")
    parser.add_argument("--kwargs-json", help="JSON object of keyword arguments.")
    parser.add_argument("--kwargs-file", help="Path to a JSON file containing keyword arguments.")
    parser.add_argument("--arrayify", action="store_true", help="Convert numeric JSON lists to `jax.numpy.array` when JAX is available.")
    parser.add_argument("--device-put", action="store_true", help="Apply `jax.device_put` to arguments before timing when JAX is available.")
    parser.add_argument("--jit", action="store_true", help="Wrap the callable with `jax.jit`.")
    parser.add_argument("--static-argnums", help="Comma-separated positional indices to mark static when using --jit.")
    parser.add_argument("--donate-argnums", help="Comma-separated positional indices to donate when using --jit.")
    parser.add_argument("--compare-eager", action="store_true", help="Also benchmark the original eager callable.")
    parser.add_argument("--repeat", type=int, default=10, help="Number of timed steady-state repetitions. Default: 10")
    parser.add_argument("--warmup", type=int, default=1, help="Warm-up calls before timed loops. Default: 1")
    parser.add_argument("--format", choices=("json", "text"), default="json", help="Output format. Default: json")
    parser.add_argument("--output", help="Write the report to this file instead of stdout.")
    return parser.parse_args()


def main() -> int:
    args = parse_args()

    try:
        module = load_module(args.module, args.file)
        fn = resolve_attr(module, args.function)
        if not callable(fn):
            raise TypeError(f"Resolved object is not callable: {args.function}")

        raw_args = load_json_arg(args.args_json, args.args_file, [])
        raw_kwargs = load_json_arg(args.kwargs_json, args.kwargs_file, {})
        if not isinstance(raw_args, list):
            raise TypeError("Positional arguments JSON must decode to a list.")
        if not isinstance(raw_kwargs, dict):
            raise TypeError("Keyword arguments JSON must decode to an object.")

        jax, jnp = maybe_import_jax()

        proc_args = list(raw_args)
        proc_kwargs = dict(raw_kwargs)
        if args.arrayify:
            if jax is None or jnp is None:
                raise RuntimeError("--arrayify requires JAX to be importable.")
            proc_args = [tree_arrayify(v, jnp) for v in proc_args]
            proc_kwargs = {k: tree_arrayify(v, jnp) for k, v in proc_kwargs.items()}

        if args.device_put:
            if jax is None:
                raise RuntimeError("--device-put requires JAX to be importable.")
            proc_args = [tree_device_put(v, jax) for v in proc_args]
            proc_kwargs = {k: tree_device_put(v, jax) for k, v in proc_kwargs.items()}

        static_argnums = parse_int_tuple(args.static_argnums)
        donate_argnums = parse_int_tuple(args.donate_argnums)

        report: dict[str, Any] = {
            "source": args.module or args.file,
            "callable": args.function,
            "repeat": args.repeat,
            "warmup": args.warmup,
            "jax_available": jax is not None,
            "jit_requested": args.jit,
            "compare_eager": args.compare_eager,
            "static_argnums": static_argnums,
            "donate_argnums": donate_argnums,
        }

        if args.compare_eager or not args.jit:
            for _ in range(args.warmup):
                run_once(fn, proc_args, proc_kwargs, jax)
            eager_times = benchmark(fn, proc_args, proc_kwargs, args.repeat, jax)
            report["eager"] = {
                "times_ms": eager_times,
                "summary": summary(eager_times),
            }

        if args.jit:
            if jax is None:
                raise RuntimeError("--jit requires JAX to be importable.")
            jit_kwargs = {}
            if static_argnums is not None:
                jit_kwargs["static_argnums"] = static_argnums
            if donate_argnums is not None:
                jit_kwargs["donate_argnums"] = donate_argnums
            fn_jit = jax.jit(fn, **jit_kwargs)

            t0 = time.perf_counter()
            run_once(fn_jit, proc_args, proc_kwargs, jax)
            first_call_ms = (time.perf_counter() - t0) * 1e3

            for _ in range(max(args.warmup - 1, 0)):
                run_once(fn_jit, proc_args, proc_kwargs, jax)

            jit_times = benchmark(fn_jit, proc_args, proc_kwargs, args.repeat, jax)
            report["jit"] = {
                "first_call_ms": first_call_ms,
                "times_ms": jit_times,
                "summary": summary(jit_times),
            }

        if "eager" in report and "jit" in report:
            eager_mean = report["eager"]["summary"]["mean_ms"]
            jit_mean = report["jit"]["summary"]["mean_ms"]
            report["speedup_vs_eager_mean"] = (eager_mean / jit_mean) if jit_mean else None

    except Exception as exc:
        sys.stderr.write(f"Error: {type(exc).__name__}: {exc}\n")
        return 2

    if args.format == "json":
        text = json.dumps(report, indent=2, sort_keys=True)
    else:
        lines = [
            f"Callable: {report['source']}::{report['callable']}",
            f"Repeat: {report['repeat']}",
            f"Warmup: {report['warmup']}",
            f"JAX available: {report['jax_available']}",
            f"JIT requested: {report['jit_requested']}",
            f"Static argnums: {report['static_argnums']}",
            f"Donate argnums: {report['donate_argnums']}",
            "",
        ]
        if "eager" in report:
            lines.append("Eager")
            for key, value in report["eager"]["summary"].items():
                lines.append(f"  {key}: {value:.3f}")
            lines.append("")
        if "jit" in report:
            lines.append("JIT")
            lines.append(f"  first_call_ms: {report['jit']['first_call_ms']:.3f}")
            for key, value in report["jit"]["summary"].items():
                lines.append(f"  {key}: {value:.3f}")
            lines.append("")
        if "speedup_vs_eager_mean" in report:
            lines.append(f"speedup_vs_eager_mean: {report['speedup_vs_eager_mean']:.3f}")
        text = "\n".join(lines)

    if args.output:
        Path(args.output).write_text(text + ("" if text.endswith("\n") else "\n"), encoding="utf-8")
    else:
        sys.stdout.write(text)
        if not text.endswith("\n"):
            sys.stdout.write("\n")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())

\
#!/usr/bin/env python3
"""Inspect JAX tracing/lowering artefacts for a callable."""

from __future__ import annotations

import argparse
import hashlib
import importlib
import importlib.util
import json
import sys
from pathlib import Path
from typing import Any


def load_module(module_name: str | None, file_path: str | None) -> Any:
    if bool(module_name) == bool(file_path):
        raise ValueError("Exactly one of --module or --file is required.")
    if module_name:
        return importlib.import_module(module_name)

    path = Path(file_path or "")
    if not path.exists():
        raise FileNotFoundError(f"Module file not found: {path}")
    spec = importlib.util.spec_from_file_location(path.stem, path)
    if spec is None or spec.loader is None:
        raise ImportError(f"Could not load module spec from: {path}")
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


def resolve_attr(obj: Any, dotted: str) -> Any:
    current = obj
    for part in dotted.split("."):
        current = getattr(current, part)
    return current


def load_json_arg(raw: str | None, file_path: str | None, default: Any) -> Any:
    if raw is not None and file_path is not None:
        raise ValueError("Choose either the inline JSON form or the file form, not both.")
    if raw is not None:
        return json.loads(raw)
    if file_path is not None:
        return json.loads(Path(file_path).read_text(encoding="utf-8"))
    return default


def numeric_tree(value: Any) -> bool:
    if isinstance(value, (int, float, bool)):
        return True
    if isinstance(value, list):
        return all(numeric_tree(v) for v in value)
    return False


def maybe_import_jax():
    try:
        jax = importlib.import_module("jax")
        jnp = importlib.import_module("jax.numpy")
        return jax, jnp
    except Exception:
        return None, None


def tree_arrayify(value: Any, jnp_module: Any) -> Any:
    if isinstance(value, dict):
        return {k: tree_arrayify(v, jnp_module) for k, v in value.items()}
    if isinstance(value, list):
        if numeric_tree(value):
            return jnp_module.array(value)
        return [tree_arrayify(v, jnp_module) for v in value]
    return value


def parse_int_tuple(raw: str | None) -> tuple[int, ...] | None:
    if raw is None or raw == "":
        return None
    return tuple(int(part.strip()) for part in raw.split(",") if part.strip())


def sha256_text(text: str | None) -> str | None:
    if text is None:
        return None
    return hashlib.sha256(text.encode("utf-8")).hexdigest()


def to_text(obj: Any) -> str:
    if obj is None:
        return ""
    if isinstance(obj, str):
        return obj
    for getter in (
        lambda: obj.as_text(),
        lambda: obj.operation.get_asm(enable_debug_info=True),
        lambda: obj.operation.get_asm(),
    ):
        try:
            return getter()
        except Exception:
            pass
    return str(obj)


def preview_text(text: str, max_lines: int, max_chars: int) -> str:
    lines = text.splitlines()
    clipped = "\n".join(lines[:max_lines])
    if len(clipped) > max_chars:
        clipped = clipped[:max_chars]
    return clipped


def write_artifact(path: Path, text: str) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(text, encoding="utf-8")


def safe_repr(value: Any) -> str:
    try:
        return repr(value)
    except Exception as exc:  # pragma: no cover - best effort only
        return f"<repr failed: {type(exc).__name__}: {exc}>"


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Inspect `eval_shape`, jaxpr, lowering, and compiler IR for a callable.",
        formatter_class=argparse.RawTextHelpFormatter,
        epilog="""Exit codes:
  0 success
  2 operational error

Examples:
  python3 scripts/jax_compile_probe.py --file my_module.py --function step --args-file args.json --arrayify
  python3 scripts/jax_compile_probe.py --module pkg.train --function step --args-file args.json --jit --write-dir probe_out
  python3 scripts/jax_compile_probe.py --file my_module.py --function fn --args-json '[[1, 2, 3]]' --arrayify --dialects stablehlo,hlo
""",
    )
    source = parser.add_mutually_exclusive_group(required=True)
    source.add_argument("--module", help="Import path for the module containing the callable.")
    source.add_argument("--file", help="Path to a Python file containing the callable.")
    parser.add_argument("--function", required=True, help="Callable name or dotted attribute path.")
    parser.add_argument("--args-json", help="JSON list of positional arguments.")
    parser.add_argument("--args-file", help="Path to a JSON file containing positional arguments.")
    parser.add_argument("--kwargs-json", help="JSON object of keyword arguments.")
    parser.add_argument("--kwargs-file", help="Path to a JSON file containing keyword arguments.")
    parser.add_argument("--arrayify", action="store_true", help="Convert numeric JSON lists to `jax.numpy.array` when JAX is available.")
    parser.add_argument("--jit", action="store_true", help="Wrap the callable with `jax.jit` before lowering.")
    parser.add_argument("--static-argnums", help="Comma-separated positional indices to mark static when using --jit.")
    parser.add_argument("--donate-argnums", help="Comma-separated positional indices to donate when using --jit.")
    parser.add_argument("--compile", action="store_true", help="Compile the lowered program and report compile timing when possible.")
    parser.add_argument("--dialects", default="stablehlo", help="Comma-separated compiler IR dialects to attempt. Default: stablehlo")
    parser.add_argument("--max-preview-lines", type=int, default=80, help="Maximum preview lines per artefact. Default: 80")
    parser.add_argument("--max-preview-chars", type=int, default=12000, help="Maximum preview characters per artefact. Default: 12000")
    parser.add_argument("--write-dir", help="Optional directory for full artefacts such as jaxpr and IR text files.")
    parser.add_argument("--format", choices=("json", "text"), default="json", help="Output format. Default: json")
    parser.add_argument("--output", help="Write the report to a file instead of stdout.")
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    try:
        jax, jnp = maybe_import_jax()
        if jax is None or jnp is None:
            raise RuntimeError("JAX is not importable in this environment.")

        module = load_module(args.module, args.file)
        fn = resolve_attr(module, args.function)
        if not callable(fn):
            raise TypeError(f"Resolved object is not callable: {args.function}")

        raw_args = load_json_arg(args.args_json, args.args_file, [])
        raw_kwargs = load_json_arg(args.kwargs_json, args.kwargs_file, {})
        if not isinstance(raw_args, list):
            raise TypeError("Positional arguments JSON must decode to a list.")
        if not isinstance(raw_kwargs, dict):
            raise TypeError("Keyword arguments JSON must decode to an object.")

        proc_args = list(raw_args)
        proc_kwargs = dict(raw_kwargs)
        if args.arrayify:
            proc_args = [tree_arrayify(v, jnp) for v in proc_args]
            proc_kwargs = {k: tree_arrayify(v, jnp) for k, v in proc_kwargs.items()}

        static_argnums = parse_int_tuple(args.static_argnums)
        donate_argnums = parse_int_tuple(args.donate_argnums)

        report: dict[str, Any] = {
            "source": args.module or args.file,
            "callable": args.function,
            "jit_requested": args.jit,
            "static_argnums": static_argnums,
            "donate_argnums": donate_argnums,
            "dialects_requested": [d.strip() for d in args.dialects.split(",") if d.strip()],
        }

        try:
            eval_shape = jax.eval_shape(fn, *proc_args, **proc_kwargs)
            report["eval_shape_repr"] = safe_repr(eval_shape)
        except Exception as exc:
            report["eval_shape_error"] = f"{type(exc).__name__}: {exc}"

        try:
            jaxpr_obj = jax.make_jaxpr(fn)(*proc_args, **proc_kwargs)
            jaxpr_text = str(jaxpr_obj)
            report["jaxpr"] = {
                "lines": len(jaxpr_text.splitlines()),
                "sha256": sha256_text(jaxpr_text),
                "preview": preview_text(jaxpr_text, args.max_preview_lines, args.max_preview_chars),
            }
            if args.write_dir:
                write_artifact(Path(args.write_dir) / "jaxpr.txt", jaxpr_text)
        except Exception as exc:
            report["jaxpr_error"] = f"{type(exc).__name__}: {exc}"

        lowered_fn = fn
        if args.jit:
            jit_kwargs = {}
            if static_argnums is not None:
                jit_kwargs["static_argnums"] = static_argnums
            if donate_argnums is not None:
                jit_kwargs["donate_argnums"] = donate_argnums
            lowered_fn = jax.jit(fn, **jit_kwargs)

        try:
            if not hasattr(lowered_fn, "lower"):
                raise TypeError("Target does not support `.lower(...)`; try using --jit.")
            lowered = lowered_fn.lower(*proc_args, **proc_kwargs)
            lowered_text = to_text(lowered)
            report["lowering"] = {
                "sha256": sha256_text(lowered_text),
                "preview": preview_text(lowered_text, args.max_preview_lines, args.max_preview_chars),
            }
            if args.write_dir:
                write_artifact(Path(args.write_dir) / "lowering.txt", lowered_text)

            ir_reports = {}
            for dialect in report["dialects_requested"]:
                try:
                    ir_obj = lowered.compiler_ir(dialect=dialect)
                    ir_text = to_text(ir_obj)
                    ir_reports[dialect] = {
                        "sha256": sha256_text(ir_text),
                        "preview": preview_text(ir_text, args.max_preview_lines, args.max_preview_chars),
                    }
                    if args.write_dir:
                        suffix = "mlir" if dialect in {"stablehlo", "mhlo"} else "txt"
                        write_artifact(Path(args.write_dir) / f"{dialect}.{suffix}", ir_text)
                except Exception as exc:
                    ir_reports[dialect] = {"error": f"{type(exc).__name__}: {exc}"}
            report["compiler_ir"] = ir_reports

            if args.compile:
                import time

                t0 = time.perf_counter()
                compiled = lowered.compile()
                compile_ms = (time.perf_counter() - t0) * 1e3
                report["compile"] = {
                    "ok": compiled is not None,
                    "elapsed_ms": compile_ms,
                }
        except Exception as exc:
            report["lowering_error"] = f"{type(exc).__name__}: {exc}"

    except Exception as exc:
        sys.stderr.write(f"Error: {type(exc).__name__}: {exc}\n")
        return 2

    if args.format == "json":
        text = json.dumps(report, indent=2, sort_keys=True)
    else:
        lines = [
            f"Callable: {report['source']}::{report['callable']}",
            f"JIT requested: {report['jit_requested']}",
            f"Static argnums: {report['static_argnums']}",
            f"Donate argnums: {report['donate_argnums']}",
            "",
        ]
        if "eval_shape_repr" in report:
            lines.append("eval_shape:")
            lines.append(report["eval_shape_repr"])
            lines.append("")
        if "eval_shape_error" in report:
            lines.append(f"eval_shape_error: {report['eval_shape_error']}")
            lines.append("")
        if "jaxpr" in report:
            lines.append(f"jaxpr lines: {report['jaxpr']['lines']}")
            lines.append(report["jaxpr"]["preview"])
            lines.append("")
        if "lowering" in report:
            lines.append("lowering preview:")
            lines.append(report["lowering"]["preview"])
            lines.append("")
        if "compiler_ir" in report:
            for dialect, info in report["compiler_ir"].items():
                lines.append(f"{dialect}:")
                if "error" in info:
                    lines.append(f"  error: {info['error']}")
                else:
                    lines.append(info["preview"])
                lines.append("")
        if "compile" in report:
            lines.append(f"compile elapsed_ms: {report['compile']['elapsed_ms']:.3f}")
        text = "\n".join(lines)

    if args.output:
        Path(args.output).write_text(text + ("" if text.endswith("\n") else "\n"), encoding="utf-8")
    else:
        sys.stdout.write(text)
        if not text.endswith("\n"):
            sys.stdout.write("\n")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())

\
#!/usr/bin/env python3
"""Emit a structured report about the local JAX environment.

Design goals:
- non-interactive
- JSON by default
- standard-library only
- useful even when JAX is not importable
"""

from __future__ import annotations

import argparse
import importlib
import importlib.metadata
import json
import os
import platform
import sys
import time
from typing import Any

PACKAGE_NAMES = (
    "jax",
    "jaxlib",
    "numpy",
    "scipy",
    "flax",
    "optax",
    "equinox",
    "orbax-checkpoint",
)

CONFIG_KEYS = (
    "jax_enable_x64",
    "jax_default_matmul_precision",
    "jax_debug_nans",
    "jax_debug_infs",
    "jax_debug_key_reuse",
    "jax_platform_name",
    "jax_default_prng_impl",
    "jax_transfer_guard",
    "jax_compilation_cache_dir",
)

ENV_PREFIXES = (
    "JAX_",
    "XLA_",
    "CUDA_",
    "NVIDIA_",
    "ROCM",
    "HIP_",
    "TPU",
    "NCCL_",
)

ENV_NAMES = {
    "CUDA_VISIBLE_DEVICES",
    "NVIDIA_VISIBLE_DEVICES",
    "XLA_FLAGS",
    "PYTHONPATH",
    "LD_LIBRARY_PATH",
    "PATH",
}


def package_info(name: str) -> dict[str, Any]:
    out: dict[str, Any] = {"installed": False}
    try:
        out["version"] = importlib.metadata.version(name)
        out["installed"] = True
    except importlib.metadata.PackageNotFoundError:
        out["version"] = None
    except Exception as exc:  # pragma: no cover - best effort only
        out["version"] = None
        out["error"] = f"{type(exc).__name__}: {exc}"
    return out


def selected_environment() -> dict[str, str]:
    env = {}
    for key, value in os.environ.items():
        if key in ENV_NAMES or any(key.startswith(prefix) for prefix in ENV_PREFIXES):
            env[key] = value
    return dict(sorted(env.items()))


def read_config_value(jax_module: Any, key: str) -> Any:
    cfg = getattr(jax_module, "config", None)
    if cfg is None:
        return None

    # Try the most stable public-ish access patterns first.
    for getter in (
        lambda: getattr(cfg, "values", {}).get(key),
        lambda: cfg.read(key),  # type: ignore[attr-defined]
        lambda: getattr(cfg, key),
    ):
        try:
            value = getter()
            if value is not None:
                return value
        except Exception:
            pass
    return None


def maybe_key(jax_module: Any, seed: int):
    try:
        return jax_module.random.key(seed)
    except Exception:
        return jax_module.random.PRNGKey(seed)


def smoke_test(jax_module: Any, jnp_module: Any) -> dict[str, Any]:
    report: dict[str, Any] = {"ok": False}
    t0 = time.perf_counter()

    key = maybe_key(jax_module, 0)
    x = jax_module.random.normal(key, (256, 256), dtype=jnp_module.float32)
    y = (x @ x.T).block_until_ready()

    @jax_module.jit
    def loss_fn(z):
        return jnp_module.sum(jnp_module.tanh(z @ z.T))

    loss = loss_fn(x)
    loss.block_until_ready()
    grad = jax_module.grad(lambda z: jnp_module.sum(jnp_module.sin(z)))(x)
    jax_module.block_until_ready(grad)

    report["ok"] = True
    report["elapsed_ms"] = (time.perf_counter() - t0) * 1e3
    report["matmul_shape"] = tuple(int(v) for v in y.shape)
    report["loss_dtype"] = str(loss.dtype)
    return report


def load_jax_report(run_smoke_test: bool) -> dict[str, Any]:
    report: dict[str, Any] = {"imported": False}
    try:
        jax = importlib.import_module("jax")
        jnp = importlib.import_module("jax.numpy")
    except Exception as exc:
        report["import_error"] = f"{type(exc).__name__}: {exc}"
        return report

    report["imported"] = True
    report["version"] = getattr(jax, "__version__", None)

    try:
        report["default_backend"] = jax.default_backend()
    except Exception:
        report["default_backend"] = None

    try:
        report["process_count"] = int(jax.process_count())
        report["process_index"] = int(jax.process_index())
        report["device_count"] = int(jax.device_count())
        report["local_device_count"] = int(jax.local_device_count())
    except Exception as exc:
        report["process_error"] = f"{type(exc).__name__}: {exc}"

    devices = []
    try:
        for dev in jax.devices():
            devices.append(
                {
                    "id": getattr(dev, "id", None),
                    "platform": getattr(dev, "platform", None),
                    "device_kind": getattr(dev, "device_kind", None),
                    "process_index": getattr(dev, "process_index", None),
                    "memory_limit": getattr(dev, "memory_limit", None),
                }
            )
    except Exception as exc:
        report["devices_error"] = f"{type(exc).__name__}: {exc}"
    report["devices"] = devices

    cfg = {}
    for key in CONFIG_KEYS:
        cfg[key] = read_config_value(jax, key)
    report["config"] = cfg

    if run_smoke_test:
        try:
            report["smoke_test"] = smoke_test(jax, jnp)
        except Exception as exc:
            report["smoke_test"] = {"ok": False, "error": f"{type(exc).__name__}: {exc}"}

    return report


def build_report(run_smoke_test: bool) -> dict[str, Any]:
    return {
        "python": {
            "version": sys.version.split()[0],
            "executable": sys.executable,
        },
        "platform": {
            "system": platform.system(),
            "release": platform.release(),
            "machine": platform.machine(),
            "platform": platform.platform(),
        },
        "packages": {name: package_info(name) for name in PACKAGE_NAMES},
        "environment": selected_environment(),
        "jax": load_jax_report(run_smoke_test),
    }


def format_text(report: dict[str, Any]) -> str:
    lines = []
    lines.append("Python")
    lines.append(f"  version: {report['python']['version']}")
    lines.append(f"  executable: {report['python']['executable']}")
    lines.append("")
    lines.append("Platform")
    for key, value in report["platform"].items():
        lines.append(f"  {key}: {value}")
    lines.append("")
    lines.append("Packages")
    for name, info in report["packages"].items():
        status = info.get("version") if info.get("installed") else "not installed"
        lines.append(f"  {name}: {status}")
    lines.append("")
    lines.append("Environment")
    for key, value in report["environment"].items():
        lines.append(f"  {key}={value}")
    lines.append("")

    jax_info = report["jax"]
    lines.append("JAX")
    if not jax_info.get("imported"):
        lines.append(f"  import_failed: {jax_info.get('import_error')}")
        return "\n".join(lines)

    lines.append(f"  version: {jax_info.get('version')}")
    lines.append(f"  default_backend: {jax_info.get('default_backend')}")
    lines.append(f"  process_count: {jax_info.get('process_count')}")
    lines.append(f"  process_index: {jax_info.get('process_index')}")
    lines.append(f"  device_count: {jax_info.get('device_count')}")
    lines.append(f"  local_device_count: {jax_info.get('local_device_count')}")
    lines.append("")
    lines.append("Devices")
    for device in jax_info.get("devices", []):
        lines.append(
            "  - id={id} platform={platform} kind={device_kind} process_index={process_index} memory_limit={memory_limit}".format(
                **device
            )
        )
    lines.append("")
    lines.append("Config")
    for key, value in jax_info.get("config", {}).items():
        lines.append(f"  {key}: {value}")

    smoke = jax_info.get("smoke_test")
    if smoke:
        lines.append("")
        lines.append("Smoke test")
        for key, value in smoke.items():
            lines.append(f"  {key}: {value}")

    return "\n".join(lines)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Emit a structured report about the local JAX environment.",
        formatter_class=argparse.RawTextHelpFormatter,
        epilog="""Exit codes:
  0 success
  2 operational error

Examples:
  python3 scripts/jax_env_report.py
  python3 scripts/jax_env_report.py --smoke-test --format text
  python3 scripts/jax_env_report.py --output env.json
""",
    )
    parser.add_argument("--smoke-test", action="store_true", help="Run a small JAX smoke test if JAX is importable.")
    parser.add_argument("--format", choices=("json", "text"), default="json", help="Output format. Default: json")
    parser.add_argument("--output", help="Write output to a file instead of stdout.")
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    try:
        report = build_report(args.smoke_test)
        text = json.dumps(report, indent=2, sort_keys=True) if args.format == "json" else format_text(report)
        if args.output:
            with open(args.output, "w", encoding="utf-8") as f:
                f.write(text)
                if not text.endswith("\n"):
                    f.write("\n")
        else:
            sys.stdout.write(text)
            if not text.endswith("\n"):
                sys.stdout.write("\n")
        return 0
    except Exception as exc:
        sys.stderr.write(f"Error: {type(exc).__name__}: {exc}\n")
        return 2


if __name__ == "__main__":
    raise SystemExit(main())