From 8b53ebbd9b9760a49e4fda2fc2962cb3d1688f17 Mon Sep 17 00:00:00 2001
From: Tarry Singh <tarry.singh@deepkapha.com>
Date: Thu, 26 Sep 2024 12:02:04 +0000
Subject: [PATCH] pjrt: bind PJRT_Device_AddressableMemories.

---
 pjrt/pjrt.zig | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/pjrt/pjrt.zig b/pjrt/pjrt.zig
index 9dbd5aa..2b0e65b 100644
--- a/pjrt/pjrt.zig
+++ b/pjrt/pjrt.zig
@@ -502,6 +502,13 @@ pub const Device = opaque {
         }) catch unreachable;
         return @intCast(ret.local_hardware_id);
     }
+
+    pub fn addressableMemories(self: *const Device, api: *const Api) ApiError![]const *Memory {
+        const ret = try api.call(.PJRT_Device_AddressableMemories, .{
+            .device = self.inner(),
+        });
+        return @ptrCast(ret.memories[0..ret.num_memories]);
+    }
 };
 
 pub const DeviceDescription = opaque {