Home
last modified time | relevance | path

Searched refs:functional_tensor (Results 1 – 25 of 38) sorted by relevance

12

/aosp_15_r20/external/pytorch/aten/src/ATen/
H A DFunctionalTensorWrapper.cpp632 void replace_(const Tensor& functional_tensor, const Tensor& other) { in replace_() argument
633 TORCH_INTERNAL_ASSERT_DEBUG_ONLY(isFunctionalTensor(functional_tensor)); in replace_()
634 unsafeGetFunctionalWrapper(functional_tensor)->replace_(other); in replace_()
637 void replace_(const ITensorListRef functional_tensor, ITensorListRef other) { in replace_() argument
638 TORCH_INTERNAL_ASSERT_DEBUG_ONLY(functional_tensor.size() == other.size()); in replace_()
639 auto functional_tensor_it = functional_tensor.begin(); in replace_()
641 for (C10_UNUSED const auto i : c10::irange(functional_tensor.size())) { in replace_()
646 void propagate_xla_data(const Tensor& functional_tensor, const Tensor& other) { in propagate_xla_data() argument
647 TORCH_INTERNAL_ASSERT_DEBUG_ONLY(isFunctionalTensor(functional_tensor)); in propagate_xla_data()
648 if (functional_tensor.key_set().has(c10::DispatchKey::XLA)) { in propagate_xla_data()
[all …]
H A DFunctionalTensorWrapper.h330 TORCH_API void replace_(const Tensor& functional_tensor, const Tensor& other);
332 const ITensorListRef functional_tensor,
335 TORCH_API void commit_update(const Tensor& functional_tensor);
336 TORCH_API void commit_update(ITensorListRef functional_tensor);
338 TORCH_API void unsafe_reset_storage(const Tensor& functional_tensor);
341 const Tensor& functional_tensor);
344 const Tensor& functional_tensor);
347 const Tensor& functional_tensor);
352 const Tensor& functional_tensor,
355 const ITensorListRef functional_tensor,
/aosp_15_r20/external/pytorch/torch/_functorch/_aot_autograd/
H A Dinput_output_analysis.py19 from torch._subclasses.functional_tensor import FunctionalTensor
70 functional_tensor=o.functional_tensor,
234 functional_tensor=o.functional_tensor,
H A Dcollect_metadata_analysis.py22 from torch._subclasses.functional_tensor import FunctionalTensor, FunctionalTensorMode
590 functional_tensor = None
622 functional_tensor = FunctionalTensorMetadataEq(o.elem)
630 functional_tensor=functional_tensor,
H A Dfunctional_utils.py17 from torch._subclasses.functional_tensor import FunctionalTensor
249 functional_tensor = target_functional_tensor.tensor
252 functional_tensor, aliased_base_tensor
H A Druntime_wrappers.py173 self.functional_tensor = info.functional_tensor
182 self.functional_tensor,
210 self.functional_tensor = info.functional_tensor
219 self.functional_tensor,
H A Ddispatch_and_compile_graph.py17 from torch._subclasses.functional_tensor import FunctionalTensorMode
134 buffer, torch._subclasses.functional_tensor.FunctionalTensor
H A Dschemas.py102 functional_tensor: Optional[FunctionalTensorMetadataEq] = None variable in OutputAliasInfo
529 new_out = dataclasses.replace(out, functional_tensor=None)
H A Dtraced_function_transforms.py549 from torch._subclasses.functional_tensor import (
/aosp_15_r20/external/pytorch/torch/_higher_order_ops/
H A Dcond.py6 import torch._subclasses.functional_tensor
29 from torch._subclasses.functional_tensor import disable_functional_mode
H A Dstrict_mode.py3 import torch._subclasses.functional_tensor
H A Dflex_attention.py388 ctx: torch._subclasses.functional_tensor.BaseFunctionalizeAPI,
488 from torch._subclasses.functional_tensor import disable_functional_mode
954 ctx: torch._subclasses.functional_tensor.BaseFunctionalizeAPI,
H A Dauto_functionalize.py333 from torch._subclasses.functional_tensor import PythonFunctionalizeAPI
415 from torch._subclasses.functional_tensor import PythonFunctionalizeAPI
H A Dmap.py16 from torch._subclasses.functional_tensor import disable_functional_mode
H A Deffects.py261 from torch._subclasses.functional_tensor import PythonFunctionalizeAPI
H A Dutils.py218 from torch._subclasses.functional_tensor import FunctionalTensor
H A Dassociative_scan.py8 import torch._subclasses.functional_tensor
/aosp_15_r20/external/pytorch/torch/autograd/
H A Dgraph.py547 torch._subclasses.functional_tensor.FunctionalTensor,
562 torch._subclasses.functional_tensor.FunctionalTensor,
/aosp_15_r20/external/pytorch/torch/_dynamo/
H A Dtensor_version_op.py5 from torch._subclasses.functional_tensor import FunctionalTensorMode
/aosp_15_r20/external/pytorch/torch/_library/
H A Dtriton.py108 from .._subclasses.functional_tensor import FunctionalTensorMode
/aosp_15_r20/external/pytorch/torch/utils/
H A D_python_dispatch.py233 from torch._subclasses.functional_tensor import FunctionalTensorMode
660 from torch._subclasses.functional_tensor import FunctionalTensor
/aosp_15_r20/external/pytorch/torch/
H A D_ops.py145 from torch._subclasses.functional_tensor import (
163 self.py_impl(torch._subclasses.functional_tensor.FunctionalTensorMode)(
540 from torch._subclasses.functional_tensor import FunctionalTensorMode
/aosp_15_r20/external/pytorch/test/
H A Dtest_functionalization.py11 from torch._subclasses.functional_tensor import (
482 functional_tensor = torch.ones(2, dtype=torch.long)
483 out = x[functional_tensor, nonfunctional_tensor]
/aosp_15_r20/external/pytorch/torch/_export/
H A D__init__.py114 from torch._subclasses.functional_tensor import FunctionalTensor
/aosp_15_r20/external/pytorch/torch/nested/_internal/
H A Dnested_tensor.py17 from torch._subclasses.functional_tensor import mb_unwrap_functional_tensor

12