Movatterモバイル変換

[0]ホーム
URL:

Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly

 Sign up
Appearance settings
/jaxPublic

Commitbe2883a

Browse files
[Mosaic GPU] Add a small section about.at to the Pallas quickstart doc.
PiperOrigin-RevId: 837159903
1 parenta5c92a2 commitbe2883a

File tree

2 files changed

+62
-2
lines changed

2 files changed

+62
-2
lines changed

‎docs/pallas/quickstart.ipynb‎

Lines changed: 35 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -89,7 +89,41 @@
8989
"We then write `x + y` to `o_ref`.\n",
9090
"Mutation has not historically been supported in JAX -- `jax.Array`s are immutable!\n",
9191
"`Ref`s are new (experimental) types that allow mutation under certain circumstances.\n",
92-
"We can interpret writing to a `Ref` as mutating its underlying buffer."
92+
"We can interpret writing to a `Ref` as mutating its underlying buffer.\n",
93+
"\n",
94+
"**Indexing and Slicing `Ref`s with `.at`**\n",
95+
"\n",
96+
"In addition to accessing the entire underlying buffer through a reference, it\n",
97+
"is possible to also access only a slice by using the `.at` property. Using\n",
98+
"`x_ref.at[slice]` does not immediately read or write data; it\n",
99+
"creates a new `Ref` object that points to a slice of the original buffer. For\n",
100+
"example `ref.at[0:128]` creates a view of the first 128 elements; `ref.at[::2]`\n",
101+
"creates a strided view.\n",
102+
"\n",
103+
"Once you have a new `Ref` that represents a slice you can read it or write to it\n",
104+
"with the usual syntax. Here is simple example:"
105+
]
106+
},
107+
{
108+
"cell_type":"code",
109+
"execution_count":null,
110+
"id":"b2563d6a",
111+
"metadata": {},
112+
"outputs": [],
113+
"source": [
114+
"def add_sliced_kernel(x_ref, y_ref, o_ref):\n",
115+
" mid = x_ref.shape[0] // 2\n",
116+
"\n",
117+
" x_left = x_ref[:mid][...]\n",
118+
" x_right = x_ref[mid:][...]\n",
119+
" y_left = y_ref[:mid][...]\n",
120+
" y_right = y_ref[mid:][...]\n",
121+
"\n",
122+
" # The output shape is (4, mid).\n",
123+
" o_ref.at[0][...] = x_left + y_left\n",
124+
" o_ref.at[1][...] = x_left + y_right\n",
125+
" o_ref.at[2][...] = x_right + y_left\n",
126+
" o_ref.at[3][...] = x_right + y_right"
93127
]
94128
},
95129
{

‎docs/pallas/quickstart.md‎

Lines changed: 27 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -72,7 +72,33 @@ Mutation has not historically been supported in JAX -- `jax.Array`s are immutabl
7272
`Ref`s are new (experimental) types that allow mutation under certain circumstances.
7373
We can interpret writing to a`Ref` as mutating its underlying buffer.
7474

75-
+++
75+
**Indexing and Slicing`Ref`s with`.at`**
76+
77+
In addition to accessing the entire underlying buffer through a reference, it
78+
is possible to also access only a slice by using the`.at` property. Using
79+
`x_ref.at[slice]` does not immediately read or write data; it
80+
creates a new`Ref` object that points to a slice of the original buffer. For
81+
example`ref.at[0:128]` creates a view of the first 128 elements;`ref.at[::2]`
82+
creates a strided view.
83+
84+
Once you have a new`Ref` that represents a slice you can read it or write to it
85+
with the usual syntax. Here is simple example:
86+
87+
```{code-cell} ipython3
88+
def add_sliced_kernel(x_ref, y_ref, o_ref):
89+
mid = x_ref.shape[0] // 2
90+
91+
x_left = x_ref[:mid][...]
92+
x_right = x_ref[mid:][...]
93+
y_left = y_ref[:mid][...]
94+
y_right = y_ref[mid:][...]
95+
96+
# The output shape is (4, mid).
97+
o_ref.at[0][...] = x_left + y_left
98+
o_ref.at[1][...] = x_left + y_right
99+
o_ref.at[2][...] = x_right + y_left
100+
o_ref.at[3][...] = x_right + y_right
101+
```
76102

77103
So we've written what we call a "kernel", which we define as a program that will
78104
run as an atomic unit of execution on an accelerator,

0 commit comments

Comments
 (0)
[8]ページ先頭
©2009-2025 Movatter.jp