maxray

the problem with doing weird metaprogramming shit is having to deal with other people's weird metaprogramming shit

Trace and modify the result of (almost) every single expression executed in a Python program in realtime, requiring zero changes to your source code. Very WIP.

from maxray import transform, xray, maxray

from torch import tensor, Tensor, device


def move_tensor(x, ctx):
    if isinstance(x, Tensor) and x.device != device("cuda"):
        return x.to("cuda")
    return x


@transform(move_tensor)
def show_multiply(a, b):
    print(a @ b)

# Source code is rewritten to be equivalent to:

def _show_multiply(a, b):
    move_tensor(move_tensor(print, ...)(move_tensor(a, ...) @ move_tensor(b, ...)), ...)

# ---

show_multiply(
    tensor([[0.0, 1.0], [1.0, 1.0]], device="cpu"),
    tensor([[1.0], [1.0]], device="cuda"),
)  # Without the decorator, you'd expect `RuntimeError: Expected all tensors to be on the same device`

# tensor([[1.],
#         [2.]], device='cuda:0')

The ctx argument contains context information about the location of the original source code, which may be useful to build editor/LSP integrations.

The *xray decorators will recursively trace and patch every single callable they encounter until reaching either builtins, native, or generated code.

Usage

This package comes with 2 CLI commands (see --help for limited docs):

• xpy: TUI for debugging and observability. Replace python your_script.py with xpy your_script.py to apply transformations (via -W args) globally
  • -W takes file:symbol args... (be careful with shell quoting) generated by maxray template <new_script.py> and hot-reloads it on every filesystem change. Can be used multiple times to compose transformations (applied sequentially)
  • --restrict to only transform "your" (same source file or module) source code (i.e. does not descend into external packages). Significantly faster, more robust, and probably more useful - should probably be the default...
  • -m is equivalent to python -m some_module
• maxray
  • template: see --help for options

Examples

• callgraph: draws the runtime function call graph
  • compared to tools relying on static analysis, can exactly resolve calls like [fn_a, fn_b][rand() < 0.5]()

pip install maxray[all]
xpy -W 'callgraph:Draw --labels' <your_script.py>

• capture: dump all runtime captured information to an Arrow IPC file for analysis
• rerun: replace all uses of print with logging to a Rerun viewer, attaching source file and line information for filtering
• tensorcheck: global nan and inf checking for tensors
• strings: grepping of runtime values to find and rewrite URLs and filesystem paths
• progress: tqdm, but like, everywhere. Automatically shows a progress bar and the last yielded value for every for _ in iterable expected to take longer than 1s to complete
• ...: Write your own! Argument to -W can be any source file or installed package/module

Issues

If running a script or module with xpy deterministically results in differing output or exceptions compared to when run with python (CPython >= 3.11), that's a bug. A stopgap solution is to restrict the offending modules to prevent them from being transformed.