mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 00:21:07 +01:00
4048d4cdd2
Adds a prototype tracer with no caching support and the `ElementwiseUnaryPythonRefInfo` class. A reference for `floor` is added to test the latter, and the elementwise binary reference inputs are extended to also return noncontiguous inputs. The SampleInput transform operation has been updated to return an actual SampleInput instead of a tuple to facilitate uniform handling of (transformed) SampleInputs. Pull Request resolved: https://github.com/pytorch/pytorch/pull/76388 Approved by: https://github.com/ngimel
14 lines
289 B
Python
14 lines
289 B
Python
from torch.fx import GraphModule
|
|
from torch._prims.context import PrimContext
|
|
|
|
|
|
def execute(ctx: PrimContext, *args, **kwargs):
|
|
"""
|
|
Prototype ATen executor.
|
|
|
|
Just executes the context's graph.
|
|
"""
|
|
|
|
gm = GraphModule({}, ctx.graph)
|
|
return gm.forward(*args, **kwargs)
|