From a1f2edc968507836638bdb224af7e310e7099cd9 Mon Sep 17 00:00:00 2001 From: Jacob Burnim Date: Wed, 25 Sep 2024 13:39:39 -0700 Subject: [PATCH] Fix make_remote_async_copy -> make_async_remote_copy in async doc. --- docs/pallas/async_note.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/pallas/async_note.md b/docs/pallas/async_note.md index 96370ee48625..42e32a074fd7 100644 --- a/docs/pallas/async_note.md +++ b/docs/pallas/async_note.md @@ -37,7 +37,7 @@ Now imagine we aren’t using XLA’s `ppermute` but have our own custom Pallas ```py def ppermute_kernel(x_ref, y_ref, send_sem, recv_sem): right_neighbor = ... - descriptor = pltpu.make_remote_async_copy(x_ref, y_ref, send_sem, recv_sem, device_id=right_neighbor) + descriptor = pltpu.make_async_remote_copy(x_ref, y_ref, send_sem, recv_sem, device_id=right_neighbor) descriptor.start() descriptor.wait_send() descriptor.wait_recv() @@ -54,7 +54,7 @@ def add_one(x_ref, z_ref): def ppermute_add_one_kernel(x_ref, y_ref, z_ref, send_sem, recv_sem): right_neighbor = ... - descriptor = pltpu.make_remote_async_copy(x_ref, y_ref, send_sem, recv_sem, device_id=right_neighbor) + descriptor = pltpu.make_async_remote_copy(x_ref, y_ref, send_sem, recv_sem, device_id=right_neighbor) descriptor.start() # Explicitly schedule inner kernel between start/wait