diff --git a/MODULE.bazel b/MODULE.bazel
index 8582201..1c854e0 100644
--- a/MODULE.bazel
+++ b/MODULE.bazel
@@ -6,7 +6,7 @@ bazel_dep(name = "abseil-cpp", version = "20240722.0.bcr.2")
 bazel_dep(name = "aspect_bazel_lib", version = "2.14.0")
 bazel_dep(name = "aspect_rules_py", version = "1.3.2")
 bazel_dep(name = "bazel_skylib", version = "1.7.1")
-bazel_dep(name = "libxev", version = "20250313.0-5773f46")
+bazel_dep(name = "libxev", version = "20250718.0-9f785d2")
 bazel_dep(name = "patchelf", version = "0.18.0")
 bazel_dep(name = "pcre2", version = "10.43")
 bazel_dep(name = "platforms", version = "0.0.11")
@@ -23,7 +23,7 @@ bazel_dep(name = "toolchains_llvm_bootstrapped", version = "0.2.4")
 bazel_dep(name = "toolchains_protoc", version = "0.4.1")
 bazel_dep(name = "with_cfg.bzl", version = "0.9.1")
 bazel_dep(name = "xla", version = "20250710.0-22ea002")
-bazel_dep(name = "zig-protobuf", version = "20250318.0-930153e")
+bazel_dep(name = "zig-protobuf", version = "20250716.0-97f1e31")
 bazel_dep(name = "zig-yaml", version = "20240903.0-83d5fdf")
 
 bazel_dep(name = "buildifier_prebuilt", version = "8.0.3", dev_dependency = True)
diff --git a/third_party/modules/libxev/20250718.0-9f785d2/MODULE.bazel b/third_party/modules/libxev/20250718.0-9f785d2/MODULE.bazel
new file mode 100644
index 0000000..82907dd
--- /dev/null
+++ b/third_party/modules/libxev/20250718.0-9f785d2/MODULE.bazel
@@ -0,0 +1,7 @@
+module(
+    name = "libxev",
+    version = "20250718.0-9f785d2",
+    compatibility_level = 1,
+)
+
+bazel_dep(name = "rules_zig", version = "20250613.0-567662a")
diff --git a/third_party/modules/libxev/20250718.0-9f785d2/overlay/BUILD.bazel b/third_party/modules/libxev/20250718.0-9f785d2/overlay/BUILD.bazel
new file mode 100644
index 0000000..872ef7f
--- /dev/null
+++ b/third_party/modules/libxev/20250718.0-9f785d2/overlay/BUILD.bazel
@@ -0,0 +1,13 @@
+load("@rules_zig//zig:defs.bzl", "zig_library")
+
+zig_library(
+    name = "xev",
+    srcs = glob([
+        "src/*.zig",
+        "src/backend/*.zig",
+        "src/linux/*.zig",
+        "src/watcher/*.zig",
+    ]),
+    main = "src/main.zig",
+    visibility = ["//visibility:public"],
+)
diff --git a/third_party/modules/libxev/20250718.0-9f785d2/overlay/MODULE.bazel b/third_party/modules/libxev/20250718.0-9f785d2/overlay/MODULE.bazel
new file mode 100644
index 0000000..82907dd
--- /dev/null
+++ b/third_party/modules/libxev/20250718.0-9f785d2/overlay/MODULE.bazel
@@ -0,0 +1,7 @@
+module(
+    name = "libxev",
+    version = "20250718.0-9f785d2",
+    compatibility_level = 1,
+)
+
+bazel_dep(name = "rules_zig", version = "20250613.0-567662a")
diff --git a/third_party/modules/libxev/20250718.0-9f785d2/source.json b/third_party/modules/libxev/20250718.0-9f785d2/source.json
new file mode 100644
index 0000000..daf7444
--- /dev/null
+++ b/third_party/modules/libxev/20250718.0-9f785d2/source.json
@@ -0,0 +1,11 @@
+{
+    "strip_prefix": "libxev-9f785d202ddccf7a625e799250579253977978b6",
+    "url": "https://github.com/mitchellh/libxev/archive/9f785d202ddccf7a625e799250579253977978b6.tar.gz",
+    "integrity": "sha256-HaExxsAsATjH/pCP9NMo/EjVrG2KVX3SiTtZzmS8R90=",
+    "overlay": {
+        "MODULE.bazel": "",
+        "BUILD.bazel": ""
+    },
+    "patch_strip": 1
+}
+
diff --git a/third_party/modules/libxev/metadata.json b/third_party/modules/libxev/metadata.json
index 59deb2c..6d9687b 100644
--- a/third_party/modules/libxev/metadata.json
+++ b/third_party/modules/libxev/metadata.json
@@ -20,6 +20,7 @@
         "20250124.0-31eed4e",
         "20250222.0-07bcffa",
         "20250313.0-5773f46",
+        "20250716.0-9f785d2",
     ],
     "yanked_versions": {}
 }
diff --git a/third_party/modules/zig-protobuf/20250716.0-97f1e31/MODULE.bazel b/third_party/modules/zig-protobuf/20250716.0-97f1e31/MODULE.bazel
new file mode 100644
index 0000000..b71798c
--- /dev/null
+++ b/third_party/modules/zig-protobuf/20250716.0-97f1e31/MODULE.bazel
@@ -0,0 +1,8 @@
+module(
+    name = "zig-protobuf",
+    version = "20250716.0-97f1e31",
+    compatibility_level = 1,
+)
+
+bazel_dep(name = "rules_zig", version = "20250613.0-567662a")
+bazel_dep(name = "rules_proto", version = "6.0.0-rc1")
diff --git a/third_party/modules/zig-protobuf/20250716.0-97f1e31/overlay/BUILD.bazel b/third_party/modules/zig-protobuf/20250716.0-97f1e31/overlay/BUILD.bazel
new file mode 100644
index 0000000..832c759
--- /dev/null
+++ b/third_party/modules/zig-protobuf/20250716.0-97f1e31/overlay/BUILD.bazel
@@ -0,0 +1,32 @@
+load("@rules_proto//proto:defs.bzl", "proto_lang_toolchain")
+load("@rules_zig//zig:defs.bzl", "BINARY_KIND", "zig_binary", "zig_library")
+
+zig_library(
+    name = "protobuf",
+    import_name = "protobuf",
+    main = "src/protobuf.zig",
+    visibility = ["//visibility:public"],
+)
+
+zig_binary(
+    name = "generator",
+    srcs = [
+        "bootstrapped-generator/FullName.zig",
+        "bootstrapped-generator/google/protobuf/compiler/plugin.pb.zig",
+        "bootstrapped-generator/google/protobuf/descriptor.pb.zig",
+    ],
+    kind = BINARY_KIND.exe,
+    main = "bootstrapped-generator/main.zig",
+    visibility = ["//visibility:public"],
+    deps = [":protobuf"],
+)
+
+proto_lang_toolchain(
+    name = "zig_toolchain",
+    command_line = "--zig_out=$(OUT)",
+    output_files = "multiple",
+    plugin = ":generator",
+    plugin_format_flag = "--plugin=protoc-gen-zig=%s",
+    runtime = ":protobuf",
+    visibility = ["//visibility:public"],
+)
diff --git a/third_party/modules/zig-protobuf/20250716.0-97f1e31/overlay/MODULE.bazel b/third_party/modules/zig-protobuf/20250716.0-97f1e31/overlay/MODULE.bazel
new file mode 100644
index 0000000..b71798c
--- /dev/null
+++ b/third_party/modules/zig-protobuf/20250716.0-97f1e31/overlay/MODULE.bazel
@@ -0,0 +1,8 @@
+module(
+    name = "zig-protobuf",
+    version = "20250716.0-97f1e31",
+    compatibility_level = 1,
+)
+
+bazel_dep(name = "rules_zig", version = "20250613.0-567662a")
+bazel_dep(name = "rules_proto", version = "6.0.0-rc1")
diff --git a/third_party/modules/zig-protobuf/20250716.0-97f1e31/source.json b/third_party/modules/zig-protobuf/20250716.0-97f1e31/source.json
new file mode 100644
index 0000000..4975cdb
--- /dev/null
+++ b/third_party/modules/zig-protobuf/20250716.0-97f1e31/source.json
@@ -0,0 +1,9 @@
+{
+    "strip_prefix": "zig-protobuf-97f1e314d9bbf600fbaadee5bb08c8856e0006de",
+    "url": "https://github.com/zml/zig-protobuf/archive/97f1e314d9bbf600fbaadee5bb08c8856e0006de.tar.gz",
+    "integrity": "sha256-iAqsLIKhl15JInJvwMY8NGVll6eUBXzdLW6Xv8wANGY=",
+    "overlay": {
+        "MODULE.bazel": "",
+        "BUILD.bazel": ""
+    }
+}
diff --git a/third_party/modules/zig-protobuf/metadata.json b/third_party/modules/zig-protobuf/metadata.json
index 87b4f2b..4a03836 100644
--- a/third_party/modules/zig-protobuf/metadata.json
+++ b/third_party/modules/zig-protobuf/metadata.json
@@ -8,12 +8,13 @@
         }
     ],
     "repository": [
-        "gwenzek/zig-protobuf"
+        "zml/zig-protobuf"
     ],
     "versions": [
         "20240722.0-c644d11",
         "20250213.0-5304067",
-        "20250318.0-930153e"
+        "20250318.0-930153e",
+        "20250716.0-97f1e31",
     ],
     "yanked_versions": {}
 }
diff --git a/zml/module.zig b/zml/module.zig
index 8ab1aed..f031e81 100644
--- a/zml/module.zig
+++ b/zml/module.zig
@@ -841,8 +841,6 @@ fn compileModuleToPjrtExecutable(arena: std.mem.Allocator, platform: Platform, m
 
     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,
@@ -852,18 +850,18 @@ fn compileModuleToPjrtExecutable(arena: std.mem.Allocator, platform: Platform, m
             .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;
-                },
+                // Filled below.
+                .computation_devices = .{},
             },
         },
     };
+    const computation_devices = &options.executable_build_options.?.device_assignment.?.computation_devices;
+    try computation_devices.ensureTotalCapacity(arena, sharding.num_partitions);
+    const replica_device_ids = try arena.alloc(i64, sharding.num_partitions);
+    for (0..sharding.num_partitions) |i| {
+        replica_device_ids[i] = @intCast(i);
+        computation_devices.appendAssumeCapacity(.{ .replica_device_ids = .fromOwnedSlice(replica_device_ids[i .. i + 1]) });
+    }
 
     // Let the arena deinit, zig-protobuf deinit is very slow.
     try options.env_option_overrides.ensureUnusedCapacity(arena, 16);