tensor_is_on_xla() — pytorch Function Reference

Architecture documentation for the tensor_is_on_xla() function in common.py from the pytorch codebase.

Function python DynamoBenchmarks Configuration called by 1

Entity Profile

DynamoBenchmarks→ Configuration→ tensor_is_on_xla() — pytorch Function Reference

Dependency Diagram

graph TD
  82d098c4_9a2a_eaf4_0bba_655d342fc39e["tensor_is_on_xla()"]
  9c8df7bf_0e05_9bbb_5e2f_6c88f28b52d4["timed()"]
  9c8df7bf_0e05_9bbb_5e2f_6c88f28b52d4 -->|calls| 82d098c4_9a2a_eaf4_0bba_655d342fc39e
  style 82d098c4_9a2a_eaf4_0bba_655d342fc39e fill:#6366f1,stroke:#818cf8,color:#fff

Relationship Graph

Source Code

benchmarks/dynamo/common.py lines 652–660

def tensor_is_on_xla(tensors):
    def visit(x: torch.Tensor):
        nonlocal result
        if x.device.type == "xla":
            result = True

    result = False
    tree_map_only(torch.Tensor, visit, tensors)
    return result