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.

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

fix lifetimes, add tensors iterator #518

Merged
merged 2 commits into from
Sep 5, 2024
Merged

Conversation

gvilums
Copy link
Contributor

@gvilums gvilums commented Aug 19, 2024

What does this PR do?

  • Fixes the lifetimes of the TensorViews returned by the top-level SafeTensors object to depend on the underlying data instead of the SafeTensors object itself. This makes the API more flexible by allowing TensorViews to outlive the SafeTensors object that they were created from.
  • Adds a method to SafeTensors to iterate over all contained TensorViews

@galqiwi
Copy link

galqiwi commented Aug 25, 2024

I just barely created the same pr, and found out that this one already exists :)
The lifetime part is very crucial for me -- to the extent that I'm planning to use your branch in my pet project.
Maintainers, please review this PR

@Narsil
Copy link
Collaborator

Narsil commented Sep 4, 2024

Can you provide an actual example where this is impactful ?

I think adding a test showcasing the usefulness would help prevent future errors on this.

@gvilums
Copy link
Contributor Author

gvilums commented Sep 4, 2024

@Narsil The basic pattern that this allows is the following:

        let serialized = b"...";

        let tensor = {
            let loaded = SafeTensors::deserialize(serialized).unwrap();
            loaded.tensor("test").unwrap()
        };
        
        do_something_with(tensor)

We want to get all the stored tensors without being tied to the lifetime of the SafeTensors object. This is of course a naive example - in our code we need this to be able to store our deserialized tensors without having to worry about keeping the SafeTensors object alive, which can get complicated because it itself borrows from the underlying data.

I added a test to check for this simple case.

Copy link
Collaborator

@Narsil Narsil left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks a lot for the tests/clarification. That makes it clearer and we can keep this for longer.

I'll release after this.

@Narsil Narsil merged commit ba2e397 into huggingface:main Sep 5, 2024
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.

3 participants