opentau.utils.fake_tensor

Utilities for working with PyTorch FakeTensor.

This module provides a FakeTensorContext class and decorator for running code with FakeTensor mode enabled, which is useful for shape inference and testing without actual tensor computations.

Functions

run_with_fake_tensor(fn)

Decorator to run a function with FakeTensor enabled.

Classes

FakeTensorContext([allow_non_fake_inputs])

Context manager for enabling FakeTensor mode with necessary patches.

class opentau.utils.fake_tensor.FakeTensorContext(allow_non_fake_inputs: bool = True)

Bases: object

Context manager for enabling FakeTensor mode with necessary patches.

This context manager applies all necessary monkey patches for FakeTensor compatibility and manages the FakeTensorMode lifecycle.

Parameters:

allow_non_fake_inputs – If True, allow non-fake tensors as inputs. Defaults to True.

__init__(allow_non_fake_inputs: bool = True)

opentau.utils.fake_tensor.run_with_fake_tensor(fn)

Decorator to run a function with FakeTensor enabled.

Parameters:

fn – Function to wrap.

Returns:

Wrapped function that runs with FakeTensorContext enabled.