mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
5c8a9803c8
Summary: This is a follow up to #78015. This PR * introduces namespace logic for generating `NativeFunctions.h`. * adds helper function to extract namespace from string * relaxes the constraint on the levels we support for custom kernel namespace to 2 Test Plan: Yaml entry: ``` - func: unsqueeze.out(Tensor(a) self, int dim, *, Tensor(a!) out) -> Tensor(a!) variants: function device_check: NoCheck dispatch: CPU: custom_1::custom_2::unsqueeze ``` Generated `NativeFunctions.h`: ``` namespace custom_1 { namespace custom_2 { namespace native { TORCH_API at::Tensor & unsqueeze(const at::Tensor & self, int64_t dim, at::Tensor & out); } // namespace native } // namespace custom_2 } // namespace custom_1 ``` Differential Revision: D37198111 Pull Request resolved: https://github.com/pytorch/pytorch/pull/79733 Approved by: https://github.com/bdhirsh
23 lines
870 B
Python
23 lines
870 B
Python
import unittest
|
|
|
|
from torchgen.utils import NamespaceHelper
|
|
|
|
|
|
class TestNamespaceHelper(unittest.TestCase):
|
|
def test_create_from_namespaced_tuple(self) -> None:
|
|
helper = NamespaceHelper.from_namespaced_entity("aten::add")
|
|
self.assertEqual(helper.entity_name, "add")
|
|
self.assertEqual(helper.get_cpp_namespace(), "aten")
|
|
|
|
def test_default_namespace(self) -> None:
|
|
helper = NamespaceHelper.from_namespaced_entity("add")
|
|
self.assertEqual(helper.entity_name, "add")
|
|
self.assertEqual(helper.get_cpp_namespace(), "")
|
|
self.assertEqual(helper.get_cpp_namespace("default"), "default")
|
|
|
|
def test_namespace_levels_more_than_max(self) -> None:
|
|
with self.assertRaises(AssertionError):
|
|
NamespaceHelper(
|
|
namespace_str="custom_1::custom_2", entity_name="", max_level=1
|
|
)
|