Searched refs:allow_cpu_scalar_tensors (Results 1 – 4 of 4) sorted by relevance
| /aosp_15_r20/external/pytorch/torch/_prims_common/ |
| H A D | __init__.py | 427 check_same_shape(*tensors, allow_cpu_scalar_tensors=True)
522 check_same_shape(*tensors, allow_cpu_scalar_tensors=True)
710 def check_same_device(*args, allow_cpu_scalar_tensors): argument
728 if allow_cpu_scalar_tensors and is_cpu_scalar_tensor(arg):
761 def check_same_shape(*args, allow_cpu_scalar_tensors: bool):
775 if allow_cpu_scalar_tensors and is_cpu_scalar_tensor(arg):
793 def extract_shape(*args, allow_cpu_scalar_tensors: bool) -> Optional[ShapeType]:
801 if allow_cpu_scalar_tensors and is_cpu_scalar_tensor(arg):
|
| /aosp_15_r20/external/pytorch/torch/_prims/ |
| H A D | __init__.py | 402 utils.check_same_device(*args_, allow_cpu_scalar_tensors=True)
403 utils.check_same_shape(*args_, allow_cpu_scalar_tensors=True)
406 shape = utils.extract_shape(*args_, allow_cpu_scalar_tensors=True)
|
| /aosp_15_r20/external/pytorch/torch/_decomp/ |
| H A D | decompositions.py | 1576 grad_output, input, mean, rstd, allow_cpu_scalar_tensors=False
1578 utils.check_same_shape(input, grad_output, allow_cpu_scalar_tensors=False)
1579 utils.check_same_shape(mean, rstd, allow_cpu_scalar_tensors=False)
|
| /aosp_15_r20/external/pytorch/torch/_refs/ |
| H A D | __init__.py | 1928 utils.check_same_device(pred, a, b, allow_cpu_scalar_tensors=True)
2737 utils.check_same_device(*tensors, allow_cpu_scalar_tensors=False)
5665 utils.check_same_device(a, b, allow_cpu_scalar_tensors=False)
|