Is `heapq` module supported for compilation by Dynamo?

torch.compile torch._dynamo
1

heapq Python module is implementing heap-based priority queues: https://github.com/python/cpython/blob/3.11/Lib/heapq.py

Supporting these basic ops is useful for compiling Dijkstra-like loops which treat graph edges or nodes one by one based on some priority recalculation.

Ideally, these Python loops (using heapq.pop and heapq.push) could be compiled to C++ at least…

2

This graph broke for me

import heapq
import torch

@torch.compile(fullgraph=True)
def program():
  h = []
  heapq.heappush(h, 3)
  heapq.heappush(h, 1)
  heapq.heappush(h, 4)
  val = heapq.heappop(h)
  return val * torch.randn(10)

program()

@voz on github might say this is a good dynamo starter task XD

1 Like
3

Created an issue Dynamo graph break when using pyton module `heapq` (manipulates with `list`s) · Issue #106885 · pytorch/pytorch · GitHub :slight_smile:

1 Like