Home
last modified time | relevance | path

Searched refs:ConstantExtraMetadataTensor (Results 1 – 2 of 2) sorted by relevance

/aosp_15_r20/external/pytorch/torch/testing/_internal/
H A Dcustom_tensor.py9 class ConstantExtraMetadataTensor(torch.Tensor): class
40 out = ConstantExtraMetadataTensor(elem)
49 ConstantExtraMetadataTensor, lambda x: x.elem, args
53 ConstantExtraMetadataTensor, lambda x: x.elem, kwargs
61 ConstantExtraMetadataTensor(o_inner)
/aosp_15_r20/external/pytorch/test/functorch/
H A Dtest_aotdispatch.py75 from torch.testing._internal.custom_tensor import ConstantExtraMetadataTensor
896 custom_a = ConstantExtraMetadataTensor(a)
898 custom_aa = ConstantExtraMetadataTensor(custom_a)
912 self.assertTrue(isinstance(custom_aa_compile.grad, ConstantExtraMetadataTensor))
914 isinstance(custom_aa_compile.grad.elem, ConstantExtraMetadataTensor)
986 custom_a = ConstantExtraMetadataTensor(a)
988 custom_aa = ConstantExtraMetadataTensor(custom_a)