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

Change TORCH_LIBRARY to TORCH_LIBRARY_FRAGMENT #1645

Merged
merged 23 commits into from
Feb 12, 2025
Original file line number Diff line number Diff line change
Expand Up @@ -60,7 +60,7 @@
"_linear_8bit_act_" #weight_nbit "bit_weight", \
&linear_meta<weight_nbit, true>);

TORCH_LIBRARY(torchao, m) {
TORCH_LIBRARY_FRAGMENT(torchao, m) {
DEFINE_OP(1);
DEFINE_OP(2);
DEFINE_OP(3);
Expand Down
Copy link
Contributor Author

Choose a reason for hiding this comment

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

@manuelcandales how do I test the change? Is there a script I should run?

Copy link
Contributor

Choose a reason for hiding this comment

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

cd ao/torchao/experimental/ops/mps
bash build.sh
python test/test_lowbit.py

Original file line number Diff line number Diff line change
Expand Up @@ -163,7 +163,7 @@ Tensor pack_weights_cpu_kernel(const Tensor& W) {
return B;
}

TORCH_LIBRARY(torchao, m) {
TORCH_LIBRARY_FRAGMENT(torchao, m) {
m.def("_pack_weight_1bit(Tensor W) -> Tensor");
m.def("_pack_weight_2bit(Tensor W) -> Tensor");
m.def("_pack_weight_3bit(Tensor W) -> Tensor");
Expand Down
Loading