Merge pull request #23924 from jburnim:jburnim_doc_typo_fix

PiperOrigin-RevId: 678896733
jax-ml · Sep 25, 2024 · 911acf1 · 911acf1
2 parents e4ca4f5 + a1f2edc
commit 911acf1
Showing 1 changed file with 2 additions and 2 deletions.
diff --git 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