Skip to content

Document inplace_on_inputs method and remove TODO comment #1575

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
wants to merge 2 commits into
base: main
Choose a base branch
from

Conversation

Copilot
Copy link

@Copilot Copilot AI commented Aug 9, 2025

This PR addresses the documentation of the inplace_on_inputs method as requested in the issue. The method allows Op classes to create inplace versions of themselves for specified inputs, enabling memory optimizations while maintaining correctness.

Changes Made

  1. Removed TODO comment from pytensor/graph/op.py in the inplace_on_inputs method definition
  2. Added comprehensive documentation to doc/extending/inplace.rst explaining the inplace_on_inputs method

Documentation Added

The new documentation section explains:

  • Purpose: How the method creates inplace versions of ops for memory optimization
  • Current usage: Primarily used with Blockwise operations through PyTensor's rewriting system
  • Future extension: Will be extended to support core ops directly
  • Implementation guidelines: How to properly implement the method in custom ops
  • Best practices: Ensuring only allowed inputs are inplaced and returning appropriate destroy_map

The method signature and expected behavior are clearly documented:

def inplace_on_inputs(self, allowed_inplace_inputs: list[int]) -> "Op":
    """Try to return a version of self that tries to inplace in as many as `allowed_inplace_inputs`."""

This method is automatically used by PyTensor's optimization pipeline and is crucial for the efficient memory management in Blockwise operations. The documentation makes it clear that users typically don't need to call this method directly, but Op developers should understand how to implement it properly.


💬 Share your feedback on Copilot coding agent for the chance to win a $200 gift card! Click here to start the survey.

@ricardoV94 ricardoV94 changed the title [WIP] Mention this method in the Pytensor docs page on inplace (alias and destruction). It only works now with Blockwised Ops but will be extended to core ops as well. It should return a version of the Op that inplaces in as many of the requested inputs as p... Document Op.inplace_on_inputs Aug 9, 2025
Co-authored-by: ricardoV94 <28983449+ricardoV94@users.noreply.github.com>
@Copilot Copilot AI changed the title Document Op.inplace_on_inputs Document inplace_on_inputs method and remove TODO comment Aug 9, 2025
@Copilot Copilot AI requested a review from ricardoV94 August 9, 2025 03:59
Copilot finished work on behalf of ricardoV94 August 9, 2025 03:59
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants