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

docstring for jax.stages.Wrapped

Browse Source
This commit is contained in:
Roy Frostig 2022-09-01 18:30:33 -07:00
parent 0869183107
commit 4505d57a60
1 changed files with 8 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
jax/_src/stages.py
View File

@ -497,6 +497,14 @@ class Lowered(Stage):
class Wrapped(Protocol):
"""A function ready to be specialized, lowered, and compiled.
This protocol reflects the output of functions such as
``jax.jit``. Calling it results in JIT (just-in-time) lowering,
compilation, and execution. It can also be explicitly lowered prior
to compilation, and the result compiled prior to execution.
"""
def __call__(self, *args, **kwargs):
"""Executes the wrapped function, lowering and compiling as needed."""
raise NotImplementedError