Skip to content

Extend TorchAOBaseTensor docstring with subclassing and safetensors docs#3846

Merged
jerryzh168 merged 4 commits into
mainfrom
gh/jerryzh168/35/head
Feb 9, 2026
Merged

Extend TorchAOBaseTensor docstring with subclassing and safetensors docs#3846
jerryzh168 merged 4 commits into
mainfrom
gh/jerryzh168/35/head

Conversation

@jerryzh168

@jerryzh168 jerryzh168 commented Feb 9, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Stack from ghstack (oldest at bottom):

Summary:

Extend the TorchAOBaseTensor docstring to document:

  • Subclassing and op inheritance: how subclasses automatically inherit torch_dispatch and torch_function implementations from parents, with independent dispatch tables and MRO-based resolution for
    multiple inheritance. Includes a concrete Parent/Child example.
  • Safetensors support: the two steps needed to enable safetensors serialization for a new tensor subclass (define class attributes + register in allowlists), and how this enables save_pretrained/push_to_hub with
    safe_serialization=True in HF Transformers.
  • Minor fixes to the existing example: torch.Tensor → TorchAOBaseTensor, missing comma in init signature.

Test Plan:
visual inspection of doc preview for the PR

Reviewers:

Subscribers:

Tasks:

Tags:

@jerryzh168
Extend TorchAOBaseTensor docstring with subclassing and safetensors docs
b6d50de 
Summary:

Extend the TorchAOBaseTensor docstring to document:
- Subclassing and op inheritance: how subclasses automatically inherit __torch_dispatch__ and __torch_function__ implementations from parents, with independent dispatch tables and MRO-based resolution for
multiple inheritance. Includes a concrete Parent/Child example.
- Safetensors support: the two steps needed to enable safetensors serialization for a new tensor subclass (define class attributes + register in allowlists), and how this enables save_pretrained/push_to_hub with
safe_serialization=True in HF Transformers.
- Minor fixes to the existing example: torch.Tensor → TorchAOBaseTensor, missing comma in __init__ signature.

Test Plan:
visual inspection of doc preview for the PR

Reviewers:

Subscribers:

Tasks:

Tags:

[ghstack-poisoned]
@pytorch-bot

pytorch-bot Bot commented Feb 9, 2026 •
edited
Loading

Copy link
Copy Markdown

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/3846

Note: Links to docs will display an error until the docs builds have been completed.

This comment was automatically generated by Dr. CI and updates every 15 minutes.

jerryzh168 added a commit that referenced this pull request Feb 9, 2026
@jerryzh168
Extend TorchAOBaseTensor docstring with subclassing and safetensors docs
09242a1 
Summary:

Extend the TorchAOBaseTensor docstring to document:
- Subclassing and op inheritance: how subclasses automatically inherit __torch_dispatch__ and __torch_function__ implementations from parents, with independent dispatch tables and MRO-based resolution for
multiple inheritance. Includes a concrete Parent/Child example.
- Safetensors support: the two steps needed to enable safetensors serialization for a new tensor subclass (define class attributes + register in allowlists), and how this enables save_pretrained/push_to_hub with
safe_serialization=True in HF Transformers.
- Minor fixes to the existing example: torch.Tensor → TorchAOBaseTensor, missing comma in __init__ signature.

Test Plan:
visual inspection of doc preview for the PR

Reviewers:

Subscribers:

Tasks:

Tags:

ghstack-source-id: 8cbfd83
Pull Request resolved: #3846
@jerryzh168
Update on "Extend TorchAOBaseTensor docstring with subclassing and sa…
d0f6419 
…fetensors docs"

Summary:

Extend the TorchAOBaseTensor docstring to document:
- Subclassing and op inheritance: how subclasses automatically inherit __torch_dispatch__ and __torch_function__ implementations from parents, with independent dispatch tables and MRO-based resolution for
multiple inheritance. Includes a concrete Parent/Child example.
- Safetensors support: the two steps needed to enable safetensors serialization for a new tensor subclass (define class attributes + register in allowlists), and how this enables save_pretrained/push_to_hub with
safe_serialization=True in HF Transformers.
- Minor fixes to the existing example: torch.Tensor → TorchAOBaseTensor, missing comma in __init__ signature.

Test Plan:
visual inspection of doc preview for the PR

Reviewers:

Subscribers:

Tasks:

Tags:

[ghstack-poisoned]
@jerryzh168 jerryzh168 added topic: documentation Use this tag if this PR adds or improves documentation topic: for developers Use this tag if this PR is mainly developer facing labels Feb 9, 2026
jerryzh168 added a commit that referenced this pull request Feb 9, 2026
@jerryzh168
Extend TorchAOBaseTensor docstring with subclassing and safetensors docs
8ab914b 
Summary:

Extend the TorchAOBaseTensor docstring to document:
- Subclassing and op inheritance: how subclasses automatically inherit __torch_dispatch__ and __torch_function__ implementations from parents, with independent dispatch tables and MRO-based resolution for
multiple inheritance. Includes a concrete Parent/Child example.
- Safetensors support: the two steps needed to enable safetensors serialization for a new tensor subclass (define class attributes + register in allowlists), and how this enables save_pretrained/push_to_hub with
safe_serialization=True in HF Transformers.
- Minor fixes to the existing example: torch.Tensor → TorchAOBaseTensor, missing comma in __init__ signature.

Test Plan:
visual inspection of doc preview for the PR

Reviewers:

Subscribers:

Tasks:

Tags:

ghstack-source-id: baa54d8
Pull Request resolved: #3846
@meta-cla meta-cla Bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Feb 9, 2026
andrewor14
andrewor14 approved these changes Feb 9, 2026
View reviewed changes

@andrewor14 andrewor14 left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Looks good. I think this will also be helpful in a developer facing doc somewhere.

Comment thread torchao/utils.py
``optional_tensor_data_names``, ``optional_tensor_attribute_names``) on the subclass,
which is already required for the other utility functions above.

2. Register the subclass in ``torchao/prototype/safetensors/safetensors_utils.py``:

@andrewor14 andrewor14 Feb 9, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

does this belong here or in torchao/prototype/safetensors?

@jerryzh168 jerryzh168 Feb 9, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

it's fine to include here I feel, since it's for TorchAOBaseTensor users. but both are fine for me, let me know if you had a strong preference here

going to merge first, we can move this later if needed

@jerryzh168 jerryzh168 changed the base branch from gh/jerryzh168/35/base to main February 9, 2026 22:58
@github-actions

github-actions Bot commented Feb 9, 2026

Copy link
Copy Markdown
Contributor

📄 Doc Preview (workflow_dispatch)

You triggered a doc build manually via workflow_dispatch.
The pytorchbot doc preview link will not work for this build.

Use this link instead: https://docs-preview.pytorch.org/pytorch/ao/21844186169/index.html

@jerryzh168 jerryzh168 changed the base branch from main to gh/jerryzh168/35/base February 9, 2026 23:09
@jerryzh168
Update on "Extend TorchAOBaseTensor docstring with subclassing and sa…
d3ef1f1 
…fetensors docs"

Summary:

Extend the TorchAOBaseTensor docstring to document:
- Subclassing and op inheritance: how subclasses automatically inherit __torch_dispatch__ and __torch_function__ implementations from parents, with independent dispatch tables and MRO-based resolution for
multiple inheritance. Includes a concrete Parent/Child example.
- Safetensors support: the two steps needed to enable safetensors serialization for a new tensor subclass (define class attributes + register in allowlists), and how this enables save_pretrained/push_to_hub with
safe_serialization=True in HF Transformers.
- Minor fixes to the existing example: torch.Tensor → TorchAOBaseTensor, missing comma in __init__ signature.

Test Plan:
visual inspection of doc preview for the PR

Reviewers:

Subscribers:

Tasks:

Tags:

[ghstack-poisoned]
jerryzh168 added a commit that referenced this pull request Feb 9, 2026
@jerryzh168
Extend TorchAOBaseTensor docstring with subclassing and safetensors docs
1bd1723 
Summary:

Extend the TorchAOBaseTensor docstring to document:
- Subclassing and op inheritance: how subclasses automatically inherit __torch_dispatch__ and __torch_function__ implementations from parents, with independent dispatch tables and MRO-based resolution for
multiple inheritance. Includes a concrete Parent/Child example.
- Safetensors support: the two steps needed to enable safetensors serialization for a new tensor subclass (define class attributes + register in allowlists), and how this enables save_pretrained/push_to_hub with
safe_serialization=True in HF Transformers.
- Minor fixes to the existing example: torch.Tensor → TorchAOBaseTensor, missing comma in __init__ signature.

Test Plan:
visual inspection of doc preview for the PR

Reviewers:

Subscribers:

Tasks:

Tags:

ghstack-source-id: 5e3d3bd
Pull Request resolved: #3846
@jerryzh168 jerryzh168 changed the base branch from gh/jerryzh168/35/base to main February 9, 2026 23:10
@github-actions

github-actions Bot commented Feb 9, 2026

Copy link
Copy Markdown
Contributor

📄 Doc Preview (workflow_dispatch)

You triggered a doc build manually via workflow_dispatch.
The pytorchbot doc preview link will not work for this build.

Use this link instead: https://docs-preview.pytorch.org/pytorch/ao/21844523865/index.html

@jerryzh168 jerryzh168 changed the base branch from main to gh/jerryzh168/35/base February 9, 2026 23:25
@jerryzh168
Update on "Extend TorchAOBaseTensor docstring with subclassing and sa…
5fd6ba0 
…fetensors docs"

Summary:

Extend the TorchAOBaseTensor docstring to document:
- Subclassing and op inheritance: how subclasses automatically inherit __torch_dispatch__ and __torch_function__ implementations from parents, with independent dispatch tables and MRO-based resolution for
multiple inheritance. Includes a concrete Parent/Child example.
- Safetensors support: the two steps needed to enable safetensors serialization for a new tensor subclass (define class attributes + register in allowlists), and how this enables save_pretrained/push_to_hub with
safe_serialization=True in HF Transformers.
- Minor fixes to the existing example: torch.Tensor → TorchAOBaseTensor, missing comma in __init__ signature.

Test Plan:
visual inspection of doc preview for the PR

Reviewers:

Subscribers:

Tasks:

Tags:

[ghstack-poisoned]
jerryzh168 added a commit that referenced this pull request Feb 9, 2026
@jerryzh168
Extend TorchAOBaseTensor docstring with subclassing and safetensors docs
2ba41e6 
Summary:

Extend the TorchAOBaseTensor docstring to document:
- Subclassing and op inheritance: how subclasses automatically inherit __torch_dispatch__ and __torch_function__ implementations from parents, with independent dispatch tables and MRO-based resolution for
multiple inheritance. Includes a concrete Parent/Child example.
- Safetensors support: the two steps needed to enable safetensors serialization for a new tensor subclass (define class attributes + register in allowlists), and how this enables save_pretrained/push_to_hub with
safe_serialization=True in HF Transformers.
- Minor fixes to the existing example: torch.Tensor → TorchAOBaseTensor, missing comma in __init__ signature.

Test Plan:
visual inspection of doc preview for the PR

Reviewers:

Subscribers:

Tasks:

Tags:

ghstack-source-id: 3022eb3
Pull Request resolved: #3846
@jerryzh168 jerryzh168 changed the base branch from gh/jerryzh168/35/base to main February 9, 2026 23:28
@github-actions

github-actions Bot commented Feb 9, 2026

Copy link
Copy Markdown
Contributor

📄 Doc Preview (workflow_dispatch)

You triggered a doc build manually via workflow_dispatch.
The pytorchbot doc preview link will not work for this build.

Use this link instead: https://docs-preview.pytorch.org/pytorch/ao/21844996663/index.html

@jerryzh168 jerryzh168 merged commit 73f4a8c into main Feb 9, 2026
33 of 38 checks passed
supriyar pushed a commit that referenced this pull request Feb 10, 2026
@jerryzh168
Extend TorchAOBaseTensor docstring with subclassing and safetensors d…
5e3d091 
…ocs (#3846)

* Extend TorchAOBaseTensor docstring with subclassing and safetensors docs

Summary:

Extend the TorchAOBaseTensor docstring to document:
- Subclassing and op inheritance: how subclasses automatically inherit __torch_dispatch__ and __torch_function__ implementations from parents, with independent dispatch tables and MRO-based resolution for
multiple inheritance. Includes a concrete Parent/Child example.
- Safetensors support: the two steps needed to enable safetensors serialization for a new tensor subclass (define class attributes + register in allowlists), and how this enables save_pretrained/push_to_hub with
safe_serialization=True in HF Transformers.
- Minor fixes to the existing example: torch.Tensor → TorchAOBaseTensor, missing comma in __init__ signature.

Test Plan:
visual inspection of doc preview for the PR

Reviewers:

Subscribers:

Tasks:

Tags:

[ghstack-poisoned]

* Update on "Extend TorchAOBaseTensor docstring with subclassing and safetensors docs"


Summary:

Extend the TorchAOBaseTensor docstring to document:
- Subclassing and op inheritance: how subclasses automatically inherit __torch_dispatch__ and __torch_function__ implementations from parents, with independent dispatch tables and MRO-based resolution for
multiple inheritance. Includes a concrete Parent/Child example.
- Safetensors support: the two steps needed to enable safetensors serialization for a new tensor subclass (define class attributes + register in allowlists), and how this enables save_pretrained/push_to_hub with
safe_serialization=True in HF Transformers.
- Minor fixes to the existing example: torch.Tensor → TorchAOBaseTensor, missing comma in __init__ signature.

Test Plan:
visual inspection of doc preview for the PR

Reviewers:

Subscribers:

Tasks:

Tags:

[ghstack-poisoned]

* Update on "Extend TorchAOBaseTensor docstring with subclassing and safetensors docs"


Summary:

Extend the TorchAOBaseTensor docstring to document:
- Subclassing and op inheritance: how subclasses automatically inherit __torch_dispatch__ and __torch_function__ implementations from parents, with independent dispatch tables and MRO-based resolution for
multiple inheritance. Includes a concrete Parent/Child example.
- Safetensors support: the two steps needed to enable safetensors serialization for a new tensor subclass (define class attributes + register in allowlists), and how this enables save_pretrained/push_to_hub with
safe_serialization=True in HF Transformers.
- Minor fixes to the existing example: torch.Tensor → TorchAOBaseTensor, missing comma in __init__ signature.

Test Plan:
visual inspection of doc preview for the PR

Reviewers:

Subscribers:

Tasks:

Tags:

[ghstack-poisoned]

* Update on "Extend TorchAOBaseTensor docstring with subclassing and safetensors docs"


Summary:

Extend the TorchAOBaseTensor docstring to document:
- Subclassing and op inheritance: how subclasses automatically inherit __torch_dispatch__ and __torch_function__ implementations from parents, with independent dispatch tables and MRO-based resolution for
multiple inheritance. Includes a concrete Parent/Child example.
- Safetensors support: the two steps needed to enable safetensors serialization for a new tensor subclass (define class attributes + register in allowlists), and how this enables save_pretrained/push_to_hub with
safe_serialization=True in HF Transformers.
- Minor fixes to the existing example: torch.Tensor → TorchAOBaseTensor, missing comma in __init__ signature.

Test Plan:
visual inspection of doc preview for the PR

Reviewers:

Subscribers:

Tasks:

Tags:

[ghstack-poisoned]
@jerryzh168 jerryzh168 added the module: core changes affecting multiple modules, e.g. base config/tensor, observers, quant ops label Feb 13, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@andrewor14 andrewor14 andrewor14 approved these changes

@liangel-02 liangel-02 Awaiting requested review from liangel-02

Assignees

No one assigned

Labels

CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. module: core changes affecting multiple modules, e.g. base config/tensor, observers, quant ops topic: documentation Use this tag if this PR adds or improves documentation topic: for developers Use this tag if this PR is mainly developer facing

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@jerryzh168 @andrewor14