Skip to content
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.

Sign up for GitHub

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

Jump to bottom

Propagate CopySemantics from python to C++ transfer APIs so that device_put works correctly in presence of copy/donate options that user specified. #24752

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Propagate CopySemantics from python to C++ transfer APIs so that device_put works correctly in presence of copy/donate options that user specified. #24752

wants to merge 1 commit into from

Commits on Nov 7, 2024

  1. Propagate CopySemantics from python to C++ transfer APIs so that devi… 

    …ce_put works correctly in presence of copy/donate options that user specified.

This change only supports pinned_host -> pinned_host copies on the same device. HBM -> HBM copies don't work yet and donation also doesn't work in PJRT.

This CL also sets up the plumbing from JAX to PJRT so that in the future support for missing features can be added easily.

Fixes #24521

PiperOrigin-RevId: 693801738
    @yashk2810 @Google-ML-Automation
    yashk2810 authored and Google-ML-Automation committed Nov 7, 2024
    Configuration menu
    Copy the full SHA
    9033ec8 View commit details
    Browse the repository at this point in the history