Add Bazel build rules and runtime implementation for AWS Neuron/Trainium/Inferentia support.

MODULE.bazel

@@ -11,10 +11,10 @@ bazel_dep(name = "platforms", version = "0.0.10")
 bazel_dep(name = "rules_cc", version = "0.0.9")
 bazel_dep(name = "rules_pkg", version = "1.0.1")
 bazel_dep(name = "rules_proto", version = "6.0.2")
-
+bazel_dep(name = "with_cfg.bzl", version = "0.5.0")
 bazel_dep(name = "buildifier_prebuilt", version = "6.4.0", dev_dependency = True)
-
+bazel_dep(name = "aspect_bazel_lib", version = "2.9.3")
-bazel_dep(name = "aspect_bazel_lib", version = "2.8.1.1")
+bazel_dep(name = "rules_distroless", version = "0.3.8")
 
 bazel_lib_toolchains = use_extension("@aspect_bazel_lib//lib:extensions.bzl", "toolchains", dev_dependency = True)
 use_repo(bazel_lib_toolchains, "jq_toolchains")
@@ -41,6 +41,29 @@ register_toolchains(
     "@zig_sdk//toolchain:linux_arm64_gnu.2.31",
 )
 
+bazel_dep(name = "rules_python", version = "0.37.0")
+bazel_dep(name = "rules_uv", version = "0.34.0")
+bazel_dep(name = "aspect_rules_py", version = "1.0.0-rc0")
+python = use_extension("@rules_python//python/extensions:python.bzl", "python")
+python.toolchain(python_version = "3.11")
+use_repo(python, "python_versions")
+
+pip = use_extension("@rules_python//python/extensions:pip.bzl", "pip")
+pip.parse(
+    hub_name = "neuron_py_deps",
+    python_version = "3.11",
+    download_only = True,
+    extra_pip_args = [
+        "--abi=cp311",
+        "--implementation=cp",
+        "--python-version=311",
+        "--platform=linux_x86_64",
+        "--platform=manylinux2014_x86_64",
+    ],
+    requirements_lock = "//runtimes/neuron:requirements.lock.txt",
+)
+use_repo(pip, "neuron_py_deps")
+
 cpu = use_extension("//runtimes/cpu:cpu.bzl", "cpu_pjrt_plugin")
 use_repo(cpu, "libpjrt_cpu_darwin_arm64", "libpjrt_cpu_linux_amd64")
 
@@ -53,6 +76,9 @@ use_repo(rocm, "libpjrt_rocm")
 tpu = use_extension("//runtimes/tpu:tpu.bzl", "tpu_packages")
 use_repo(tpu, "libpjrt_tpu")
 
+neuron = use_extension("//runtimes/neuron:neuron.bzl", "neuron_packages")
+use_repo(neuron, "aws-neuronx-collectives", "aws-neuronx-runtime-lib")
+
 zls = use_extension("//third_party/zls:zls.bzl", "repo")
 use_repo(zls, "zls_aarch64-macos", "zls_x86_64-linux")
 
@@ -81,3 +107,11 @@ bazel_dep(name = "zigcoro", version = "20240829.0-fc1db29")
 bazel_dep(name = "sentencepiece", version = "20240618.0-d7ace0a")
 bazel_dep(name = "zig-protobuf", version = "20240722.0-c644d11")
 bazel_dep(name = "zig-yaml", version = "20240903.0-83d5fdf")
+
+apt = use_extension("@rules_distroless//apt:extensions.bzl", "apt")
+apt.install(
+    name = "neuron_bookworm",
+    lock = "//runtimes/neuron:packages.lock.json",
+    manifest = "//runtimes/neuron:packages.yaml",
+)
+use_repo(apt, "neuron_bookworm")

bazel/runfiles.bzl

@@ -0,0 +1,17 @@
+def _runfiles_to_default(ctx):
+    files = depset([], transitive = [
+        dep[DefaultInfo].default_runfiles.files
+        for dep in ctx.attr.deps
+    ])
+    return [
+        DefaultInfo(
+            files = files,
+        ),
+    ]
+
+runfiles_to_default = rule(
+    implementation = _runfiles_to_default,
+    attrs = {
+        "deps": attr.label_list(providers = [DefaultInfo]),  # We expect DefaultInfo from dependencies
+    },
+)

runtimes/BUILD.bazel

@@ -1,4 +1,6 @@
+load("@aspect_bazel_lib//lib:tar.bzl", "tar")
 load("@bazel_skylib//rules:common_settings.bzl", "bool_flag")
+load("@bazel_skylib//rules:write_file.bzl", "write_file")
 load("@rules_zig//zig:defs.bzl", "zig_library")
 
 RUNTIMES = {
@@ -6,6 +8,7 @@ RUNTIMES = {
     "cuda": False,
     "rocm": False,
     "tpu": False,
+    "neuron": False,
 }
 
 [
@@ -25,6 +28,39 @@ RUNTIMES = {
     for runtime in RUNTIMES.keys()
 ]
 
+write_file(
+    name = "zml.txt",
+    out = "zml.txt",
+    content = ["ZML loves you <3"],
+)
+
+tar(
+    name = "empty",
+    srcs = [":zml.txt"],
+    mtree = "auto",
+)
+
+filegroup(
+    name = "layers",
+    srcs = [":empty"] + select({
+        ":cpu.enabled": ["//runtimes/cpu:layers"],
+        "//conditions:default": [],
+    }) + select({
+        ":cuda.enabled": ["//runtimes/cuda:layers"],
+        "//conditions:default": [],
+    }) + select({
+        ":neuron.enabled": ["//runtimes/neuron:layers"],
+        "//conditions:default": [],
+    }) + select({
+        ":rocm.enabled": ["//runtimes/rocm:layers"],
+        "//conditions:default": [],
+    }) + select({
+        ":tpu.enabled": ["//runtimes/tpu:layers"],
+        "//conditions:default": [],
+    }),
+    visibility = ["//visibility:public"],
+)
+
 zig_library(
     name = "runtimes",
     main = "runtimes.zig",

runtimes/cpu/BUILD.bazel

@@ -25,3 +25,9 @@ zig_library(
         "//conditions:default": [":empty"],
     }),
 )
+
+filegroup(
+    name = "layers",
+    srcs = [],
+    visibility = ["//visibility:public"],
+)

runtimes/cuda/BUILD.bazel

@@ -25,3 +25,9 @@ zig_library(
         "//conditions:default": [":empty"],
     }),
 )
+
+filegroup(
+    name = "layers",
+    srcs = [],
+    visibility = ["//visibility:public"],
+)

runtimes/neuron/BUILD.bazel

@@ -0,0 +1,156 @@
+load("@bazel_skylib//rules:select_file.bzl", "select_file")
+load("@rules_python//python:pip.bzl", "compile_pip_requirements")
+load("@rules_python//python/entry_points:py_console_script_binary.bzl", "py_console_script_binary")
+load("@rules_uv//uv:pip.bzl", "pip_compile")
+load("@rules_zig//zig:defs.bzl", "zig_library")
+load("@zml//bazel:cc_import.bzl", "cc_import")
+load("@zml//bazel:runfiles.bzl", "runfiles_to_default")
+load(":neuron.bzl", "py_binary_with_script")
+load(":pyenv.bzl", "pyenv_zig")
+
+pyenv_zig(
+    name = "libneuronxla_pyenv_zig",
+    deps = [":libneuronxla"],
+)
+
+zig_library(
+    name = "libneuronxla_pyenv",
+    main = ":libneuronxla_pyenv_zig",
+)
+
+pip_compile(
+    name = "update_requirements",
+    args = [
+        "--generate-hashes",
+        "--emit-index-url",
+        "--emit-find-links",
+        "--no-strip-extras",
+        "--index-strategy unsafe-best-match",
+        "--upgrade",
+    ],
+    python_platform = "x86_64-unknown-linux-gnu",
+    requirements_in = "requirements.in",
+    requirements_txt = "requirements.lock.txt",
+)
+
+py_console_script_binary(
+    name = "neuronx-cc",
+    binary_rule = py_binary_with_script,
+    pkg = "@neuron_py_deps//neuronx_cc",
+    visibility = ["//visibility:public"],
+)
+
+py_binary_with_script(
+    name = "libneuronxla",
+    srcs = ["empty.py"],
+    main = "empty.py",
+    deps = [
+        ":neuronx-cc",
+        "@neuron_py_deps//libneuronxla",
+    ],
+)
+
+cc_library(
+    name = "libpython",
+    hdrs = ["libpython.h"],
+    visibility = ["//visibility:public"],
+    deps = [
+        "@rules_python//python/cc:current_py_cc_headers",
+        "@rules_python//python/cc:current_py_cc_libs",
+    ],
+)
+
+compile_pip_requirements(
+    name = "requirements",
+    src = "requirements.in",
+    py_binary = py_binary_with_script,
+    requirements_txt = "requirements_lock.txt",
+)
+
+runfiles_to_default(
+    name = "libneuronxla_files",
+    deps = ["@neuron_py_deps//libneuronxla:pkg"],
+)
+
+select_file(
+    name = "libneuronpjrt_so",
+    srcs = ":libneuronxla_files",
+    subpath = "site-packages/libneuronxla/libneuronpjrt.so",
+)
+
+cc_import(
+    name = "libpjrt_neuron_",
+    add_needed = [
+        "libpython3.11.so.1.0",
+        "libzmlxneuron.so.0",
+    ],
+    data = [":libneuronxla"],
+    shared_library = ":libneuronpjrt_so",
+    soname = "libpjrt_neuron.so",
+    deps = [
+        ":zmlxneuron",
+        "@aws-neuronx-runtime-lib",
+        "@rules_python//python/cc:current_py_cc_libs",
+    ],
+)
+
+cc_library(
+    name = "zmlxneuron_lib",
+    srcs = ["zmlxneuron.c"],
+    linkopts = ["-ldl"],
+)
+
+cc_shared_library(
+    name = "zmlxneuron_",
+    shared_lib_name = "libzmlxneuron.so.0",
+    deps = [":zmlxneuron_lib"],
+)
+
+cc_import(
+    name = "zmlxneuron",
+    shared_library = ":zmlxneuron_",
+)
+
+alias(
+    name = "bootstrap",
+    actual = "@rules_python//python/config_settings:bootstrap_impl",
+)
+
+cc_library(
+    name = "empty",
+)
+
+cc_library(
+    name = "libpjrt_neuron",
+    defines = ["ZML_RUNTIME_NEURON"],
+    deps = [":libpjrt_neuron_"],
+)
+
+zig_library(
+    name = "neuron",
+    import_name = "runtimes/neuron",
+    main = "neuron.zig",
+    visibility = ["//visibility:public"],
+    deps = [
+        "//pjrt",
+    ] + select({
+        "//runtimes:neuron.enabled": [
+            ":libneuronxla_pyenv",
+            ":libpjrt_neuron",
+            ":libpython",
+            "//async",
+            "@rules_zig//zig/runfiles",
+        ],
+        "//conditions:default": [":empty"],
+    }),
+)
+
+filegroup(
+    name = "layers",
+    srcs = [
+        "@neuron_bookworm//bash/amd64",
+        "@neuron_bookworm//libxml2/amd64",
+        "@neuron_bookworm//zlib1g/amd64",
+    ],
+    visibility = ["//visibility:public"],
+)

runtimes/neuron/libpython.h

@@ -0,0 +1,2 @@
+#define PY_SSIZE_T_CLEAN
+#include <Python.h>

runtimes/neuron/neuron.bzl

@@ -0,0 +1,65 @@
+load("@python_versions//3.11:defs.bzl", _py_binary = "py_binary")
+load("@rules_python//python:defs.bzl", "PyInfo")
+load("@with_cfg.bzl", "with_cfg")
+load("//bazel:http_deb_archive.bzl", "http_deb_archive")
+
+BASE_URL = "https://apt.repos.neuron.amazonaws.com"
+STRIP_PREFIX = "opt/aws/neuron"
+
+_PACKAGES = {
+    "aws-neuronx-runtime-lib": """\
+load("@zml//bazel:cc_import.bzl", "cc_import")
+cc_import(
+    name = "aws-neuronx-runtime-lib",
+    shared_library = "lib/libnrt.so.1",
+    rename_dynamic_symbols = {
+        "dlopen": "zmlxneuron_dlopen",
+    },
+    visibility = ["@zml//runtimes/neuron:__subpackages__"],
+    deps = ["@aws-neuronx-collectives//:libnccom"],
+)
+""",
+    "aws-neuronx-collectives": """\
+cc_import(
+    name = "libnccom",
+    shared_library = "lib/libnccom.so.2",
+    visibility = ["@aws-neuronx-runtime-lib//:__subpackages__"],
+)
+
+cc_import(
+    name = "libnccom-net",
+    shared_library = "lib/libnccom-net.so.0",
+)
+""",
+}
+
+def _neuron_impl(mctx):
+    packages_json = json.decode(mctx.read(Label("@zml//runtimes/neuron:packages.lock.json")))
+    packages = {
+        pkg["name"]: pkg
+        for pkg in packages_json["packages"]
+    }
+    for pkg_name, build_file_content in _PACKAGES.items():
+        pkg = packages[pkg_name]
+        http_deb_archive(
+            name = pkg_name,
+            urls = [pkg["url"]],
+            sha256 = pkg["sha256"],
+            strip_prefix = STRIP_PREFIX,
+            build_file_content = build_file_content,
+        )
+
+    return mctx.extension_metadata(
+        reproducible = True,
+        root_module_direct_deps = "all",
+        root_module_direct_dev_deps = [],
+    )
+
+neuron_packages = module_extension(
+    implementation = _neuron_impl,
+)
+
+py_binary_with_script, _py_binary_internal = with_cfg(
+    kind = _py_binary,
+    extra_providers = [PyInfo],
+).set(Label("@rules_python//python/config_settings:bootstrap_impl"), "script").build()

runtimes/neuron/neuron.zig

@@ -0,0 +1,138 @@
+const builtin = @import("builtin");
+const asynk = @import("async");
+const pjrt = @import("pjrt");
+const c = @import("c");
+const std = @import("std");
+const runfiles = @import("runfiles");
+const bazel_builtin = @import("bazel_builtin");
+const libneuronxla_pyenv = @import("libneuronxla_pyenv");
+
+pub fn isEnabled() bool {
+    return @hasDecl(c, "ZML_RUNTIME_NEURON");
+}
+
+fn hasNeuronDevice() bool {
+    asynk.File.access("/dev/neuron0", .{ .mode = .read_only }) catch return false;
+    return true;
+}
+
+fn isRunningOnEC2() !bool {
+    const AmazonEC2 = "Amazon EC2";
+
+    var f = try asynk.File.open("/sys/devices/virtual/dmi/id/sys_vendor", .{ .mode = .read_only });
+    defer f.close() catch {};
+
+    var buf = [_]u8{0} ** AmazonEC2.len;
+    _ = try f.reader().readAll(&buf);
+
+    return std.mem.eql(u8, &buf, AmazonEC2);
+}
+
+fn toWchar(str: []const u8, out: [:0]c.wchar_t) [:0]c.wchar_t {
+    const len = c.mbstowcs(out.ptr, str.ptr, str.len);
+    out[len] = 0;
+    return out[0..len :0];
+}
+
+fn pyErrorOrExit(status: c.PyStatus) void {
+    if (c.PyStatus_Exception(status) != 0) {
+        if (c.PyStatus_IsExit(status) != 0) {
+            std.process.exit(@intCast(status.exitcode));
+        }
+        c.Py_ExitStatusException(status);
+    }
+}
+
+fn initialize() !void {
+    var arena = std.heap.ArenaAllocator.init(std.heap.c_allocator);
+    const allocator = arena.allocator();
+    defer arena.deinit();
+
+    {
+        var preconfig: c.PyPreConfig = undefined;
+        c.PyPreConfig_InitIsolatedConfig(&preconfig);
+        preconfig.utf8_mode = 1;
+        pyErrorOrExit(c.Py_PreInitialize(&preconfig));
+    }
+
+    var config: c.PyConfig = undefined;
+    c.PyConfig_InitIsolatedConfig(&config);
+    defer c.PyConfig_Clear(&config);
+
+    var r_ = try runfiles.Runfiles.create(.{ .allocator = allocator }) orelse return error.Unavailable;
+    const r = r_.withSourceRepo(bazel_builtin.current_repository);
+
+    var buf: [std.fs.max_path_bytes]u8 = undefined;
+    var wbuf: [std.fs.max_path_bytes:0]c.wchar_t = undefined;
+
+    {
+        const path = (try r.rlocation(libneuronxla_pyenv.home, &buf)).?;
+        const wpath = toWchar(std.fs.path.dirname(path).?, &wbuf);
+        pyErrorOrExit(c.PyConfig_SetString(&config, &config.home, wpath.ptr));
+    }
+
+    {
+        config.module_search_paths_set = 1;
+        for (libneuronxla_pyenv.modules) |module| {
+            const path = (try r.rlocation(module, &buf)).?;
+            const wline = toWchar(std.fs.path.dirname(path).?, &wbuf);
+            pyErrorOrExit(c.PyWideStringList_Append(&config.module_search_paths, wline.ptr));
+        }
+    }
+
+    {
+        const neuronx_cc = (try r.rlocation("zml/runtimes/neuron/neuronx-cc/neuronx-cc", &buf)).?;
+        const neuronx_cc_path = std.fs.path.dirname(neuronx_cc).?;
+        const path = std.posix.getenv("PATH") orelse "";
+        const new_path = try std.fmt.allocPrintZ(allocator, "{s}:{s}", .{ neuronx_cc_path, path });
+        _ = c.setenv("PATH", new_path.ptr, 1);
+    }
+
+    pyErrorOrExit(c.Py_InitializeFromConfig(&config));
+
+    // release the GIL
+    _ = c.PyEval_SaveThread();
+}
+
+fn comptimeStrJoin(comptime separator: [:0]const u8, comptime slices: []const [:0]const u8) [:0]const u8 {
+    comptime var ret = slices[0];
+    inline for (slices[1..]) |slice| {
+        ret = ret ++ separator ++ slice;
+    }
+    return ret;
+}
+
+pub fn setNeuronCCFlags() void {
+    _ = c.setenv("NEURON_CC_FLAGS", comptimeStrJoin(" ", &.{
+        // 30% faster, no visible speed difference on llama
+        "--optlevel=1",
+        // generic is the default, but it fails on transformers, force it
+        "--model-type=transformer",
+        // disable it, we do our own
+        "--auto-cast=none",
+        "--enable-fast-loading-neuron-binaries",
+    }), 1);
+
+    // Enable stochastic rounding
+    // https://awsdocs-neuron.readthedocs-hosted.com/en/latest/general/arch/neuron-features/rounding-modes.html
+    _ = c.setenv("NEURON_RT_STOCHASTIC_ROUNDING_EN", "1", 1);
+}
+
+pub fn load() !*const pjrt.Api {
+    if (comptime !isEnabled()) {
+        return error.Unavailable;
+    }
+    if (comptime builtin.os.tag != .linux) {
+        return error.Unavailable;
+    }
+    if (!(isRunningOnEC2() catch false)) {
+        return error.Unavailable;
+    }
+    if (!hasNeuronDevice()) {
+        return error.Unavailable;
+    }
+
+    setNeuronCCFlags();
+    try initialize();
+    return try pjrt.Api.loadFrom("libpjrt_neuron.so");
+}

runtimes/neuron/packages.lock.json

@@ -0,0 +1,214 @@
+{
+	"packages": [
+		{
+			"arch": "amd64",
+			"dependencies": [
+				{
+					"key": "debianutils_5.7-0.5_deb12u1_amd64",
+					"name": "debianutils",
+					"version": "5.7-0.5~deb12u1"
+				},
+				{
+					"key": "libc6_2.36-9-p-deb12u9_amd64",
+					"name": "libc6",
+					"version": "2.36-9+deb12u9"
+				},
+				{
+					"key": "libgcc-s1_12.2.0-14_amd64",
+					"name": "libgcc-s1",
+					"version": "12.2.0-14"
+				},
+				{
+					"key": "gcc-12-base_12.2.0-14_amd64",
+					"name": "gcc-12-base",
+					"version": "12.2.0-14"
+				},
+				{
+					"key": "base-files_12.4-p-deb12u8_amd64",
+					"name": "base-files",
+					"version": "12.4+deb12u8"
+				},
+				{
+					"key": "libtinfo6_6.4-4_amd64",
+					"name": "libtinfo6",
+					"version": "6.4-4"
+				}
+			],
+			"key": "bash_5.2.15-2-p-b7_amd64",
+			"name": "bash",
+			"sha256": "d3c65202568384a588ea32017b243fe64ca054943a008cba4e5bd8f9e79975c4",
+			"url": "https://deb.debian.org/debian/pool/main/b/bash/bash_5.2.15-2+b7_amd64.deb",
+			"version": "5.2.15-2+b7"
+		},
+		{
+			"arch": "amd64",
+			"dependencies": [],
+			"key": "debianutils_5.7-0.5_deb12u1_amd64",
+			"name": "debianutils",
+			"sha256": "55f951359670eb3236c9e2ccd5fac9ccb3db734f5a22aff21589e7a30aee48c9",
+			"url": "https://deb.debian.org/debian/pool/main/d/debianutils/debianutils_5.7-0.5~deb12u1_amd64.deb",
+			"version": "5.7-0.5~deb12u1"
+		},
+		{
+			"arch": "amd64",
+			"dependencies": [],
+			"key": "libc6_2.36-9-p-deb12u9_amd64",
+			"name": "libc6",
+			"sha256": "3a043b9dbf1a8c8b5a9e2268253e6dca8c7d431bd7a202fea96364abf374fea9",
+			"url": "https://deb.debian.org/debian/pool/main/g/glibc/libc6_2.36-9+deb12u9_amd64.deb",
+			"version": "2.36-9+deb12u9"
+		},
+		{
+			"arch": "amd64",
+			"dependencies": [],
+			"key": "libgcc-s1_12.2.0-14_amd64",
+			"name": "libgcc-s1",
+			"sha256": "f3d1d48c0599aea85b7f2077a01d285badc42998c1a1e7473935d5cf995c8141",
+			"url": "https://deb.debian.org/debian/pool/main/g/gcc-12/libgcc-s1_12.2.0-14_amd64.deb",
+			"version": "12.2.0-14"
+		},
+		{
+			"arch": "amd64",
+			"dependencies": [],
+			"key": "gcc-12-base_12.2.0-14_amd64",
+			"name": "gcc-12-base",
+			"sha256": "1a03df5a57833d65b5bb08cfa19d50e76f29088dc9e64fb934af42d9023a0807",
+			"url": "https://deb.debian.org/debian/pool/main/g/gcc-12/gcc-12-base_12.2.0-14_amd64.deb",
+			"version": "12.2.0-14"
+		},
+		{
+			"arch": "amd64",
+			"dependencies": [],
+			"key": "base-files_12.4-p-deb12u8_amd64",
+			"name": "base-files",
+			"sha256": "b961a69d1b79f093597d89a96f296ddca17d3bab9d166361c42f07a6beb4f685",
+			"url": "https://deb.debian.org/debian/pool/main/b/base-files/base-files_12.4+deb12u8_amd64.deb",
+			"version": "12.4+deb12u8"
+		},
+		{
+			"arch": "amd64",
+			"dependencies": [],
+			"key": "libtinfo6_6.4-4_amd64",
+			"name": "libtinfo6",
+			"sha256": "072d908f38f51090ca28ca5afa3b46b2957dc61fe35094c0b851426859a49a51",
+			"url": "https://deb.debian.org/debian/pool/main/n/ncurses/libtinfo6_6.4-4_amd64.deb",
+			"version": "6.4-4"
+		},
+		{
+			"arch": "amd64",
+			"dependencies": [
+				{
+					"key": "libc6_2.36-9-p-deb12u9_amd64",
+					"name": "libc6",
+					"version": "2.36-9+deb12u9"
+				},
+				{
+					"key": "libgcc-s1_12.2.0-14_amd64",
+					"name": "libgcc-s1",
+					"version": "12.2.0-14"
+				},
+				{
+					"key": "gcc-12-base_12.2.0-14_amd64",
+					"name": "gcc-12-base",
+					"version": "12.2.0-14"
+				}
+			],
+			"key": "zlib1g_1-1.2.13.dfsg-1_amd64",
+			"name": "zlib1g",
+			"sha256": "d7dd1d1411fedf27f5e27650a6eff20ef294077b568f4c8c5e51466dc7c08ce4",
+			"url": "https://deb.debian.org/debian/pool/main/z/zlib/zlib1g_1.2.13.dfsg-1_amd64.deb",
+			"version": "1:1.2.13.dfsg-1"
+		},
+		{
+			"arch": "amd64",
+			"dependencies": [
+				{
+					"key": "zlib1g_1-1.2.13.dfsg-1_amd64",
+					"name": "zlib1g",
+					"version": "1:1.2.13.dfsg-1"
+				},
+				{
+					"key": "libc6_2.36-9-p-deb12u9_amd64",
+					"name": "libc6",
+					"version": "2.36-9+deb12u9"
+				},
+				{
+					"key": "libgcc-s1_12.2.0-14_amd64",
+					"name": "libgcc-s1",
+					"version": "12.2.0-14"
+				},
+				{
+					"key": "gcc-12-base_12.2.0-14_amd64",
+					"name": "gcc-12-base",
+					"version": "12.2.0-14"
+				},
+				{
+					"key": "liblzma5_5.4.1-0.2_amd64",
+					"name": "liblzma5",
+					"version": "5.4.1-0.2"
+				},
+				{
+					"key": "libicu72_72.1-3_amd64",
+					"name": "libicu72",
+					"version": "72.1-3"
+				},
+				{
+					"key": "libstdc-p--p-6_12.2.0-14_amd64",
+					"name": "libstdc++6",
+					"version": "12.2.0-14"
+				}
+			],
+			"key": "libxml2_2.9.14-p-dfsg-1.3_deb12u1_amd64",
+			"name": "libxml2",
+			"sha256": "35b76cb7038fc1c940204a4f05f33ffb79d027353ce469397d9adcf8f9b3e1a7",
+			"url": "https://deb.debian.org/debian/pool/main/libx/libxml2/libxml2_2.9.14+dfsg-1.3~deb12u1_amd64.deb",
+			"version": "2.9.14+dfsg-1.3~deb12u1"
+		},
+		{
+			"arch": "amd64",
+			"dependencies": [],
+			"key": "liblzma5_5.4.1-0.2_amd64",
+			"name": "liblzma5",
+			"sha256": "d4b7736e58512a2b047f9cb91b71db5a3cf9d3451192fc6da044c77bf51fe869",
+			"url": "https://deb.debian.org/debian/pool/main/x/xz-utils/liblzma5_5.4.1-0.2_amd64.deb",
+			"version": "5.4.1-0.2"
+		},
+		{
+			"arch": "amd64",
+			"dependencies": [],
+			"key": "libicu72_72.1-3_amd64",
+			"name": "libicu72",
+			"sha256": "e239c1c9f52bee0ff627f291552d63691b765ec7c5cdf6de7c7ae4dec0275857",
+			"url": "https://deb.debian.org/debian/pool/main/i/icu/libicu72_72.1-3_amd64.deb",
+			"version": "72.1-3"
+		},
+		{
+			"arch": "amd64",
+			"dependencies": [],
+			"key": "libstdc-p--p-6_12.2.0-14_amd64",
+			"name": "libstdc++6",
+			"sha256": "9b1b269020cec6aced3b39f096f7b67edd1f0d4ab24f412cb6506d0800e19cbf",
+			"url": "https://deb.debian.org/debian/pool/main/g/gcc-12/libstdc++6_12.2.0-14_amd64.deb",
+			"version": "12.2.0-14"
+		},
+		{
+			"arch": "amd64",
+			"dependencies": [],
+			"key": "aws-neuronx-runtime-lib_2.22.19.0-5856c0b42_amd64",
+			"name": "aws-neuronx-runtime-lib",
+			"sha256": "f952dad4f615f2881a0bf93ba5132d308fdfa697a78a09e2e1653a6c4704222d",
+			"url": "https://apt.repos.neuron.amazonaws.com/pool/main/a/aws-neuronx-runtime-lib/aws-neuronx-runtime-lib_2.22.19.0-5856c0b42_amd64.deb",
+			"version": "2.22.19.0-5856c0b42"
+		},
+		{
+			"arch": "amd64",
+			"dependencies": [],
+			"key": "aws-neuronx-collectives_2.22.33.0-d2128d1aa_amd64",
+			"name": "aws-neuronx-collectives",
+			"sha256": "1eadead90a7daf6056a10a09d0ac3b72a3e885e6fedacd14b53a11a30d9774b3",
+			"url": "https://apt.repos.neuron.amazonaws.com/pool/main/a/aws-neuronx-collectives/aws-neuronx-collectives_2.22.33.0-d2128d1aa_amd64.deb",
+			"version": "2.22.33.0-d2128d1aa"
+		}
+	],
+	"version": 1
+}

runtimes/neuron/packages.yaml

@@ -0,0 +1,25 @@
+# Packages for examples/debian_snapshot.
+#
+#  Anytime this file is changed, the lockfile needs to be regenerated.
+#
+#  To generate the ubuntu_base.lock.json run the following command
+#
+#     bazel run @neuron_bookworm//:lock
+#
+version: 1
+
+sources:
+  - channel: bookworm main
+    url: https://deb.debian.org/debian
+  - channel: jammy main
+    url: https://apt.repos.neuron.amazonaws.com
+
+archs:
+  - "amd64"
+
+packages:
+  - "bash"
+  - "zlib1g"
+  - "libxml2"
+  - "aws-neuronx-runtime-lib"
+  - "aws-neuronx-collectives"

runtimes/neuron/pyenv.bzl

@@ -0,0 +1,41 @@
+load("@bazel_skylib//lib:paths.bzl", "paths")
+load("@rules_python//python:defs.bzl", "PyInfo")
+
+def runfile_path(ctx, runfile):
+    return paths.normalize(ctx.workspace_name + "/" + runfile.short_path)
+
+def _pyenv_zig_impl(ctx):
+    py_toolchain = ctx.toolchains["@rules_python//python:toolchain_type"].py3_runtime
+    content = "pub const home: [:0]const u8 = {};\n".format(repr(runfile_path(ctx, py_toolchain.interpreter)))
+
+    modules = []
+    for file in py_toolchain.files.to_list():
+        if file.basename == "__hello__.py":
+            modules.append(repr(runfile_path(ctx, file)))
+
+    imports = depset([], transitive = [
+        dep[PyInfo].imports
+        for dep in ctx.attr.deps
+    ])
+    modules.extend([
+        repr("{}/__init__.py".format(imp))
+        for imp in imports.to_list()
+    ])
+    content += "pub const modules: []const [:0]const u8 = &.{{ {} }};\n".format(", ".join(modules))
+
+    f = ctx.actions.declare_file("{}.zig".format(ctx.label.name))
+    ctx.actions.write(f, content)
+
+    return [
+        DefaultInfo(
+            files = depset([f]),
+        ),
+    ]
+
+pyenv_zig = rule(
+    implementation = _pyenv_zig_impl,
+    attrs = {
+        "deps": attr.label_list(providers = [PyInfo]),
+    },
+    toolchains = ["@rules_python//python:toolchain_type"],
+)

runtimes/neuron/requirements.in

@@ -0,0 +1,4 @@
+--find-links https://mirror.zml.ai/pypi/aws-neuronx-runtime-discovery/index.html
+--extra-index-url=https://pip.repos.neuron.amazonaws.com
+libneuronxla
+neuronx-cc

runtimes/neuron/requirements.lock.txt

@@ -0,0 +1,356 @@
+# This file was autogenerated by uv via the following command:
+#    bazel run @@//runtimes/neuron:update_requirements
+--index-url https://pypi.org/simple
+--extra-index-url https://pip.repos.neuron.amazonaws.com
+--find-links https://mirror.zml.ai/pypi/aws-neuronx-runtime-discovery/index.html
+
+aws-neuronx-runtime-discovery==2.9 \
+    --hash=sha256:87c291f833609bb5558708461d68426db82bf7eb80c5ca3426c3261b5622ffb6
+    # via libneuronxla
+boto3==1.35.69 \
+    --hash=sha256:20945912130cca1505f45819cd9b7183a0e376e91a1221a0b1f50c80d35fd7e2 \
+    --hash=sha256:40db86c7732a310b282f595251995ecafcbd62009a57e47a22683862e570cc7a
+    # via libneuronxla
+botocore==1.35.69 \
+    --hash=sha256:cad8d9305f873404eee4b197d84e60a40975d43cbe1ab63abe893420ddfe6e3c \
+    --hash=sha256:f9f23dd76fb247d9b0e8d411d2995e6f847fc451c026f1e58e300f815b0b36eb
+    # via
+    #   boto3
+    #   libneuronxla
+    #   s3transfer
+certifi==2024.8.30 \
+    --hash=sha256:922820b53db7a7257ffbda3f597266d435245903d80737e34f8a45ff3e3230d8 \
+    --hash=sha256:bec941d2aa8195e248a60b31ff9f0558284cf01a52591ceda73ea9afffd69fd9
+    # via requests
+charset-normalizer==3.4.0 \
+    --hash=sha256:0099d79bdfcf5c1f0c2c72f91516702ebf8b0b8ddd8905f97a8aecf49712c621 \
+    --hash=sha256:0713f3adb9d03d49d365b70b84775d0a0d18e4ab08d12bc46baa6132ba78aaf6 \
+    --hash=sha256:07afec21bbbbf8a5cc3651aa96b980afe2526e7f048fdfb7f1014d84acc8b6d8 \
+    --hash=sha256:0b309d1747110feb25d7ed6b01afdec269c647d382c857ef4663bbe6ad95a912 \
+    --hash=sha256:0d99dd8ff461990f12d6e42c7347fd9ab2532fb70e9621ba520f9e8637161d7c \
+    --hash=sha256:0de7b687289d3c1b3e8660d0741874abe7888100efe14bd0f9fd7141bcbda92b \
+    --hash=sha256:1110e22af8ca26b90bd6364fe4c763329b0ebf1ee213ba32b68c73de5752323d \
+    --hash=sha256:130272c698667a982a5d0e626851ceff662565379baf0ff2cc58067b81d4f11d \
+    --hash=sha256:136815f06a3ae311fae551c3df1f998a1ebd01ddd424aa5603a4336997629e95 \
+    --hash=sha256:14215b71a762336254351b00ec720a8e85cada43b987da5a042e4ce3e82bd68e \
+    --hash=sha256:1db4e7fefefd0f548d73e2e2e041f9df5c59e178b4c72fbac4cc6f535cfb1565 \
+    --hash=sha256:1ffd9493de4c922f2a38c2bf62b831dcec90ac673ed1ca182fe11b4d8e9f2a64 \
+    --hash=sha256:2006769bd1640bdf4d5641c69a3d63b71b81445473cac5ded39740a226fa88ab \
+    --hash=sha256:20587d20f557fe189b7947d8e7ec5afa110ccf72a3128d61a2a387c3313f46be \
+    --hash=sha256:223217c3d4f82c3ac5e29032b3f1c2eb0fb591b72161f86d93f5719079dae93e \
+    --hash=sha256:27623ba66c183eca01bf9ff833875b459cad267aeeb044477fedac35e19ba907 \
+    --hash=sha256:285e96d9d53422efc0d7a17c60e59f37fbf3dfa942073f666db4ac71e8d726d0 \
+    --hash=sha256:2de62e8801ddfff069cd5c504ce3bc9672b23266597d4e4f50eda28846c322f2 \
+    --hash=sha256:2f6c34da58ea9c1a9515621f4d9ac379871a8f21168ba1b5e09d74250de5ad62 \
+    --hash=sha256:309a7de0a0ff3040acaebb35ec45d18db4b28232f21998851cfa709eeff49d62 \
+    --hash=sha256:35c404d74c2926d0287fbd63ed5d27eb911eb9e4a3bb2c6d294f3cfd4a9e0c23 \
+    --hash=sha256:3710a9751938947e6327ea9f3ea6332a09bf0ba0c09cae9cb1f250bd1f1549bc \
+    --hash=sha256:3d59d125ffbd6d552765510e3f31ed75ebac2c7470c7274195b9161a32350284 \
+    --hash=sha256:40d3ff7fc90b98c637bda91c89d51264a3dcf210cade3a2c6f838c7268d7a4ca \
+    --hash=sha256:425c5f215d0eecee9a56cdb703203dda90423247421bf0d67125add85d0c4455 \
+    --hash=sha256:43193c5cda5d612f247172016c4bb71251c784d7a4d9314677186a838ad34858 \
+    --hash=sha256:44aeb140295a2f0659e113b31cfe92c9061622cadbc9e2a2f7b8ef6b1e29ef4b \
+    --hash=sha256:47334db71978b23ebcf3c0f9f5ee98b8d65992b65c9c4f2d34c2eaf5bcaf0594 \
+    --hash=sha256:4796efc4faf6b53a18e3d46343535caed491776a22af773f366534056c4e1fbc \
+    --hash=sha256:4a51b48f42d9358460b78725283f04bddaf44a9358197b889657deba38f329db \
+    --hash=sha256:4b67fdab07fdd3c10bb21edab3cbfe8cf5696f453afce75d815d9d7223fbe88b \
+    --hash=sha256:4ec9dd88a5b71abfc74e9df5ebe7921c35cbb3b641181a531ca65cdb5e8e4dea \
+    --hash=sha256:4f9fc98dad6c2eaa32fc3af1417d95b5e3d08aff968df0cd320066def971f9a6 \
+    --hash=sha256:54b6a92d009cbe2fb11054ba694bc9e284dad30a26757b1e372a1fdddaf21920 \
+    --hash=sha256:55f56e2ebd4e3bc50442fbc0888c9d8c94e4e06a933804e2af3e89e2f9c1c749 \
+    --hash=sha256:5726cf76c982532c1863fb64d8c6dd0e4c90b6ece9feb06c9f202417a31f7dd7 \
+    --hash=sha256:5d447056e2ca60382d460a604b6302d8db69476fd2015c81e7c35417cfabe4cd \
+    --hash=sha256:5ed2e36c3e9b4f21dd9422f6893dec0abf2cca553af509b10cd630f878d3eb99 \
+    --hash=sha256:5ff2ed8194587faf56555927b3aa10e6fb69d931e33953943bc4f837dfee2242 \
+    --hash=sha256:62f60aebecfc7f4b82e3f639a7d1433a20ec32824db2199a11ad4f5e146ef5ee \
+    --hash=sha256:63bc5c4ae26e4bc6be6469943b8253c0fd4e4186c43ad46e713ea61a0ba49129 \
+    --hash=sha256:6b40e8d38afe634559e398cc32b1472f376a4099c75fe6299ae607e404c033b2 \
+    --hash=sha256:6b493a043635eb376e50eedf7818f2f322eabbaa974e948bd8bdd29eb7ef2a51 \
+    --hash=sha256:6dba5d19c4dfab08e58d5b36304b3f92f3bd5d42c1a3fa37b5ba5cdf6dfcbcee \
+    --hash=sha256:6fd30dc99682dc2c603c2b315bded2799019cea829f8bf57dc6b61efde6611c8 \
+    --hash=sha256:707b82d19e65c9bd28b81dde95249b07bf9f5b90ebe1ef17d9b57473f8a64b7b \
+    --hash=sha256:7706f5850360ac01d80c89bcef1640683cc12ed87f42579dab6c5d3ed6888613 \
+    --hash=sha256:7782afc9b6b42200f7362858f9e73b1f8316afb276d316336c0ec3bd73312742 \
+    --hash=sha256:79983512b108e4a164b9c8d34de3992f76d48cadc9554c9e60b43f308988aabe \
+    --hash=sha256:7f683ddc7eedd742e2889d2bfb96d69573fde1d92fcb811979cdb7165bb9c7d3 \
+    --hash=sha256:82357d85de703176b5587dbe6ade8ff67f9f69a41c0733cf2425378b49954de5 \
+    --hash=sha256:84450ba661fb96e9fd67629b93d2941c871ca86fc38d835d19d4225ff946a631 \
+    --hash=sha256:86f4e8cca779080f66ff4f191a685ced73d2f72d50216f7112185dc02b90b9b7 \
+    --hash=sha256:8cda06946eac330cbe6598f77bb54e690b4ca93f593dee1568ad22b04f347c15 \
+    --hash=sha256:8ce7fd6767a1cc5a92a639b391891bf1c268b03ec7e021c7d6d902285259685c \
+    --hash=sha256:8ff4e7cdfdb1ab5698e675ca622e72d58a6fa2a8aa58195de0c0061288e6e3ea \
+    --hash=sha256:9289fd5dddcf57bab41d044f1756550f9e7cf0c8e373b8cdf0ce8773dc4bd417 \
+    --hash=sha256:92a7e36b000bf022ef3dbb9c46bfe2d52c047d5e3f3343f43204263c5addc250 \
+    --hash=sha256:92db3c28b5b2a273346bebb24857fda45601aef6ae1c011c0a997106581e8a88 \
+    --hash=sha256:95c3c157765b031331dd4db3c775e58deaee050a3042fcad72cbc4189d7c8dca \
+    --hash=sha256:980b4f289d1d90ca5efcf07958d3eb38ed9c0b7676bf2831a54d4f66f9c27dfa \
+    --hash=sha256:9ae4ef0b3f6b41bad6366fb0ea4fc1d7ed051528e113a60fa2a65a9abb5b1d99 \
+    --hash=sha256:9c98230f5042f4945f957d006edccc2af1e03ed5e37ce7c373f00a5a4daa6149 \
+    --hash=sha256:9fa2566ca27d67c86569e8c85297aaf413ffab85a8960500f12ea34ff98e4c41 \
+    --hash=sha256:a14969b8691f7998e74663b77b4c36c0337cb1df552da83d5c9004a93afdb574 \
+    --hash=sha256:a8aacce6e2e1edcb6ac625fb0f8c3a9570ccc7bfba1f63419b3769ccf6a00ed0 \
+    --hash=sha256:a8e538f46104c815be19c975572d74afb53f29650ea2025bbfaef359d2de2f7f \
+    --hash=sha256:aa41e526a5d4a9dfcfbab0716c7e8a1b215abd3f3df5a45cf18a12721d31cb5d \
+    --hash=sha256:aa693779a8b50cd97570e5a0f343538a8dbd3e496fa5dcb87e29406ad0299654 \
+    --hash=sha256:ab22fbd9765e6954bc0bcff24c25ff71dcbfdb185fcdaca49e81bac68fe724d3 \
+    --hash=sha256:ab2e5bef076f5a235c3774b4f4028a680432cded7cad37bba0fd90d64b187d19 \
+    --hash=sha256:ab973df98fc99ab39080bfb0eb3a925181454d7c3ac8a1e695fddfae696d9e90 \
+    --hash=sha256:af73657b7a68211996527dbfeffbb0864e043d270580c5aef06dc4b659a4b578 \
+    --hash=sha256:b197e7094f232959f8f20541ead1d9862ac5ebea1d58e9849c1bf979255dfac9 \
+    --hash=sha256:b295729485b06c1a0683af02a9e42d2caa9db04a373dc38a6a58cdd1e8abddf1 \
+    --hash=sha256:b8831399554b92b72af5932cdbbd4ddc55c55f631bb13ff8fe4e6536a06c5c51 \
+    --hash=sha256:b8dcd239c743aa2f9c22ce674a145e0a25cb1566c495928440a181ca1ccf6719 \
+    --hash=sha256:bcb4f8ea87d03bc51ad04add8ceaf9b0f085ac045ab4d74e73bbc2dc033f0236 \
+    --hash=sha256:bd7af3717683bea4c87acd8c0d3d5b44d56120b26fd3f8a692bdd2d5260c620a \
+    --hash=sha256:bf4475b82be41b07cc5e5ff94810e6a01f276e37c2d55571e3fe175e467a1a1c \
+    --hash=sha256:c3e446d253bd88f6377260d07c895816ebf33ffffd56c1c792b13bff9c3e1ade \
+    --hash=sha256:c57516e58fd17d03ebe67e181a4e4e2ccab1168f8c2976c6a334d4f819fe5944 \
+    --hash=sha256:c94057af19bc953643a33581844649a7fdab902624d2eb739738a30e2b3e60fc \
+    --hash=sha256:cab5d0b79d987c67f3b9e9c53f54a61360422a5a0bc075f43cab5621d530c3b6 \
+    --hash=sha256:ce031db0408e487fd2775d745ce30a7cd2923667cf3b69d48d219f1d8f5ddeb6 \
+    --hash=sha256:cee4373f4d3ad28f1ab6290684d8e2ebdb9e7a1b74fdc39e4c211995f77bec27 \
+    --hash=sha256:d5b054862739d276e09928de37c79ddeec42a6e1bfc55863be96a36ba22926f6 \
+    --hash=sha256:dbe03226baf438ac4fda9e2d0715022fd579cb641c4cf639fa40d53b2fe6f3e2 \
+    --hash=sha256:dc15e99b2d8a656f8e666854404f1ba54765871104e50c8e9813af8a7db07f12 \
+    --hash=sha256:dcaf7c1524c0542ee2fc82cc8ec337f7a9f7edee2532421ab200d2b920fc97cf \
+    --hash=sha256:dd4eda173a9fcccb5f2e2bd2a9f423d180194b1bf17cf59e3269899235b2a114 \
+    --hash=sha256:dd9a8bd8900e65504a305bf8ae6fa9fbc66de94178c420791d0293702fce2df7 \
+    --hash=sha256:de7376c29d95d6719048c194a9cf1a1b0393fbe8488a22008610b0361d834ecf \
+    --hash=sha256:e7fdd52961feb4c96507aa649550ec2a0d527c086d284749b2f582f2d40a2e0d \
+    --hash=sha256:e91f541a85298cf35433bf66f3fab2a4a2cff05c127eeca4af174f6d497f0d4b \
+    --hash=sha256:e9e3c4c9e1ed40ea53acf11e2a386383c3304212c965773704e4603d589343ed \
+    --hash=sha256:ee803480535c44e7f5ad00788526da7d85525cfefaf8acf8ab9a310000be4b03 \
+    --hash=sha256:f09cb5a7bbe1ecae6e87901a2eb23e0256bb524a79ccc53eb0b7629fbe7677c4 \
+    --hash=sha256:f19c1585933c82098c2a520f8ec1227f20e339e33aca8fa6f956f6691b784e67 \
+    --hash=sha256:f1a2f519ae173b5b6a2c9d5fa3116ce16e48b3462c8b96dfdded11055e3d6365 \
+    --hash=sha256:f28f891ccd15c514a0981f3b9db9aa23d62fe1a99997512b0491d2ed323d229a \
+    --hash=sha256:f3e73a4255342d4eb26ef6df01e3962e73aa29baa3124a8e824c5d3364a65748 \
+    --hash=sha256:f606a1881d2663630ea5b8ce2efe2111740df4b687bd78b34a8131baa007f79b \
+    --hash=sha256:fe9f97feb71aa9896b81973a7bbada8c49501dc73e58a10fcef6663af95e5079 \
+    --hash=sha256:ffc519621dce0c767e96b9c53f09c5d215578e10b02c285809f76509a3931482
+    # via requests
+ec2-metadata==2.14.0 \
+    --hash=sha256:b2f83381722072efa28f65dc37e726c0abd3a0532f227b3fa6a2ad68460c7feb \
+    --hash=sha256:fc4f9e893da1c049761221ec3d2208cce4872953b8960e1ed1f7120ff4ca3512
+    # via neuronx-cc
+idna==3.10 \
+    --hash=sha256:12f65c9b470abda6dc35cf8e63cc574b1c52b11df2c86030af0ac09b01b13ea9 \
+    --hash=sha256:946d195a0d259cbba61165e88e65941f16e9b36ea6ddb97f00452bae8b1287d3
+    # via requests
+islpy==2023.2.5 \
+    --hash=sha256:06446395028b88701fd84b685b1c82b9de493728906f96f576a4dca69d12dd1f \
+    --hash=sha256:0c1689d9ef2ccfb3fcb3108a37c8fa21561dac3615ec05508021628a3197d81c \
+    --hash=sha256:0f723e6d8fda7608e6329ee9a786b2812afefec01f6f45e4ac812594608d2990 \
+    --hash=sha256:20484f2ef8d8f5827e18eaac064bfbd4d0e9cb877030ee1297a544bc80e126db \
+    --hash=sha256:21571052bac1a319ef4937886559a41c85eddfc29389c7f9be6d8ddaa08dbd4a \
+    --hash=sha256:24b469e1e390860cb78a13e5a933727fed42fd991eb6dbb758754410224a28a2 \
+    --hash=sha256:2a64abecd9a966e6223633315e20b3abd9eb1faa8887122c675a35a7f4885212 \
+    --hash=sha256:31c48b641824bd6b50e97659a6842c445feeb359115148c3debcbb862cb85c1f \
+    --hash=sha256:3ee7815b1d49ac1efca0a42a5c038f88e87be01b9ddf598aadea0ef84b930a46 \
+    --hash=sha256:3f50ec187bbd61cef4400af6885e65afd4092e3c4dc84b5d5c98554de6d06988 \
+    --hash=sha256:4e97e23ba815afa37b8919ebb0d1c7ec7a9cb2c1491ea7ba89a886a8d213169c \
+    --hash=sha256:58b5b86a9f5f58d0e8165c7e2215e7441eeea3466277e44f24788aef8b2b8311 \
+    --hash=sha256:662a1c6f83feb77397eaef0027c5a57b98e249ecc454df6131c8cfe1541964dc \
+    --hash=sha256:679431e17ad40d3c496d5fd3e9cbb824047775c00ade143dd59f848110b287f7 \
+    --hash=sha256:69a13fa79fe9e901376e2c75bc49770e9027c55d9699077f8417e5bbe9e8b137 \
+    --hash=sha256:6a336af9106eb92c3b7babe919d9e1e0a629143941d32d474e47cc55cf64f647 \
+    --hash=sha256:6ed6c49803e53db5b0dbba9ec2fd3f896a8b9ea68f4a1cd203e5e7490f8e1c8e \
+    --hash=sha256:773e85ba8768e14bd2b813d4a7e8e8140fe628cddddbd4fcb3452b023ed57f84 \
+    --hash=sha256:7acbb16a6f55a6020cfcbd5823bbd3dd36ca03ef0713f9f7a47404b0aad7cba4 \
+    --hash=sha256:80674c7a9aef8d2657872264dc59aeb66ce92067603ead43be51a56a6efc058a \
+    --hash=sha256:819b16c4add84b42f227b451c2f773b8ae06489313d482771fc54c8806c1a2af \
+    --hash=sha256:831b7a8f290af99e24addcb5cc5f4e73f7d047d50202ad84c7d208cadbe24f15 \
+    --hash=sha256:a66c7baf86d6fcef8e2170483f45a241dca0811e0c78cc71efa30bc74bfc6bc8 \
+    --hash=sha256:b27f1aa0ec0ab8182e47aaab4e363a42221ae00f9cba1f7cdde4d687120965f3 \
+    --hash=sha256:b3582a7d19ee47f29c8c6b40643f69e5a52c0f0489ed179408cfd4d9616210c9 \
+    --hash=sha256:b45c69990d0f7a9fbbba4a28b552c3b93c80a75d17b52032851f2dde13ebd661 \
+    --hash=sha256:c0d67b1b8644ac8cf5dee3b2efbf8029891f5db1a625ee5259e2fcf6c33a2d55 \
+    --hash=sha256:c2fdf13ae6465c511d75b42ed18b2289b49ecc7f1764e7fa24340b05d0392e80 \
+    --hash=sha256:c6665d90beb3f37b2dcb95570bf8700e58884c5ad03dd131568287a639c03a5e \
+    --hash=sha256:cea357471c345e804fa31c9f1eacfb3f67eda26b3e8fd417c33556822076565b \
+    --hash=sha256:d35750ea2a3964325a2e67196806cf02ed771c2d2b2d47fb2bf60f58bf933cc4 \
+    --hash=sha256:d5bedd79939dac4403b573abe5e13f6fc6a4502ebe8ed7de79b2fbfdd7856f3f \
+    --hash=sha256:d75b4521f3291b92fc6796ebb7eabf5c6e318d44a578390f678d7bfaedde13b8 \
+    --hash=sha256:d77acfe6653d044d14b8e93482e0282f42b947d40f5a98a13ee6b705936aaa18 \
+    --hash=sha256:dd74398baf36938abb7c2a9d9828cc1a2e549c6c80485cec8b06b29abfaad18f \
+    --hash=sha256:e33812a9922d967da814b17bc778a2c0dd1dabcd7ff685d86677bd745e8a3d65 \
+    --hash=sha256:e372461fb5ba22f9955cdf180d6187ac691f0ffa972874af181bf60501c738cd \
+    --hash=sha256:ea705d8b5244409fdc32affc689b010e88efbc8d378072bcba8edcf51c797de7 \
+    --hash=sha256:ed0c62e64d160735ed1a7fc314d4c0f501711c282bb84a00ceaf590b4c4bd61f \
+    --hash=sha256:eeec97588bd5282c983bd73f9010156e5e251e9e5307af547e2cb53b02d4639a
+    # via neuronx-cc
+jmespath==1.0.1 \
+    --hash=sha256:02e2e4cc71b5bcab88332eebf907519190dd9e6e82107fa7f83b1003a6252980 \
+    --hash=sha256:90261b206d6defd58fdd5e85f478bf633a2901798906be2ad389150c5c60edbe
+    # via
+    #   boto3
+    #   botocore
+libneuronxla==2.0.5347.0 \
+    --hash=sha256:87331978c86faaa794cdfa951404a6f3adeac1a986dc5c6e0b742048ba2c67ed
+    # via -r runtimes/neuron/requirements.in
+lockfile==0.12.2 \
+    --hash=sha256:6aed02de03cba24efabcd600b30540140634fc06cfa603822d508d5361e9f799 \
+    --hash=sha256:6c3cb24f344923d30b2785d5ad75182c8ea7ac1b6171b08657258ec7429d50fa
+    # via python-daemon
+ml-dtypes==0.2.0 \
+    --hash=sha256:022d5a4ee6be14569c2a9d1549e16f1ec87ca949681d0dca59995445d5fcdd5b \
+    --hash=sha256:1749b60348da71fd3c2ab303fdbc1965958dc50775ead41f5669c932a341cafd \
+    --hash=sha256:32107e7fa9f62db9a5281de923861325211dfff87bd23faefb27b303314635ab \
+    --hash=sha256:35b984cddbe8173b545a0e3334fe56ea1a5c3eb67c507f60d0cfde1d3fa8f8c2 \
+    --hash=sha256:36d28b8861a8931695e5a31176cad5ae85f6504906650dea5598fbec06c94606 \
+    --hash=sha256:50845af3e9a601810751b55091dee6c2562403fa1cb4e0123675cf3a4fc2c17a \
+    --hash=sha256:6488eb642acaaf08d8020f6de0a38acee7ac324c1e6e92ee0c0fea42422cb797 \
+    --hash=sha256:75015818a7fccf99a5e8ed18720cb430f3e71a8838388840f4cdf225c036c983 \
+    --hash=sha256:80d304c836d73f10605c58ccf7789c171cc229bfb678748adfb7cea2510dfd0e \
+    --hash=sha256:832a019a1b6db5c4422032ca9940a990fa104eee420f643713241b3a518977fa \
+    --hash=sha256:8faaf0897942c8253dd126662776ba45f0a5861968cf0f06d6d465f8a7bc298a \
+    --hash=sha256:bc29a0524ef5e23a7fbb8d881bdecabeb3fc1d19d9db61785d077a86cb94fab2 \
+    --hash=sha256:df6a76e1c8adf484feb138ed323f9f40a7b6c21788f120f7c78bec20ac37ee81 \
+    --hash=sha256:e70047ec2c83eaee01afdfdabee2c5b0c133804d90d0f7db4dd903360fcc537c \
+    --hash=sha256:e85ba8e24cf48d456e564688e981cf379d4c8e644db0a2f719b78de281bac2ca \
+    --hash=sha256:f00c71c8c63e03aff313bc6a7aeaac9a4f1483a921a6ffefa6d4404efd1af3d0 \
+    --hash=sha256:f08c391c2794f2aad358e6f4c70785a9a7b1df980ef4c232b3ccd4f6fe39f719
+    # via neuronx-cc
+networkx==2.8.8 \
+    --hash=sha256:230d388117af870fce5647a3c52401fcf753e94720e6ea6b4197a5355648885e \
+    --hash=sha256:e435dfa75b1d7195c7b8378c3859f0445cd88c6b0375c181ed66823a9ceb7524
+    # via neuronx-cc
+neuronx-cc==2.15.143.0+e39249ad \
+    --hash=sha256:86d0d509ff1ec344b3ea30149e55a69ec720eec7e26a634c5e9d63350608b50e
+    # via
+    #   -r runtimes/neuron/requirements.in
+    #   libneuronxla
+numpy==1.25.2 \
+    --hash=sha256:0d60fbae8e0019865fc4784745814cff1c421df5afee233db6d88ab4f14655a2 \
+    --hash=sha256:1a1329e26f46230bf77b02cc19e900db9b52f398d6722ca853349a782d4cff55 \
+    --hash=sha256:1b9735c27cea5d995496f46a8b1cd7b408b3f34b6d50459d9ac8fe3a20cc17bf \
+    --hash=sha256:2792d23d62ec51e50ce4d4b7d73de8f67a2fd3ea710dcbc8563a51a03fb07b01 \
+    --hash=sha256:3e0746410e73384e70d286f93abf2520035250aad8c5714240b0492a7302fdca \
+    --hash=sha256:4c3abc71e8b6edba80a01a52e66d83c5d14433cbcd26a40c329ec7ed09f37901 \
+    --hash=sha256:5883c06bb92f2e6c8181df7b39971a5fb436288db58b5a1c3967702d4278691d \
+    --hash=sha256:5c97325a0ba6f9d041feb9390924614b60b99209a71a69c876f71052521d42a4 \
+    --hash=sha256:60e7f0f7f6d0eee8364b9a6304c2845b9c491ac706048c7e8cf47b83123b8dbf \
+    --hash=sha256:76b4115d42a7dfc5d485d358728cdd8719be33cc5ec6ec08632a5d6fca2ed380 \
+    --hash=sha256:7dc869c0c75988e1c693d0e2d5b26034644399dd929bc049db55395b1379e044 \
+    --hash=sha256:834b386f2b8210dca38c71a6e0f4fd6922f7d3fcff935dbe3a570945acb1b545 \
+    --hash=sha256:8b77775f4b7df768967a7c8b3567e309f617dd5e99aeb886fa14dc1a0791141f \
+    --hash=sha256:90319e4f002795ccfc9050110bbbaa16c944b1c37c0baeea43c5fb881693ae1f \
+    --hash=sha256:b79e513d7aac42ae918db3ad1341a015488530d0bb2a6abcbdd10a3a829ccfd3 \
+    --hash=sha256:bb33d5a1cf360304754913a350edda36d5b8c5331a8237268c48f91253c3a364 \
+    --hash=sha256:bec1e7213c7cb00d67093247f8c4db156fd03075f49876957dca4711306d39c9 \
+    --hash=sha256:c5462d19336db4560041517dbb7759c21d181a67cb01b36ca109b2ae37d32418 \
+    --hash=sha256:c5652ea24d33585ea39eb6a6a15dac87a1206a692719ff45d53c5282e66d4a8f \
+    --hash=sha256:d7806500e4f5bdd04095e849265e55de20d8cc4b661b038957354327f6d9b295 \
+    --hash=sha256:db3ccc4e37a6873045580d413fe79b68e47a681af8db2e046f1dacfa11f86eb3 \
+    --hash=sha256:dfe4a913e29b418d096e696ddd422d8a5d13ffba4ea91f9f60440a3b759b0187 \
+    --hash=sha256:eb942bfb6f84df5ce05dbf4b46673ffed0d3da59f13635ea9b926af3deb76926 \
+    --hash=sha256:f08f2e037bba04e707eebf4bc934f1972a315c883a9e0ebfa8a7756eabf9e357 \
+    --hash=sha256:fd608e19c8d7c55021dffd43bfe5492fab8cc105cc8986f813f8c3c048b38760
+    # via
+    #   ml-dtypes
+    #   neuronx-cc
+    #   scipy
+pgzip==0.3.5 \
+    --hash=sha256:4e13ab66ecface5c51c5af51d8cd676aa51675cf85df000f501a86cf38c208c1 \
+    --hash=sha256:dd35510f59f6bd6b64e31c4baf90c10cdbb2775235fcc079b14b404fbd7f46bf
+    # via neuronx-cc
+protobuf==5.28.3 \
+    --hash=sha256:0c4eec6f987338617072592b97943fdbe30d019c56126493111cf24344c1cc24 \
+    --hash=sha256:135658402f71bbd49500322c0f736145731b16fc79dc8f367ab544a17eab4535 \
+    --hash=sha256:27b246b3723692bf1068d5734ddaf2fccc2cdd6e0c9b47fe099244d80200593b \
+    --hash=sha256:3e6101d095dfd119513cde7259aa703d16c6bbdfae2554dfe5cfdbe94e32d548 \
+    --hash=sha256:3fa2de6b8b29d12c61911505d893afe7320ce7ccba4df913e2971461fa36d584 \
+    --hash=sha256:64badbc49180a5e401f373f9ce7ab1d18b63f7dd4a9cdc43c92b9f0b481cef7b \
+    --hash=sha256:70585a70fc2dd4818c51287ceef5bdba6387f88a578c86d47bb34669b5552c36 \
+    --hash=sha256:712319fbdddb46f21abb66cd33cb9e491a5763b2febd8f228251add221981135 \
+    --hash=sha256:91fba8f445723fcf400fdbe9ca796b19d3b1242cd873907979b9ed71e4afe868 \
+    --hash=sha256:a3f6857551e53ce35e60b403b8a27b0295f7d6eb63d10484f12bc6879c715687 \
+    --hash=sha256:cee1757663fa32a1ee673434fcf3bf24dd54763c79690201208bafec62f19eed
+    # via neuronx-cc
+psutil==6.1.0 \
+    --hash=sha256:000d1d1ebd634b4efb383f4034437384e44a6d455260aaee2eca1e9c1b55f047 \
+    --hash=sha256:045f00a43c737f960d273a83973b2511430d61f283a44c96bf13a6e829ba8fdc \
+    --hash=sha256:0895b8414afafc526712c498bd9de2b063deaac4021a3b3c34566283464aff8e \
+    --hash=sha256:1209036fbd0421afde505a4879dee3b2fd7b1e14fee81c0069807adcbbcca747 \
+    --hash=sha256:1ad45a1f5d0b608253b11508f80940985d1d0c8f6111b5cb637533a0e6ddc13e \
+    --hash=sha256:353815f59a7f64cdaca1c0307ee13558a0512f6db064e92fe833784f08539c7a \
+    --hash=sha256:498c6979f9c6637ebc3a73b3f87f9eb1ec24e1ce53a7c5173b8508981614a90b \
+    --hash=sha256:5cd2bcdc75b452ba2e10f0e8ecc0b57b827dd5d7aaffbc6821b2a9a242823a76 \
+    --hash=sha256:6d3fbbc8d23fcdcb500d2c9f94e07b1342df8ed71b948a2649b5cb060a7c94ca \
+    --hash=sha256:6e2dcd475ce8b80522e51d923d10c7871e45f20918e027ab682f94f1c6351688 \
+    --hash=sha256:9118f27452b70bb1d9ab3198c1f626c2499384935aaf55388211ad982611407e \
+    --hash=sha256:9dcbfce5d89f1d1f2546a2090f4fcf87c7f669d1d90aacb7d7582addece9fb38 \
+    --hash=sha256:a8506f6119cff7015678e2bce904a4da21025cc70ad283a53b099e7620061d85 \
+    --hash=sha256:a8fb3752b491d246034fa4d279ff076501588ce8cbcdbb62c32fd7a377d996be \
+    --hash=sha256:c0e0c00aa18ca2d3b2b991643b799a15fc8f0563d2ebb6040f64ce8dc027b942 \
+    --hash=sha256:d905186d647b16755a800e7263d43df08b790d709d575105d419f8b6ef65423a \
+    --hash=sha256:ff34df86226c0227c52f38b919213157588a678d049688eded74c76c8ba4a5d0
+    # via neuronx-cc
+python-daemon==3.1.0 \
+    --hash=sha256:a66b5896f0aed5807a25c6128268eb496488b1f9c6927c487710049ba16be32a \
+    --hash=sha256:fdb621d7e5f46e74b4de1ad6b0fff6e69cd91b4f219de1476190ebdd0f4781df
+    # via neuronx-cc
+python-dateutil==2.9.0.post0 \
+    --hash=sha256:37dd54208da7e1cd875388217d5e00ebd4179249f90fb72437e91a35459a0ad3 \
+    --hash=sha256:a8b2bc7bffae282281c8140a97d3aa9c14da0b136dfe83f850eea9a5f7470427
+    # via botocore
+requests==2.31.0 \
+    --hash=sha256:58cd2187c01e70e6e26505bca751777aa9f2ee0b7f4300988b709f44e013003f \
+    --hash=sha256:942c5a758f98d790eaed1a29cb6eefc7ffb0d1cf7af05c3d2791656dbd6ad1e1
+    # via
+    #   ec2-metadata
+    #   neuronx-cc
+    #   requests-unixsocket
+requests-unixsocket==0.3.0 \
+    --hash=sha256:28304283ea9357d45fff58ad5b11e47708cfbf5806817aa59b2a363228ee971e \
+    --hash=sha256:c685c680f0809e1b2955339b1e5afc3c0022b3066f4f7eb343f43a6065fc0e5d
+    # via neuronx-cc
+s3transfer==0.10.4 \
+    --hash=sha256:244a76a24355363a68164241438de1b72f8781664920260c48465896b712a41e \
+    --hash=sha256:29edc09801743c21eb5ecbc617a152df41d3c287f67b615f73e5f750583666a7
+    # via boto3
+scipy==1.11.2 \
+    --hash=sha256:0f3261f14b767b316d7137c66cc4f33a80ea05841b9c87ad83a726205b901423 \
+    --hash=sha256:10eb6af2f751aa3424762948e5352f707b0dece77288206f227864ddf675aca0 \
+    --hash=sha256:1342ca385c673208f32472830c10110a9dcd053cf0c4b7d4cd7026d0335a6c1d \
+    --hash=sha256:214cdf04bbae7a54784f8431f976704ed607c4bc69ba0d5d5d6a9df84374df76 \
+    --hash=sha256:2b997a5369e2d30c97995dcb29d638701f8000d04df01b8e947f206e5d0ac788 \
+    --hash=sha256:2c91cf049ffb5575917f2a01da1da082fd24ed48120d08a6e7297dfcac771dcd \
+    --hash=sha256:3aeb87661de987f8ec56fa6950863994cd427209158255a389fc5aea51fa7055 \
+    --hash=sha256:4447ad057d7597476f9862ecbd9285bbf13ba9d73ce25acfa4e4b11c6801b4c9 \
+    --hash=sha256:542a757e2a6ec409e71df3d8fd20127afbbacb1c07990cb23c5870c13953d899 \
+    --hash=sha256:8d9886f44ef8c9e776cb7527fb01455bf4f4a46c455c4682edc2c2cc8cd78562 \
+    --hash=sha256:90d3b1364e751d8214e325c371f0ee0dd38419268bf4888b2ae1040a6b266b2a \
+    --hash=sha256:95763fbda1206bec41157582bea482f50eb3702c85fffcf6d24394b071c0e87a \
+    --hash=sha256:ac74b1512d38718fb6a491c439aa7b3605b96b1ed3be6599c17d49d6c60fca18 \
+    --hash=sha256:afdb0d983f6135d50770dd979df50bf1c7f58b5b33e0eb8cf5c73c70600eae1d \
+    --hash=sha256:b0620240ef445b5ddde52460e6bc3483b7c9c750275369379e5f609a1050911c \
+    --hash=sha256:b133f237bd8ba73bad51bc12eb4f2d84cbec999753bf25ba58235e9fc2096d80 \
+    --hash=sha256:b29318a5e39bd200ca4381d80b065cdf3076c7d7281c5e36569e99273867f61d \
+    --hash=sha256:b8425fa963a32936c9773ee3ce44a765d8ff67eed5f4ac81dc1e4a819a238ee9 \
+    --hash=sha256:d2b813bfbe8dec6a75164523de650bad41f4405d35b0fa24c2c28ae07fcefb20 \
+    --hash=sha256:d690e1ca993c8f7ede6d22e5637541217fc6a4d3f78b3672a6fe454dbb7eb9a7 \
+    --hash=sha256:e367904a0fec76433bf3fbf3e85bf60dae8e9e585ffd21898ab1085a29a04d16 \
+    --hash=sha256:ea932570b1c2a30edafca922345854ff2cd20d43cd9123b6dacfdecebfc1a80b \
+    --hash=sha256:f28f1f6cfeb48339c192efc6275749b2a25a7e49c4d8369a28b6591da02fbc9a \
+    --hash=sha256:f73102f769ee06041a3aa26b5841359b1a93cc364ce45609657751795e8f4a4a \
+    --hash=sha256:fa4909c6c20c3d91480533cddbc0e7c6d849e7d9ded692918c76ce5964997898
+    # via neuronx-cc
+setuptools==75.6.0 \
+    --hash=sha256:8199222558df7c86216af4f84c30e9b34a61d8ba19366cc914424cdbd28252f6 \
+    --hash=sha256:ce74b49e8f7110f9bf04883b730f4765b774ef3ef28f722cce7c273d253aaf7d
+    # via python-daemon
+six==1.16.0 \
+    --hash=sha256:1e61c37477a1626458e36f7b1d82aa5c9b094fa4802892072e49de9c60c4c926 \
+    --hash=sha256:8abb2f1d86890a2dfb989f9a77cfcfd3e47c2a354b01111771326f8aa26e0254
+    # via python-dateutil
+tqdm==4.67.1 \
+    --hash=sha256:26445eca388f82e72884e0d580d5464cd801a3ea01e63e5601bdff9ba6a48de2 \
+    --hash=sha256:f8aef9c52c08c13a65f30ea34f4e5aac3fd1a34959879d7e59e63027286627f2
+    # via neuronx-cc
+urllib3==2.2.3 \
+    --hash=sha256:ca899ca043dcb1bafa3e262d73aa25c465bfb49e0bd9dd5d59f1d0acba2f8fac \
+    --hash=sha256:e7d814a81dad81e6caf2ec9fdedb284ecc9c73076b62654547cc64ccdcae26e9
+    # via
+    #   botocore
+    #   requests

runtimes/neuron/zmlxneuron.c

@@ -0,0 +1,21 @@
+#include <dlfcn.h>
+#include <string.h>
+
+void *zmlxneuron_dlopen(const char *filename, int flags)
+{
+    if (filename != NULL)
+    {
+        char *replacements[] = {
+            "libnccom.so", "libnccom.so.2", "libnccom-net.so", "libnccom-net.so.0", NULL, NULL,
+        };
+        for (int i = 0; replacements[i] != NULL; i += 2)
+        {
+            if (strcmp(filename, replacements[i]) == 0)
+            {
+                filename = replacements[i + 1];
+                break;
+            }
+        }
+    }
+    return dlopen(filename, flags);
+}

runtimes/rocm/BUILD.bazel

@@ -41,3 +41,9 @@ zig_library(
         "//conditions:default": [":empty"],
     }),
 )
+
+filegroup(
+    name = "layers",
+    srcs = [],
+    visibility = ["//visibility:public"],
+)

runtimes/runtimes.zig

@@ -3,12 +3,14 @@ const cpu = @import("runtimes/cpu");
 const cuda = @import("runtimes/cuda");
 const rocm = @import("runtimes/rocm");
 const tpu = @import("runtimes/tpu");
+const neuron = @import("runtimes/neuron");
 
 pub const Platform = enum {
     cpu,
     cuda,
     rocm,
     tpu,
+    neuron,
 };
 
 pub fn load(tag: Platform) !*const pjrt.Api {
@@ -17,6 +19,7 @@ pub fn load(tag: Platform) !*const pjrt.Api {
         .cuda => try cuda.load(),
         .rocm => try rocm.load(),
         .tpu => try tpu.load(),
+        .neuron => try neuron.load(),
     };
 }
 
@@ -26,5 +29,6 @@ pub fn isEnabled(tag: Platform) bool {
         .cuda => cuda.isEnabled(),
         .rocm => rocm.isEnabled(),
         .tpu => tpu.isEnabled(),
+        .neuron => neuron.isEnabled(),
     };
 }

runtimes/tpu/BUILD.bazel

@@ -25,3 +25,9 @@ zig_library(
         "//conditions:default": [":empty"],
     }),
 )
+
+filegroup(
+    name = "layers",
+    srcs = [],
+    visibility = ["//visibility:public"],
+)

zml/module.zig

@@ -350,6 +350,9 @@ pub const CompilationContext = struct {
     /// Given a list of donations mapping output buffers to input buffers,
     /// generate donation attribute for each `n_args` input argument.
     fn addDonationsAttributes(self: CompilationContext, attributes: []AttributeList, donations: []const Tensor._Donation) void {
+        if (self.target() == .neuron) {
+            return;
+        }
         var n_donations: usize = 0;
         for (donations, 0..) |donation, index| {
             switch (donation) {
@@ -1210,13 +1213,31 @@ fn loadOrCompilePjrtExecutable(
 
 fn compileModuleToPjrtExecutable(arena: std.mem.Allocator, platform: Platform, module: mlir.Module, module_hash: u64) !*pjrt.LoadedExecutable {
     const sharding = platform.sharding();
+
+    // NOTE(Corendos): Hack needed because Protobuf struct are not public.
+    const DeviceAssignmentProto = @TypeOf(xla_pb.CompileOptionsProto.init().executable_build_options.?.device_assignment.?);
     var options: xla_pb.CompileOptionsProto = .{
         .executable_build_options = .{
+            .device_ordinal = -1,
             .num_replicas = sharding.num_replicas,
             .num_partitions = sharding.num_partitions,
             .use_spmd_partitioning = sharding.num_partitions > 1 or sharding.num_replicas > 1,
+            .device_assignment = .{
+                .replica_count = sharding.num_replicas,
+                .computation_count = sharding.num_partitions,
+                .computation_devices = blk: {
+                    var computation_devices = try std.ArrayListUnmanaged(DeviceAssignmentProto.ComputationDevice).initCapacity(arena, sharding.num_partitions);
+                    for (0..sharding.num_partitions) |i| {
+                        var replica_device_ids = std.ArrayListUnmanaged(i64).initCapacity(arena, 1) catch unreachable;
+                        replica_device_ids.appendAssumeCapacity(@intCast(i));
+                        computation_devices.appendAssumeCapacity(.{ .replica_device_ids = replica_device_ids });
+                    }
+                    break :blk computation_devices;
+                },
+            },
         },
     };
+
     // Let the arena deinit, zig-protobuf deinit is very slow.
     if (platform.compilation_options.xla_dump_to) |xla_dump_to| {
         try options.env_option_overrides.append(arena, .{