adding a big feature for the plugin registry, making plugins usable

leanpass/__init__.py

@@ -1,5 +1,6 @@
 from .tensor import Tensor
 from . import nn
 from . import optim
+from . import ops
 
-__all__ = ["Tensor", "nn", "optim"]
+__all__ = ["Tensor", "nn", "optim", "ops"]

leanpass/ops.py

@@ -0,0 +1,8 @@
+from .plugin_registry import get_op
+
+def __getattr__(name: str):
+    try:
+        return get_op(name)
+    except KeyError:
+        raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
+

leanpass/plugin_registry.py

@@ -0,0 +1,24 @@
+from typing import Callable, Any, Dict
+
+# The central registry holding all ops (both built-in and plugins)
+_REGISTRY: Dict[str, Callable] = {}
+
+def register_op(name: str, fn: Callable):
+    """Registers a function into LeanPass."""
+    if name in _REGISTRY:
+        pass
+    _REGISTRY[name] = fn
+
+def op(name: str):
+    """Decorator for registering ops."""
+    def decorator(fn):
+        register_op(name, fn)
+        return fn
+    return decorator
+
+def get_op(name: str) -> Callable:
+    """Retrieve an op by name."""
+    if name not in _REGISTRY:
+        raise KeyError(f"Op '{name}' not found. Did you install its plugin?")
+    return _REGISTRY[name]
+