From 4e1b7c4228ed6d8d1d5dfba7dabbc845c3507178 Mon Sep 17 00:00:00 2001 From: haleqiu Date: Sat, 14 Mar 2026 22:09:57 -0400 Subject: [PATCH 1/3] add a default warpper --- __init__.py | 115 ++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 115 insertions(+) create mode 100644 __init__.py diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..6ad2c42 --- /dev/null +++ b/__init__.py @@ -0,0 +1,115 @@ +"""Enable PyPose monkey-patching with the warp backend.""" + +from typing import Any, cast + +import pypose +from functools import partial +from . import pypose_warp + +_wp = cast(Any, pypose_warp) + +# Save original types before monkey patching (for from_matrix comparison) +_orig_SE3_type = pypose.SE3_type +_orig_SO3_type = pypose.SO3_type + +# Save original functions before monkey patching +_orig_identity_SE3 = pypose.identity_SE3 +_orig_identity_SO3 = pypose.identity_SO3 +_orig_identity_se3 = pypose.identity_se3 +_orig_identity_so3 = pypose.identity_so3 +_orig_randn_SE3 = pypose.randn_SE3 +_orig_randn_SO3 = pypose.randn_SO3 +_orig_randn_se3 = pypose.randn_se3 +_orig_randn_so3 = pypose.randn_so3 +_orig_from_matrix = pypose.from_matrix + + +# Wrapper functions that use warp types +def wp_identity_SE3(*args, **kwargs): + result = _orig_identity_SE3(*args, **kwargs) + result.ltype = _wp.warpSE3_type + return result + + +def wp_identity_SO3(*args, **kwargs): + result = _orig_identity_SO3(*args, **kwargs) + result.ltype = _wp.warpSO3_type + return result + + +def wp_identity_se3(*args, **kwargs): + result = _orig_identity_se3(*args, **kwargs) + result.ltype = _wp.warpse3_type + return result + + +def wp_identity_so3(*args, **kwargs): + result = _orig_identity_so3(*args, **kwargs) + result.ltype = _wp.warpso3_type + return result + + +def wp_randn_SE3(*args, **kwargs): + result = _orig_randn_SE3(*args, **kwargs) + result.ltype = _wp.warpSE3_type + return result + + +def wp_randn_SO3(*args, **kwargs): + result = _orig_randn_SO3(*args, **kwargs) + result.ltype = _wp.warpSO3_type + return result + + +def wp_randn_se3(*args, **kwargs): + result = _orig_randn_se3(*args, **kwargs) + result.ltype = _wp.warpse3_type + return result + + +def wp_randn_so3(*args, **kwargs): + result = _orig_randn_so3(*args, **kwargs) + result.ltype = _wp.warpso3_type + return result + + +def wp_from_matrix(mat, ltype, check=True, rtol=1e-5, atol=1e-5): + # Map warp types to original pypose types for internal comparison + if ltype is _wp.warpSE3_type: + result = _orig_from_matrix(mat, _orig_SE3_type, check=check, rtol=rtol, atol=atol) + result.ltype = _wp.warpSE3_type + elif ltype is _wp.warpSO3_type: + result = _orig_from_matrix(mat, _orig_SO3_type, check=check, rtol=rtol, atol=atol) + result.ltype = _wp.warpSO3_type + else: + # For non-warp types (Sim3, RxSO3), use original function directly + result = _orig_from_matrix(mat, ltype, check=check, rtol=rtol, atol=atol) + return result + + +# Monkey patching - Type objects +pypose.SE3_type = _wp.warpSE3_type +pypose.SO3_type = _wp.warpSO3_type +pypose.se3_type = _wp.warpse3_type +pypose.so3_type = _wp.warpso3_type + +# Monkey patching - Constructor functions +pypose.SE3 = partial(pypose.LieTensor, ltype=_wp.warpSE3_type) +pypose.SO3 = partial(pypose.LieTensor, ltype=_wp.warpSO3_type) +pypose.se3 = partial(pypose.LieTensor, ltype=_wp.warpse3_type) +pypose.so3 = partial(pypose.LieTensor, ltype=_wp.warpso3_type) + +# Monkey patching - Identity functions +pypose.identity_SE3 = wp_identity_SE3 +pypose.identity_SO3 = wp_identity_SO3 +pypose.identity_se3 = wp_identity_se3 +pypose.identity_so3 = wp_identity_so3 + +# Monkey patching - Random generation functions +pypose.randn_SE3 = wp_randn_SE3 +pypose.randn_SO3 = wp_randn_SO3 +pypose.randn_se3 = wp_randn_se3 +pypose.randn_so3 = wp_randn_so3 + +# Monkey patching - Matrix conversion +pypose.from_matrix = wp_from_matrix From 8c6806a4a2c52e21ee6d098feccac54534179adf Mon Sep 17 00:00:00 2001 From: haleqiu Date: Sat, 14 Mar 2026 22:46:44 -0400 Subject: [PATCH 2/3] Remove global monkey-patching, expose warp-backed constructors as module-level exports Warp-backed constructors (SO3, so3, SE3, se3, identity_*, randn_*, euler2SO3, mat2SO3, mat2SE3, from_matrix) are now local module attributes. __getattr__ delegates non-Lie APIs to the unpatched pypose module, allowing PyposeWarp and pypose to coexist without interference. Made-with: Cursor --- __init__.py | 150 ++++++++++++++++++++++++++++++---------------------- 1 file changed, 87 insertions(+), 63 deletions(-) diff --git a/__init__.py b/__init__.py index 6ad2c42..0cd1ccb 100644 --- a/__init__.py +++ b/__init__.py @@ -1,4 +1,10 @@ -"""Enable PyPose monkey-patching with the warp backend.""" +"""Warp-accelerated PyPose Lie operations without global monkey-patching. + +``import PyposeWarp as pp`` gives warp-backed constructors (so3, SO3, se3, SE3, +identity_*, randn_*, euler2SO3, mat2SO3, mat2SE3, from_matrix). Any attribute +not explicitly defined here (e.g. ``optim``, ``module``, ``vec2skew``) is +delegated to the *unpatched* ``pypose`` module via ``__getattr__``. +""" from typing import Any, cast @@ -8,108 +14,126 @@ _wp = cast(Any, pypose_warp) -# Save original types before monkey patching (for from_matrix comparison) _orig_SE3_type = pypose.SE3_type _orig_SO3_type = pypose.SO3_type -# Save original functions before monkey patching -_orig_identity_SE3 = pypose.identity_SE3 -_orig_identity_SO3 = pypose.identity_SO3 -_orig_identity_se3 = pypose.identity_se3 -_orig_identity_so3 = pypose.identity_so3 -_orig_randn_SE3 = pypose.randn_SE3 -_orig_randn_SO3 = pypose.randn_SO3 -_orig_randn_se3 = pypose.randn_se3 -_orig_randn_so3 = pypose.randn_so3 -_orig_from_matrix = pypose.from_matrix - - -# Wrapper functions that use warp types -def wp_identity_SE3(*args, **kwargs): - result = _orig_identity_SE3(*args, **kwargs) +# ====== Warp-backed type objects ====== + +SE3_type = _wp.warpSE3_type +SO3_type = _wp.warpSO3_type +se3_type = _wp.warpse3_type +so3_type = _wp.warpso3_type + +# ====== Warp-backed constructors ====== + +SO3 = partial(pypose.LieTensor, ltype=_wp.warpSO3_type) +so3 = partial(pypose.LieTensor, ltype=_wp.warpso3_type) +SE3 = partial(pypose.LieTensor, ltype=_wp.warpSE3_type) +se3 = partial(pypose.LieTensor, ltype=_wp.warpse3_type) + +LieTensor = pypose.LieTensor + +# ====== Warp-backed identity functions ====== + + +def identity_SE3(*args: Any, **kwargs: Any) -> pypose.LieTensor: + result = pypose.identity_SE3(*args, **kwargs) result.ltype = _wp.warpSE3_type return result -def wp_identity_SO3(*args, **kwargs): - result = _orig_identity_SO3(*args, **kwargs) +def identity_SO3(*args: Any, **kwargs: Any) -> pypose.LieTensor: + result = pypose.identity_SO3(*args, **kwargs) result.ltype = _wp.warpSO3_type return result -def wp_identity_se3(*args, **kwargs): - result = _orig_identity_se3(*args, **kwargs) +def identity_se3(*args: Any, **kwargs: Any) -> pypose.LieTensor: + result = pypose.identity_se3(*args, **kwargs) result.ltype = _wp.warpse3_type return result -def wp_identity_so3(*args, **kwargs): - result = _orig_identity_so3(*args, **kwargs) +def identity_so3(*args: Any, **kwargs: Any) -> pypose.LieTensor: + result = pypose.identity_so3(*args, **kwargs) result.ltype = _wp.warpso3_type return result -def wp_randn_SE3(*args, **kwargs): - result = _orig_randn_SE3(*args, **kwargs) +# ====== Warp-backed randn functions ====== + + +def randn_SE3(*args: Any, **kwargs: Any) -> pypose.LieTensor: + result = pypose.randn_SE3(*args, **kwargs) result.ltype = _wp.warpSE3_type return result -def wp_randn_SO3(*args, **kwargs): - result = _orig_randn_SO3(*args, **kwargs) +def randn_SO3(*args: Any, **kwargs: Any) -> pypose.LieTensor: + result = pypose.randn_SO3(*args, **kwargs) result.ltype = _wp.warpSO3_type return result -def wp_randn_se3(*args, **kwargs): - result = _orig_randn_se3(*args, **kwargs) +def randn_se3(*args: Any, **kwargs: Any) -> pypose.LieTensor: + result = pypose.randn_se3(*args, **kwargs) result.ltype = _wp.warpse3_type return result -def wp_randn_so3(*args, **kwargs): - result = _orig_randn_so3(*args, **kwargs) +def randn_so3(*args: Any, **kwargs: Any) -> pypose.LieTensor: + result = pypose.randn_so3(*args, **kwargs) result.ltype = _wp.warpso3_type return result -def wp_from_matrix(mat, ltype, check=True, rtol=1e-5, atol=1e-5): - # Map warp types to original pypose types for internal comparison +# ====== Warp-backed conversion functions ====== + + +def from_matrix(mat: Any, ltype: Any, check: bool = True, + rtol: float = 1e-5, atol: float = 1e-5) -> pypose.LieTensor: if ltype is _wp.warpSE3_type: - result = _orig_from_matrix(mat, _orig_SE3_type, check=check, rtol=rtol, atol=atol) + result = pypose.from_matrix(mat, _orig_SE3_type, check=check, rtol=rtol, atol=atol) result.ltype = _wp.warpSE3_type elif ltype is _wp.warpSO3_type: - result = _orig_from_matrix(mat, _orig_SO3_type, check=check, rtol=rtol, atol=atol) + result = pypose.from_matrix(mat, _orig_SO3_type, check=check, rtol=rtol, atol=atol) result.ltype = _wp.warpSO3_type else: - # For non-warp types (Sim3, RxSO3), use original function directly - result = _orig_from_matrix(mat, ltype, check=check, rtol=rtol, atol=atol) + result = pypose.from_matrix(mat, ltype, check=check, rtol=rtol, atol=atol) + return result + + +def euler2SO3(*args: Any, **kwargs: Any) -> pypose.LieTensor: + result = pypose.euler2SO3(*args, **kwargs) + result.ltype = _wp.warpSO3_type return result -# Monkey patching - Type objects -pypose.SE3_type = _wp.warpSE3_type -pypose.SO3_type = _wp.warpSO3_type -pypose.se3_type = _wp.warpse3_type -pypose.so3_type = _wp.warpso3_type - -# Monkey patching - Constructor functions -pypose.SE3 = partial(pypose.LieTensor, ltype=_wp.warpSE3_type) -pypose.SO3 = partial(pypose.LieTensor, ltype=_wp.warpSO3_type) -pypose.se3 = partial(pypose.LieTensor, ltype=_wp.warpse3_type) -pypose.so3 = partial(pypose.LieTensor, ltype=_wp.warpso3_type) - -# Monkey patching - Identity functions -pypose.identity_SE3 = wp_identity_SE3 -pypose.identity_SO3 = wp_identity_SO3 -pypose.identity_se3 = wp_identity_se3 -pypose.identity_so3 = wp_identity_so3 - -# Monkey patching - Random generation functions -pypose.randn_SE3 = wp_randn_SE3 -pypose.randn_SO3 = wp_randn_SO3 -pypose.randn_se3 = wp_randn_se3 -pypose.randn_so3 = wp_randn_so3 - -# Monkey patching - Matrix conversion -pypose.from_matrix = wp_from_matrix +def mat2SO3(*args: Any, **kwargs: Any) -> pypose.LieTensor: + result = pypose.mat2SO3(*args, **kwargs) + result.ltype = _wp.warpSO3_type + return result + + +def mat2SE3(*args: Any, **kwargs: Any) -> pypose.LieTensor: + result = pypose.mat2SE3(*args, **kwargs) + result.ltype = _wp.warpSE3_type + return result + + +# ====== Backend conversion utilities ====== + +from .pypose_warp import to_warp_backend, to_pypose_backend, is_warp_backend # noqa: E402 + + +# ====== Fallback delegation to unpatched pypose ====== + + +def __getattr__(name: str) -> Any: + """Delegate unresolved attributes to the unpatched pypose module.""" + return getattr(pypose, name) + + +def __dir__() -> list[str]: + """Expose both local symbols and pypose public attributes.""" + return sorted(set(globals().keys()) | set(dir(pypose))) From 6541c742a07a06dfcfb6ad28a15b1dbcd584b62a Mon Sep 17 00:00:00 2001 From: haleqiu Date: Sat, 14 Mar 2026 23:41:02 -0400 Subject: [PATCH 3/3] update global patch --- __init__.py | 180 ++++++++++++++++++++++------------------ generic.py | 10 +++ pypose_warp/__init__.py | 6 +- 3 files changed, 113 insertions(+), 83 deletions(-) create mode 100644 generic.py diff --git a/__init__.py b/__init__.py index 0cd1ccb..da10dbd 100644 --- a/__init__.py +++ b/__init__.py @@ -1,136 +1,158 @@ -"""Warp-accelerated PyPose Lie operations without global monkey-patching. +"""Warp-accelerated PyPose via global monkey-patching. -``import PyposeWarp as pp`` gives warp-backed constructors (so3, SO3, se3, SE3, -identity_*, randn_*, euler2SO3, mat2SO3, mat2SE3, from_matrix). Any attribute -not explicitly defined here (e.g. ``optim``, ``module``, ``vec2skew``) is -delegated to the *unpatched* ``pypose`` module via ``__getattr__``. +``import PyposeWarp`` patches ``pypose`` globally so that every +``import pypose as pp`` yields warp-accelerated Lie constructors. + +For code that needs the original (unpatched) constructors — e.g. the EKF +Jacobian pipeline that relies on ``torch.autograd.functional.jacobian`` — +use ``import PyposeWarp.generic as pp`` instead. """ +import sys +import types from typing import Any, cast +from functools import partial import pypose -from functools import partial from . import pypose_warp _wp = cast(Any, pypose_warp) -_orig_SE3_type = pypose.SE3_type -_orig_SO3_type = pypose.SO3_type +# ====== Save originals before patching ====== -# ====== Warp-backed type objects ====== +_PATCHED_ATTRS = [ + "so3", "SO3", "se3", "SE3", + "SO3_type", "so3_type", "SE3_type", "se3_type", + "identity_SO3", "identity_SE3", "identity_so3", "identity_se3", + "randn_SO3", "randn_SE3", "randn_so3", "randn_se3", + "from_matrix", "euler2SO3", "mat2SO3", "mat2SE3", +] -SE3_type = _wp.warpSE3_type -SO3_type = _wp.warpSO3_type -se3_type = _wp.warpse3_type -so3_type = _wp.warpso3_type +_orig: dict[str, Any] = {name: getattr(pypose, name) for name in _PATCHED_ATTRS} -# ====== Warp-backed constructors ====== +# ====== Build PyposeWarp.generic module ====== -SO3 = partial(pypose.LieTensor, ltype=_wp.warpSO3_type) -so3 = partial(pypose.LieTensor, ltype=_wp.warpso3_type) -SE3 = partial(pypose.LieTensor, ltype=_wp.warpSE3_type) -se3 = partial(pypose.LieTensor, ltype=_wp.warpse3_type) +generic = types.ModuleType("PyposeWarp.generic") +generic.__doc__ = ( + "Original (unpatched) PyPose constructors. Use this in files that " + "need torch.autograd.functional.jacobian compatibility." +) +generic.__dict__.update(_orig) +generic.__package__ = "PyposeWarp" -LieTensor = pypose.LieTensor -# ====== Warp-backed identity functions ====== +def _generic_getattr(name: str) -> Any: + """Delegate non-patched attributes to the real pypose module.""" + return getattr(pypose, name) -def identity_SE3(*args: Any, **kwargs: Any) -> pypose.LieTensor: - result = pypose.identity_SE3(*args, **kwargs) - result.ltype = _wp.warpSE3_type - return result +generic.__getattr__ = _generic_getattr # type: ignore[attr-defined] +sys.modules["PyposeWarp.generic"] = generic +# ====== Warp-backed wrapper functions ====== -def identity_SO3(*args: Any, **kwargs: Any) -> pypose.LieTensor: - result = pypose.identity_SO3(*args, **kwargs) - result.ltype = _wp.warpSO3_type - return result +def _wp_identity(orig_fn: Any, warp_ltype: Any) -> Any: + def wrapper(*args: Any, **kwargs: Any) -> pypose.LieTensor: + result = orig_fn(*args, **kwargs) + result.ltype = warp_ltype + return result + return wrapper -def identity_se3(*args: Any, **kwargs: Any) -> pypose.LieTensor: - result = pypose.identity_se3(*args, **kwargs) - result.ltype = _wp.warpse3_type - return result - -def identity_so3(*args: Any, **kwargs: Any) -> pypose.LieTensor: - result = pypose.identity_so3(*args, **kwargs) - result.ltype = _wp.warpso3_type - return result - - -# ====== Warp-backed randn functions ====== +def _wp_randn(orig_fn: Any, warp_ltype: Any) -> Any: + def wrapper(*args: Any, **kwargs: Any) -> pypose.LieTensor: + result = orig_fn(*args, **kwargs) + result.ltype = warp_ltype + return result + return wrapper -def randn_SE3(*args: Any, **kwargs: Any) -> pypose.LieTensor: - result = pypose.randn_SE3(*args, **kwargs) - result.ltype = _wp.warpSE3_type +def wp_from_matrix(mat: Any, ltype: Any, check: bool = True, + rtol: float = 1e-5, atol: float = 1e-5) -> pypose.LieTensor: + if ltype is _wp.warpSE3_type: + result = _orig["from_matrix"](mat, _orig["SE3_type"], check=check, rtol=rtol, atol=atol) + result.ltype = _wp.warpSE3_type + elif ltype is _wp.warpSO3_type: + result = _orig["from_matrix"](mat, _orig["SO3_type"], check=check, rtol=rtol, atol=atol) + result.ltype = _wp.warpSO3_type + else: + result = _orig["from_matrix"](mat, ltype, check=check, rtol=rtol, atol=atol) return result -def randn_SO3(*args: Any, **kwargs: Any) -> pypose.LieTensor: - result = pypose.randn_SO3(*args, **kwargs) +def wp_euler2SO3(*args: Any, **kwargs: Any) -> pypose.LieTensor: + result = _orig["euler2SO3"](*args, **kwargs) result.ltype = _wp.warpSO3_type return result -def randn_se3(*args: Any, **kwargs: Any) -> pypose.LieTensor: - result = pypose.randn_se3(*args, **kwargs) - result.ltype = _wp.warpse3_type +def wp_mat2SO3(*args: Any, **kwargs: Any) -> pypose.LieTensor: + result = _orig["mat2SO3"](*args, **kwargs) + result.ltype = _wp.warpSO3_type return result -def randn_so3(*args: Any, **kwargs: Any) -> pypose.LieTensor: - result = pypose.randn_so3(*args, **kwargs) - result.ltype = _wp.warpso3_type +def wp_mat2SE3(*args: Any, **kwargs: Any) -> pypose.LieTensor: + result = _orig["mat2SE3"](*args, **kwargs) + result.ltype = _wp.warpSE3_type return result -# ====== Warp-backed conversion functions ====== +# ====== Global monkey-patching ====== +# Type objects +pypose.SE3_type = _wp.warpSE3_type +pypose.SO3_type = _wp.warpSO3_type +pypose.se3_type = _wp.warpse3_type +pypose.so3_type = _wp.warpso3_type -def from_matrix(mat: Any, ltype: Any, check: bool = True, - rtol: float = 1e-5, atol: float = 1e-5) -> pypose.LieTensor: - if ltype is _wp.warpSE3_type: - result = pypose.from_matrix(mat, _orig_SE3_type, check=check, rtol=rtol, atol=atol) - result.ltype = _wp.warpSE3_type - elif ltype is _wp.warpSO3_type: - result = pypose.from_matrix(mat, _orig_SO3_type, check=check, rtol=rtol, atol=atol) - result.ltype = _wp.warpSO3_type - else: - result = pypose.from_matrix(mat, ltype, check=check, rtol=rtol, atol=atol) - return result - +# Constructors +pypose.SO3 = partial(pypose.LieTensor, ltype=_wp.warpSO3_type) +pypose.so3 = partial(pypose.LieTensor, ltype=_wp.warpso3_type) +pypose.SE3 = partial(pypose.LieTensor, ltype=_wp.warpSE3_type) +pypose.se3 = partial(pypose.LieTensor, ltype=_wp.warpse3_type) -def euler2SO3(*args: Any, **kwargs: Any) -> pypose.LieTensor: - result = pypose.euler2SO3(*args, **kwargs) - result.ltype = _wp.warpSO3_type - return result +# Identity functions +pypose.identity_SE3 = _wp_identity(_orig["identity_SE3"], _wp.warpSE3_type) +pypose.identity_SO3 = _wp_identity(_orig["identity_SO3"], _wp.warpSO3_type) +pypose.identity_se3 = _wp_identity(_orig["identity_se3"], _wp.warpse3_type) +pypose.identity_so3 = _wp_identity(_orig["identity_so3"], _wp.warpso3_type) +# Randn functions +pypose.randn_SE3 = _wp_randn(_orig["randn_SE3"], _wp.warpSE3_type) +pypose.randn_SO3 = _wp_randn(_orig["randn_SO3"], _wp.warpSO3_type) +pypose.randn_se3 = _wp_randn(_orig["randn_se3"], _wp.warpse3_type) +pypose.randn_so3 = _wp_randn(_orig["randn_so3"], _wp.warpso3_type) -def mat2SO3(*args: Any, **kwargs: Any) -> pypose.LieTensor: - result = pypose.mat2SO3(*args, **kwargs) - result.ltype = _wp.warpSO3_type - return result +# Conversion functions +pypose.from_matrix = wp_from_matrix +pypose.euler2SO3 = wp_euler2SO3 +pypose.mat2SO3 = wp_mat2SO3 +pypose.mat2SE3 = wp_mat2SE3 +# ====== Patch pypose internal module bindings ====== +# These modules bind `so3`, `identity_SO3` etc. at import time via +# ``from .. import so3, identity_SO3``, so we must patch those bindings +# directly after global patching. -def mat2SE3(*args: Any, **kwargs: Any) -> pypose.LieTensor: - result = pypose.mat2SE3(*args, **kwargs) - result.ltype = _wp.warpSE3_type - return result +import pypose.module.imu_preintegrator as _preint # noqa: E402 +_preint.so3 = pypose.so3 +_preint.SO3 = pypose.SO3 +_preint.identity_SO3 = pypose.identity_SO3 +_preint.LieTensor = pypose.LieTensor -# ====== Backend conversion utilities ====== +# ====== Re-exports ====== from .pypose_warp import to_warp_backend, to_pypose_backend, is_warp_backend # noqa: E402 -# ====== Fallback delegation to unpatched pypose ====== +# ====== Module-level delegation ====== def __getattr__(name: str) -> Any: - """Delegate unresolved attributes to the unpatched pypose module.""" + """Delegate unresolved attributes to the (now patched) pypose module.""" return getattr(pypose, name) diff --git a/generic.py b/generic.py new file mode 100644 index 0000000..ef6e4a1 --- /dev/null +++ b/generic.py @@ -0,0 +1,10 @@ +"""Original (unpatched) PyPose constructors for autograd-compatible code. + +This module is populated by ``PyposeWarp.__init__`` before monkey-patching. +Use ``import PyposeWarp.generic as pp`` in files that need +``torch.autograd.functional.jacobian`` compatibility (e.g. EKF NLS pipeline). + +Attributes not explicitly saved (e.g. ``module``, ``optim``, ``vec2skew``) +are delegated to the real ``pypose`` module via ``__getattr__``; these are +unaffected by monkey-patching so they work identically either way. +""" diff --git a/pypose_warp/__init__.py b/pypose_warp/__init__.py index 2968ce0..2bfede8 100644 --- a/pypose_warp/__init__.py +++ b/pypose_warp/__init__.py @@ -19,6 +19,7 @@ ] _PP_TO_WP = {pp_ltype : wp_ltype for pp_ltype, wp_ltype in _BACKEND_LIST} _WP_TO_PP = {wp_ltype : pp_ltype for pp_ltype, wp_ltype in _BACKEND_LIST} +_ORIG_PP_TYPES = frozenset(pp_ltype for pp_ltype, _ in _BACKEND_LIST) def to_warp_backend(x: pp.LieTensor) -> pp.LieTensor: @@ -46,7 +47,4 @@ def is_warp_backend(x: pp.LieTensor) -> bool: def is_pypose_backend(x: pp.LieTensor) -> bool: - return x.ltype in { - pp.SE3_type, pp.SO3_type, pp.RxSO3_type, pp.Sim3_type, - pp.se3_type, pp.so3_type, pp.rxso3_type, pp.sim3_type - } + return x.ltype in _ORIG_PP_TYPES