Split implements and implements_torch_function

[Krishn1412]

1. Added two registers _ATEN_OP_TABLE and _TORCH_FN_TABLE instead of one.

2. Split the decorator into two.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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]

thanks! there are some more things to do I think.

(1). we'll need to modify all callsites that's using implements on some aten function to use implements_torch_function instead
(2). need to make sure these two things can compose,
e.g.

ao/torchao/quantization/quantize_/workflows/float8/float8_tensor.py

Line 232 in bc2c83e

@implements([torch.nn.functional.linear, aten.linear.default])

is using both, so:

@implements(aten.linear)
@implements_torch_function(F.linear)
def _(...):
     ...

should work, please add a test in

ao/test/test_utils.py

Line 89 in f03a737

class TestTorchAOBaseTensor(unittest.TestCase):

to make sure

[jerryzh168]

nit: btw, we don't need a list if it's a single op, i.e. we can do @implements(aten.linear.default)

[jerryzh168]

these are testing implementation details, I think it might be better to change these to test the user code instead, for example:

class MyTensor(TorchAOBaseTensor):
    ...
    # something like this https://github.com/pytorch/ao/blob/87769675a3e12209c4c30cd4e8563de7099d9d21/test/test_utils.py#L239-L252
    
counter = 0

def fake_linear(f, types, args, kwargs):
       x, w, b = args
       counter += 1
       return torch.matmul(x, w.T) + (b if b is not None else 0)

linear = torch.nn.Linear(...)
# swap the weight to MyTensor: https://github.com/pytorch/ao/blob/87769675a3e12209c4c30cd4e8563de7099d9d21/test/test_utils.py#L254-L256
linear.weight = MyTensor(...)

# trigger F.linear
linear(*example_inputs)

# check the value of counter to make sure the function called       

# trigger aten.linear
with torch.inference_mode():
    linear(*example_inputs)

# check the value of counter to make sure the function called       

[jerryzh168]

looks great, thanks!

[jerryzh168]

seems like there is a conflict with main, please rebase

[jerryzh168]

nit: p.data is not a recommended API I think, just initialize the parameter with MyTensor would be better:

l.weight = torch.nn.Parameter(MyTensor(orig_w, "attr", None))

[jerryzh168]

please fix ruff errors as well