frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-20 21:14:14 +08:00
Code Issues Packages Projects Releases Wiki Activity

[functorch] update compile example imports (pytorch/functorch#834)

Browse Source
This commit is contained in:
Andre
2022-05-26 12:09:18 -07:00
committed by Jon Janzen
parent 36fa9d8295
commit 1feff6bb69
3 changed files with 8 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
functorch/examples/compilation/simple_function.py
View File

@ -4,7 +4,8 @@
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.
from functorch import grad, nnc_jit, make_fx, make_nnc
from functorch import grad, make_fx
from functorch.compile import nnc_jit
import torch
import time
@ -16,9 +17,7 @@ def f(x):
inp = torch.randn(100)
grad_pt = grad(f)
grad_fx = make_fx(grad_pt)(inp)
grad_nnc = nnc_jit(grad_pt, skip_specialization=True)
loopnest = make_nnc(grad_pt)(inp)
print(loopnest)
grad_nnc = nnc_jit(grad_pt)
def bench(name, f, iters=10000, warmup=3):