From 9297dfb538e6e23cf45f6d8389b31c10645117bb Mon Sep 17 00:00:00 2001 From: Benson Ma Date: Wed, 22 Jan 2025 15:26:57 -0800 Subject: [PATCH] [fbgemm_gpu] Try nested ops namespaces - Try nested ops namespaces --- fbgemm_gpu/experimental/example/example/utils.py | 15 +++++++++++++++ .../example/test/add_tensors_float_test.py | 7 +++++++ 2 files changed, 22 insertions(+) diff --git a/fbgemm_gpu/experimental/example/example/utils.py b/fbgemm_gpu/experimental/example/example/utils.py index ffd312695b..23167bc38e 100644 --- a/fbgemm_gpu/experimental/example/example/utils.py +++ b/fbgemm_gpu/experimental/example/example/utils.py @@ -10,10 +10,25 @@ import torch +lib = torch.library.Library("fbgemm.examples", "FRAGMENT") +lib.define( + """add_tensors_float( + Tensor x, + Tensor y + ) -> Tensor + """ +) +lib.impl("add_tensors_float", torch.ops.fbgemm.add_tensors_float, "CPU") + + def add_tensors(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: return torch.ops.fbgemm.add_tensors_float(a, b) +def add_tensors_nested_ns(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: + return torch.ops.fbgemm.examples.add_tensors_float(a, b) + + def sgemm( alpha: float, TA: torch.Tensor, TB: torch.Tensor, beta: float, TC: torch.Tensor ) -> torch.Tensor: diff --git a/fbgemm_gpu/experimental/example/test/add_tensors_float_test.py b/fbgemm_gpu/experimental/example/test/add_tensors_float_test.py index b76b80fbca..eea39885eb 100644 --- a/fbgemm_gpu/experimental/example/test/add_tensors_float_test.py +++ b/fbgemm_gpu/experimental/example/test/add_tensors_float_test.py @@ -20,6 +20,13 @@ def test_add_tensors_float(self) -> None: c = utils.add_tensors(a, b) torch.testing.assert_close(c.cpu(), expected.cpu()) + def test_add_tensors_float_nested_ns(self) -> None: + a = torch.tensor([1, 2, 3]) + b = torch.tensor([4, 5, 6]) + expected = torch.tensor([5, 7, 9], dtype=torch.float) + c = utils.add_tensors_nested_ns(a, b) + torch.testing.assert_close(c.cpu(), expected.cpu()) + if __name__ == "__main__": unittest.main()