opentau.utils.monkey_patch

Monkey patches to inject behaviour into computational graph construction. This eliminates the need to trace and modify source code of 3rd party libraries, such as transformers or even PyTorch itself.

Where necessary, we can do

>>> from opentau.utils.monkey_patch import torch_cumsum_patch, torch_pow_patch
>>> torch_cumsum_patch()  # Apply the patch to handle bool tensors in cumsum
>>> torch_pow_patch()  # Apply the patch to handle mixed number-tensor exponentiation in pow

to apply the patches.

Running the same monkey patch twice will not have any effect, as the 2nd call will return immediately.

Note: currently, there is no way to undo a monkey patch once it has been applied, which can be a future to-do. For now, only apply the patch when its implications are understood.

Functions

gym_is_gymnasium_patch()	Monkey patch to make import gym equivalent to import gymnasium as gym.
patches_applied()	Get a list of all patches that have been applied.
torch_cumsum_patch()	Override torch.cumsum to handle bool tensors correctly.
torch_fake_tensor_beta_validate_args_patch()	Fix torch.distributions.Beta to work in FakeTensorMode.
torch_fake_tensor_is_inf_patch()	Patch torch.isinf to work with FakeTensor.
torch_fake_tensor_module_to_patch()	Prefer overwriting (not in-place updating) module params on .to() conversion.
torch_fake_tensor_to_numpy_patch()	Enable .numpy() calls on FakeTensor to return random numpy arrays.
torch_full_patch()	Override torch.full to convert bool fill values to int.
torch_load_patch()	Override torch.load to handle weights_only argument.
torch_pow_patch()	Override torch.pow to ensure both base and exponent are tensors.

opentau.utils.monkey_patch.gym_is_gymnasium_patch()

Monkey patch to make import gym equivalent to import gymnasium as gym.

This patch is necessary because the original gym package is incompatible with numpy >= 2.0. It redirects gym imports to use gymnasium instead.

opentau.utils.monkey_patch.patches_applied()

Get a list of all patches that have been applied.

Returns:

List of patch function names that have been applied.

opentau.utils.monkey_patch.torch_cumsum_patch()

Override torch.cumsum to handle bool tensors correctly.

PyTorch allows cumsum on bool tensors, but ONNX Runtime does not. This patch converts bool tensors to int64 before calling cumsum.

opentau.utils.monkey_patch.torch_fake_tensor_beta_validate_args_patch()

Fix torch.distributions.Beta to work in FakeTensorMode.

This patch sets validate_args=False by default for Beta distributions, which is required for FakeTensorMode compatibility.

opentau.utils.monkey_patch.torch_fake_tensor_is_inf_patch()

Patch torch.isinf to work with FakeTensor.

This patch provides a mock implementation of torch.isinf that returns a mock object compatible with FakeTensor operations.

opentau.utils.monkey_patch.torch_fake_tensor_module_to_patch()

Prefer overwriting (not in-place updating) module params on .to() conversion.

Historically this flag made Module.to(device) behave sensibly in FakeTensorMode (see https://github.com/pytorch/pytorch/issues/119665). On modern torch (confirmed on the locked 2.10) it is no longer sufficient: nn.Module._apply hardcodes the torch.utils.swap_tensors path for any FakeTensor param (... or isinstance(param, FakeTensor)), and swap_tensors rejects them because FakeTensorMode tracks them with weakrefs — neither this flag nor set_swap_module_params_on_conversion can steer around it. So the supported pattern is to construct modules directly on the target device inside FakeTensorMode (nn.Linear(..., device=device)) rather than building on CPU and calling .to(device) afterwards.

The flag is retained as a harmless, defensive default (it is a no-op for FakeTensor params); it only affects the now-unused Module._apply conversion path.

opentau.utils.monkey_patch.torch_fake_tensor_to_numpy_patch()

Enable .numpy() calls on FakeTensor to return random numpy arrays.

This patch allows .numpy() to be called on FakeTensor instances, returning numpy arrays with random values. Note that calling .numpy() multiple times on the same FakeTensor may return different values.

opentau.utils.monkey_patch.torch_full_patch()

Override torch.full to convert bool fill values to int.

This patch ensures that True/False are converted to 1/0 before reaching the C++ level, improving compatibility with certain backends.

opentau.utils.monkey_patch.torch_load_patch()

Override torch.load to handle weights_only argument.

This patch ensures that torch.load properly handles the weights_only argument for PyTorch versions >= 2.6, setting it to False by default if not explicitly provided.

opentau.utils.monkey_patch.torch_pow_patch()

Override torch.pow to ensure both base and exponent are tensors.

This patch converts scalar arguments to tensors before calling pow, ensuring compatibility with ONNX export.