mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-16 03:46:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fix typo in pmap docstring

Browse Source 
Docstring states:
>  If the pmapped function is called with fewer positional arguments than indicated by **`static_argnums`** then an error is raised.

However `static_argnums` is not an argument that exists - I believe this should be corrected to `static_broadcasted_argnums`.

PiperOrigin-RevId: 595731210
This commit is contained in:
Tom Cobley 2024-01-04 09:49:14 -08:00 committed by jax authors
parent 326d1d27ef
commit ebc7af95df
1 changed files with 4 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
jax/_src/api.py
View File

@ -1415,10 +1415,10 @@ def pmap(
Operations that only depend on static arguments will be constant-folded. Operations that only depend on static arguments will be constant-folded.
Calling the pmapped function with different values for these constants Calling the pmapped function with different values for these constants
will trigger recompilation. If the pmapped function is called with fewer will trigger recompilation. If the pmapped function is called with fewer
positional arguments than indicated by ``static_argnums`` then an error is positional arguments than indicated by ``static_broadcasted_argnums`` then
raised. Each of the static arguments will be broadcasted to all devices. an error is raised. Each of the static arguments will be broadcasted to
Arguments that are not arrays or containers thereof must be marked as all devices. Arguments that are not arrays or containers thereof must be
static. Defaults to (). marked as static. Defaults to ().
Static arguments must be hashable, meaning both ``__hash__`` and Static arguments must be hashable, meaning both ``__hash__`` and
``__eq__`` are implemented, and should be immutable. ``__eq__`` are implemented, and should be immutable.