Simple Python file that exemplifies how to add an op to NF4Tensor (#51)

* Simple Python file that exemplifies how to add an op to NF4Tensor

* Wording

* Move test location

Author: cpuhrsch

test/modules/test_nf4_linear.py renamed to test/dtypes/test_nf4.py

@@ -40,6 +40,10 @@ def noop_detach(func, *args, **kwargs):
 def _to_copy(func, *args, **kwargs):
     return args[0][0].get_original_weight().to(args[1]['dtype'])
 
+@implements([torch.ops.aten.to.dtype])
+def to_dtype(func, *args, **kwargs):
+    return args[0][0].get_original_weight().to(args[0][1])
+
 
 @implements(
     [

torchao/dtypes/nf4tensor.py

@@ -0,0 +1,39 @@
+import torch
+import torchao
+from torchao.dtypes import to_nf4
+
+# To create coverage for a new nf4 op we first attempt to run it
+
+# Construct a small nf4 Tensor of desired shaped
+a = torch.randn(64)
+a[0] = 0
+
+# Don't forget to pick block and scalar shapes that work for your shape
+a_nf4 = to_nf4(a, 32, 2)
+
+# Trust is good, print better
+print(f"a: {a}")
+print(f"a_nf4: {a_nf4}")
+
+
+# If GELU is not supported you'll get the following error
+# NotImplementedError: NF4Tensor dispatch: attempting to run aten.gelu.default, this is not supported
+# torch.nn.functional.gelu(a_nf4)
+
+# Next you can add this function using the implements decorator
+@torchao.dtypes.nf4tensor.implements([torch.ops.aten.gelu.default])
+def gelu(func, *args, **kwargs):
+    # The torch dispatch convention is to pass all args and kwargs via the
+    # args input.
+    # args[0] here corresponds to the original *args
+    # args[1] here corresponds to the original *kwargs
+    # We're getting the first argument of the original args
+    inp = args[0][0]
+    # There's a way very inefficient way to implement it
+    return to_nf4(torch.nn.functional.gelu(inp.to(torch.float32)), inp.block_size, inp.scaler_block_size)
+
+print(f"gelu(a): {torch.nn.functional.gelu(a)}")
+print(f"gelu(a_nf4): {torch.nn.functional.gelu(a_nf4)}")
+
+# We collect these implementations in torchao.dtypes.nf4tensor, but you can also
+# just roll your own.