Searched refs:FunctionalTensorMode (Results 1 – 13 of 13) sorted by relevance
| /external/pytorch/torch/_dynamo/ |
| D | tensor_version_op.py | 5 from torch._subclasses.functional_tensor import FunctionalTensorMode
52 @_tensor_version.py_impl(FunctionalTensorMode)
57 @_unsafe_set_version_counter.py_impl(FunctionalTensorMode)
|
| /external/pytorch/torch/_library/ |
| D | triton.py | 108 from .._subclasses.functional_tensor import FunctionalTensorMode
128 result.register_torch_dispatch(FunctionalTensorMode, functional_decomp)
|
| /external/pytorch/torch/_subclasses/ |
| D | functional_tensor.py | 314 class FunctionalTensorMode(TorchDispatchMode): class
605 def dispatch_functionalize(func, mode: FunctionalTensorMode = FunctionalTensorMode()): argument
674 self, mode: Optional[FunctionalTensorMode] = None, pre_dispatch: bool = False argument
677 self.mode = mode if mode else FunctionalTensorMode()
|
| /external/pytorch/torch/_functorch/_aot_autograd/ |
| D | dispatch_and_compile_graph.py | 17 from torch._subclasses.functional_tensor import FunctionalTensorMode
48 with enable_python_dispatcher(), FunctionalTensorMode(
|
| D | collect_metadata_analysis.py | 22 from torch._subclasses.functional_tensor import FunctionalTensor, FunctionalTensorMode
165 mode = FunctionalTensorMode(_allow_token_discovery=True)
424 with FunctionalTensorMode():
|
| /external/pytorch/torch/utils/ |
| D | _python_dispatch.py | 233 from torch._subclasses.functional_tensor import FunctionalTensorMode
249 if isinstance(i, FunctionalTensorMode):
259 isinstance(old, FunctionalTensorMode)
|
| /external/pytorch/torch/ |
| D | _ops.py | 163 self.py_impl(torch._subclasses.functional_tensor.FunctionalTensorMode)(
540 from torch._subclasses.functional_tensor import FunctionalTensorMode
547 FunctionalTensorMode,
561 elif isinstance(mode, FunctionalTensorMode):
|
| /external/pytorch/test/inductor/ |
| D | test_triton_kernels.py | 122 FunctionalTensorMode,
145 with FunctionalTensorMode():
175 FunctionalTensorMode,
180 with FunctionalTensorMode():
189 with FunctionalTensorMode():
200 with FunctionalTensorMode():
220 with FunctionalTensorMode():
|
| /external/executorch/exir/ |
| D | pass_base.py | 38 from torch._subclasses.functional_tensor import FunctionalTensor, FunctionalTensorMode
86 FunctionalTensorMode
|
| /external/pytorch/test/ |
| D | test_functionalization.py | 14 FunctionalTensorMode,
2117 with maybe_disable, FunctionalTensorMode():
2267 with maybe_disable, FunctionalTensorMode():
|
| D | test_proxy_tensor.py | 14 from torch._subclasses.functional_tensor import FunctionalTensor, FunctionalTensorMode
453 a = FunctionalTensorMode(pre_dispatch=True)
478 a = FunctionalTensorMode(pre_dispatch=True)
|
| /external/pytorch/test/functorch/ |
| D | test_control_flow.py | 15 FunctionalTensorMode,
1677 mode = FunctionalTensorMode()
|
| /external/pytorch/test/dynamo/ |
| D | test_repros.py | 5200 FunctionalTensorMode,
5211 with FunctionalTensorMode():
|