kiwi-array 0.2.47__py3-none-any.whl

kiwi_array/__init__.py

@@ -0,0 +1,9 @@
+from __future__ import annotations
+
+__version__ = "0.2.47"
+
+
+def load_ipython_extension(ipython):
+    from .notebook import load_ipython_extension as load
+
+    load(ipython)

kiwi_array/bridge.py

@@ -0,0 +1,801 @@
+from __future__ import annotations
+
+import ctypes
+import ctypes.util
+import os
+import platform
+import shutil
+import shlex
+import subprocess
+from dataclasses import dataclass
+from importlib import metadata as importlib_metadata
+from pathlib import Path
+from typing import Any, Optional, Sequence, Union
+
+try:
+    import numpy as _np
+except ImportError:  # pragma: no cover - optional dependency
+    _np = None
+
+LIB_ENV_VAR = "KIWI_BRIDGE_LIB"
+LEGACY_LIB_ENV_VAR = "KIWI_JUPYTER_BRIDGE_LIB"
+RUNTIME_ENV_VAR = "KIWI_RUNTIME"
+CUDA_DRIVER_LIB_ENV_VAR = "KIWI_CUDA_DRIVER_LIB"
+CUDA_PRELOAD_DRIVER_ENV_VAR = "KIWI_CUDA_PRELOAD_DRIVER"
+SOURCE_ROOT_ENV_VAR = "KIWI_SOURCE_ROOT"
+ZIG_ENV_VAR = "KIWI_ZIG"
+MLX_PREFIX_ENV_VAR = "KIWI_MLX_PREFIX"
+MLX_C_INCLUDE_ENV_VAR = "KIWI_MLX_C_INCLUDE"
+RUNTIME_ENTRY_POINT_GROUP = "kiwi_array.runtimes"
+
+DEVICE_VALUES = {
+    "auto": 0,
+    "cpu": 1,
+    "gpu": 2,
+}
+
+STATUS_NAMES = {
+    0: "ok",
+    1: "parse",
+    2: "type",
+    3: "name",
+    4: "domain",
+    5: "rank",
+    6: "nyi",
+    7: "length",
+    8: "index",
+    9: "mlx",
+    10: "device",
+    11: "error",
+    12: "oom",
+}
+
+AUTOGRAD_PATH_NAMES = {
+    0: "none",
+    1: "mlx",
+    2: "finite_difference",
+}
+
+PathLikeArg = Union[os.PathLike[str], str]
+
+
+class _CKiwiEvalResult(ctypes.Structure):
+    _fields_ = [
+        ("status", ctypes.c_int),
+        ("echoed", ctypes.c_bool),
+        ("autograd_path", ctypes.c_int),
+        ("text_ptr", ctypes.c_void_p),
+        ("text_len", ctypes.c_size_t),
+        ("display_mime_ptr", ctypes.c_void_p),
+        ("display_mime_len", ctypes.c_size_t),
+        ("display_data_ptr", ctypes.c_void_p),
+        ("display_data_len", ctypes.c_size_t),
+    ]
+
+
+@dataclass(frozen=True)
+class KiwiEvalResult:
+    status: str
+    echoed: bool
+    autograd_path: str
+    text: Optional[str]
+    display_mime: Optional[str] = None
+    display_data: Optional[str] = None
+
+
+class KiwiBridgeError(RuntimeError):
+    pass
+
+
+@dataclass(frozen=True)
+class RuntimeDescriptor:
+    name: str
+    priority: int
+    backend: str
+    accelerator: Optional[str]
+    native_dir: Optional[Path]
+    lib_dir: Optional[Path]
+    bin_dir: Optional[Path] = None
+    source: str = "unknown"
+
+    def bridge_library_path(self) -> Optional[Path]:
+        if self.native_dir is None:
+            return None
+        return self.native_dir / platform_library_name("kiwi_bridge")
+
+    def cli_path(self) -> Optional[Path]:
+        if self.bin_dir is None:
+            return None
+        return self.bin_dir / platform_executable_name("kiwi")
+
+
+def _find_source_root() -> Optional[Path]:
+    override = os.environ.get(SOURCE_ROOT_ENV_VAR)
+    if override:
+        return Path(override).expanduser()
+
+    for parent in Path(__file__).resolve().parents:
+        if (parent / "build.zig").is_file() and (parent / "src" / "kiwi_bridge.zig").is_file():
+            return parent
+    return None
+
+
+def implementation_root() -> Path:
+    root = _find_source_root()
+    if root is None:
+        raise KiwiBridgeError(
+            "Kiwi source root not found. Set KIWI_SOURCE_ROOT when building the "
+            "bridge from source, or install a platform wheel that carries the "
+            "native Kiwi runtime."
+        )
+    return root
+
+
+def _has_workspace_dependency_layout() -> bool:
+    root = _find_source_root()
+    if root is None:
+        return False
+    return (root.parents[1] / "vendor" / "mlx-c").exists()
+
+
+def repo_root() -> Path:
+    root = implementation_root()
+    return root.parents[1] if _has_workspace_dependency_layout() else root
+
+
+def _package_root() -> Optional[Path]:
+    return Path(__file__).resolve().parent
+
+
+def _packaged_native_dir() -> Optional[Path]:
+    root = _package_root()
+    return None if root is None else root / "native"
+
+
+def _packaged_library_dir() -> Optional[Path]:
+    root = _package_root()
+    return None if root is None else root / "lib"
+
+
+def _is_platform_dynamic_library_name(name: str) -> bool:
+    system = platform.system()
+    if system == "Darwin":
+        return name.endswith(".dylib")
+    if system == "Windows":
+        return name.endswith(".dll")
+    return ".so" in name
+
+
+def _dir_has_payload(path: Optional[Path]) -> bool:
+    if path is None or not path.is_dir():
+        return False
+    return any(
+        child.is_file()
+        and child.name not in {".gitignore", ".gitkeep"}
+        and _is_platform_dynamic_library_name(child.name)
+        for child in path.iterdir()
+    )
+
+
+def _path_from_descriptor(value: object) -> Optional[Path]:
+    if value in (None, ""):
+        return None
+    return Path(os.fspath(value)).expanduser()
+
+
+def _coerce_runtime_descriptor(raw: object, source: str) -> RuntimeDescriptor:
+    if isinstance(raw, RuntimeDescriptor):
+        return raw
+    if not isinstance(raw, dict):
+        raise TypeError(f"runtime descriptor from {source} must be a dict")
+    return RuntimeDescriptor(
+        name=str(raw["name"]),
+        priority=int(raw.get("priority", 0)),
+        backend=str(raw.get("backend", "unknown")),
+        accelerator=None if raw.get("accelerator") is None else str(raw.get("accelerator")),
+        native_dir=_path_from_descriptor(raw.get("native_dir")),
+        lib_dir=_path_from_descriptor(raw.get("lib_dir")),
+        bin_dir=_path_from_descriptor(raw.get("bin_dir")),
+        source=source,
+    )
+
+
+def _embedded_runtime_descriptor() -> RuntimeDescriptor:
+    root = _package_root()
+    return RuntimeDescriptor(
+        name="host",
+        priority=10,
+        backend="host",
+        accelerator=None,
+        native_dir=root / "native",
+        lib_dir=root / "lib",
+        bin_dir=root / "bin",
+        source="kiwi_array",
+    )
+
+
+def _entry_point_runtime_descriptors() -> list[RuntimeDescriptor]:
+    try:
+        entry_points = importlib_metadata.entry_points()
+    except Exception:
+        return []
+    if hasattr(entry_points, "select"):
+        selected = entry_points.select(group=RUNTIME_ENTRY_POINT_GROUP)
+    else:  # pragma: no cover - Python <3.10 compatibility path
+        selected = entry_points.get(RUNTIME_ENTRY_POINT_GROUP, [])
+
+    descriptors: list[RuntimeDescriptor] = []
+    for entry_point in selected:
+        try:
+            loaded = entry_point.load()
+            raw = loaded() if callable(loaded) else loaded
+            descriptors.append(_coerce_runtime_descriptor(raw, entry_point.name))
+        except Exception:
+            continue
+    return descriptors
+
+
+def discover_runtime_descriptors() -> list[RuntimeDescriptor]:
+    descriptors = [_embedded_runtime_descriptor(), *_entry_point_runtime_descriptors()]
+    by_name: dict[str, RuntimeDescriptor] = {}
+    for descriptor in descriptors:
+        current = by_name.get(descriptor.name)
+        if current is None or descriptor.priority > current.priority:
+            by_name[descriptor.name] = descriptor
+    return sorted(by_name.values(), key=lambda item: (-item.priority, item.name))
+
+
+def _runtime_selection() -> str:
+    value = os.environ.get(RUNTIME_ENV_VAR, "auto").strip()
+    return value or "auto"
+
+
+def _selected_runtime_descriptors() -> list[RuntimeDescriptor]:
+    selection = _runtime_selection()
+    descriptors = discover_runtime_descriptors()
+    if selection == "auto":
+        return descriptors
+    return [descriptor for descriptor in descriptors if descriptor.name == selection]
+
+
+def active_runtime_descriptor() -> Optional[RuntimeDescriptor]:
+    for descriptor in _selected_runtime_descriptors():
+        bridge = descriptor.bridge_library_path()
+        if bridge is not None and bridge.exists():
+            return descriptor
+    return None
+
+
+def zig_executable() -> str:
+    override = os.environ.get(ZIG_ENV_VAR) or os.environ.get("ZIG")
+    if override:
+        return str(Path(override).expanduser())
+    bundled = repo_root() / "tools" / "zig"
+    if bundled.exists():
+        return str(bundled)
+    path_zig = shutil.which("zig")
+    return path_zig if path_zig is not None else "zig"
+
+
+def default_mlx_prefix() -> Path:
+    if _has_workspace_dependency_layout():
+        return repo_root() / ".artifacts" / "mlx" / "macos-default-install"
+    return implementation_root() / ".deps" / "mlx"
+
+
+def default_mlx_c_include() -> Path:
+    if _has_workspace_dependency_layout():
+        return repo_root() / "vendor" / "mlx-c"
+    return implementation_root() / ".deps" / "mlx-c"
+
+
+def mlx_prefix() -> Path:
+    return Path(
+        os.environ.get(MLX_PREFIX_ENV_VAR)
+        or default_mlx_prefix()
+    ).expanduser()
+
+
+def mlx_c_include() -> Path:
+    return Path(
+        os.environ.get(MLX_C_INCLUDE_ENV_VAR)
+        or default_mlx_c_include()
+    ).expanduser()
+
+
+def runtime_library_search_dir() -> Path:
+    source_root = _find_source_root()
+    if _runtime_selection() == "auto" and source_root is not None:
+        source_bridge = source_root / "zig-out" / "lib" / platform_library_name("kiwi_bridge")
+        if source_bridge.exists():
+            return mlx_prefix() / "lib"
+
+    descriptor = active_runtime_descriptor()
+    if descriptor is not None and descriptor.lib_dir is not None and _dir_has_payload(descriptor.lib_dir):
+        return descriptor.lib_dir
+    return mlx_prefix() / "lib"
+
+
+def runtime_library_env_var() -> Optional[str]:
+    system = platform.system()
+    if system == "Darwin":
+        return "DYLD_LIBRARY_PATH"
+    if system == "Linux":
+        return "LD_LIBRARY_PATH"
+    return None
+
+
+def build_args(optimize: str) -> list[str]:
+    bridge_mlx_prefix = mlx_prefix()
+    bridge_mlx_c_include = mlx_c_include()
+    return [
+        zig_executable(),
+        "build",
+        f"-Doptimize={optimize}",
+        "-Dpublic-cli=true",
+        "-Dcli-name=kiwi",
+        "-Dinstall-sdk=true",
+        f"-Dmlx-prefix={bridge_mlx_prefix}",
+        f"-Dmlx-c-include={bridge_mlx_c_include}",
+    ]
+
+
+def platform_library_name(base_name: str) -> str:
+    system = platform.system()
+    if system == "Darwin":
+        return f"lib{base_name}.dylib"
+    if system == "Windows":
+        return f"{base_name}.dll"
+    return f"lib{base_name}.so"
+
+
+def platform_executable_name(base_name: str) -> str:
+    return f"{base_name}.exe" if platform.system() == "Windows" else base_name
+
+
+def default_cli_candidates() -> list[Path]:
+    candidates: list[Path] = []
+    source_root = _find_source_root()
+    source_cli = None
+    if source_root is not None:
+        source_cli = source_root / "zig-out" / "bin" / platform_executable_name("kiwi")
+
+    if _runtime_selection() == "auto" and source_cli is not None:
+        candidates.append(source_cli)
+
+    for descriptor in _selected_runtime_descriptors():
+        cli = descriptor.cli_path()
+        if cli is not None:
+            candidates.append(cli)
+
+    if _runtime_selection() != "auto" and source_cli is not None:
+        candidates.append(source_cli)
+    return candidates
+
+
+def find_cli_path() -> Path:
+    for candidate in default_cli_candidates():
+        if candidate.exists():
+            return candidate
+    joined = ", ".join(str(candidate) for candidate in default_cli_candidates())
+    raise FileNotFoundError(
+        "Kiwi CLI not found. Install a runtime payload such as "
+        "`kiwi-array-host`, `kiwi-array-cpu`, `kiwi-array-metal`, or "
+        "`kiwi-array-cuda12`. Searched: "
+        f"{joined}"
+    )
+
+
+def default_library_candidates() -> list[Path]:
+    candidates: list[Path] = []
+    for env_var in (LIB_ENV_VAR, LEGACY_LIB_ENV_VAR):
+        env_path = os.environ.get(env_var)
+        if env_path:
+            candidates.append(Path(env_path).expanduser())
+
+    source_root = _find_source_root()
+    source_bridge = None
+    if source_root is not None:
+        source_bridge = source_root / "zig-out" / "lib" / platform_library_name("kiwi_bridge")
+
+    if _runtime_selection() == "auto" and source_bridge is not None:
+        candidates.append(source_bridge)
+
+    for descriptor in _selected_runtime_descriptors():
+        bridge = descriptor.bridge_library_path()
+        if bridge is not None:
+            candidates.append(bridge)
+
+    if _runtime_selection() != "auto" and source_bridge is not None:
+        candidates.append(source_bridge)
+    return candidates
+
+
+def find_library_path(path: Optional[PathLikeArg] = None) -> Path:
+    if path is not None:
+        candidate = Path(path).expanduser()
+        if candidate.exists():
+            return candidate
+        raise FileNotFoundError(f"Kiwi bridge library not found at {candidate}")
+
+    for candidate in default_library_candidates():
+        if candidate.exists():
+            return candidate
+    joined = ", ".join(str(candidate) for candidate in default_library_candidates())
+    source_root = _find_source_root()
+    if source_root is None:
+        hint = "Install a platform wheel that carries the native Kiwi runtime."
+    else:
+        hint = f"Build it with `{shlex.join(build_args('ReleaseFast'))}` from {source_root}."
+    raise FileNotFoundError(f"Kiwi bridge library not found. {hint} Searched: {joined}")
+
+
+def build_bridge(optimize: str = "ReleaseFast") -> None:
+    subprocess.run(
+        build_args(optimize),
+        cwd=implementation_root(),
+        check=True,
+    )
+
+
+def _configure_library(lib: ctypes.CDLL) -> ctypes.CDLL:
+    lib.kiwi_session_create.argtypes = [ctypes.c_int]
+    lib.kiwi_session_create.restype = ctypes.c_void_p
+
+    lib.kiwi_session_destroy.argtypes = [ctypes.c_void_p]
+    lib.kiwi_session_destroy.restype = None
+
+    lib.kiwi_session_eval.argtypes = [ctypes.c_void_p, ctypes.c_char_p, ctypes.c_size_t]
+    lib.kiwi_session_eval.restype = _CKiwiEvalResult
+
+    lib.kiwi_session_set_global_float_array.argtypes = [
+        ctypes.c_void_p,
+        ctypes.c_char_p,
+        ctypes.c_size_t,
+        ctypes.POINTER(ctypes.c_float),
+        ctypes.POINTER(ctypes.c_int32),
+        ctypes.c_size_t,
+    ]
+    lib.kiwi_session_set_global_float_array.restype = ctypes.c_int
+
+    if hasattr(lib, "kiwi_session_set_global_mlx_float_array"):
+        lib.kiwi_session_set_global_mlx_float_array.argtypes = [
+            ctypes.c_void_p,
+            ctypes.c_char_p,
+            ctypes.c_size_t,
+            ctypes.POINTER(ctypes.c_float),
+            ctypes.POINTER(ctypes.c_int32),
+            ctypes.c_size_t,
+        ]
+        lib.kiwi_session_set_global_mlx_float_array.restype = ctypes.c_int
+
+    lib.kiwi_session_set_global_int_array.argtypes = [
+        ctypes.c_void_p,
+        ctypes.c_char_p,
+        ctypes.c_size_t,
+        ctypes.POINTER(ctypes.c_int32),
+        ctypes.POINTER(ctypes.c_int32),
+        ctypes.c_size_t,
+    ]
+    lib.kiwi_session_set_global_int_array.restype = ctypes.c_int
+
+    lib.kiwi_session_set_global_bool_array.argtypes = [
+        ctypes.c_void_p,
+        ctypes.c_char_p,
+        ctypes.c_size_t,
+        ctypes.POINTER(ctypes.c_bool),
+        ctypes.POINTER(ctypes.c_int32),
+        ctypes.c_size_t,
+    ]
+    lib.kiwi_session_set_global_bool_array.restype = ctypes.c_int
+
+    lib.kiwi_eval_result_free.argtypes = [_CKiwiEvalResult]
+    lib.kiwi_eval_result_free.restype = None
+
+    lib.kiwi_status_name.argtypes = [ctypes.c_int]
+    lib.kiwi_status_name.restype = ctypes.c_char_p
+    return lib
+
+
+def _preload_runtime_dependencies(descriptor: Optional[RuntimeDescriptor] = None) -> None:
+    _preload_cuda_driver()
+
+    search_dirs = []
+    if descriptor is not None and descriptor.lib_dir is not None:
+        search_dirs.append(descriptor.lib_dir)
+    search_dirs.append(runtime_library_search_dir())
+    packaged_lib = _packaged_library_dir()
+    if packaged_lib is not None:
+        search_dirs.append(packaged_lib)
+    source_root = _find_source_root()
+    if source_root is not None:
+        search_dirs.append(source_root / "zig-out" / "lib")
+
+    seen: set[Path] = set()
+    for directory in search_dirs:
+        if directory in seen:
+            continue
+        seen.add(directory)
+        for base_name in ("mlx", "duckdb"):
+            library = directory / platform_library_name(base_name)
+            if library.exists():
+                ctypes.CDLL(str(library), mode=getattr(ctypes, "RTLD_GLOBAL", 0))
+
+
+def _preload_cuda_driver() -> None:
+    if platform.system() != "Linux":
+        return
+
+    preload_mode = os.environ.get(CUDA_PRELOAD_DRIVER_ENV_VAR, "0").strip().lower()
+    if preload_mode in {"", "0", "false", "no", "off"}:
+        return
+
+    candidates = []
+    override = os.environ.get(CUDA_DRIVER_LIB_ENV_VAR)
+    if override:
+        candidates.append(Path(override).expanduser())
+
+    found = ctypes.util.find_library("cuda")
+    if found:
+        candidates.append(Path(found))
+
+    candidates.extend(
+        [
+            Path("/usr/lib64-nvidia/libcuda.so.1"),
+            Path("/usr/local/nvidia/lib64/libcuda.so.1"),
+            Path("/usr/lib/x86_64-linux-gnu/libcuda.so.1"),
+        ]
+    )
+    if preload_mode in {"compat", "cuda-compat"}:
+        candidates.append(Path("/usr/local/cuda/compat/libcuda.so.1"))
+
+    for candidate in candidates:
+        try:
+            if candidate.is_absolute() and not candidate.exists():
+                continue
+            ctypes.CDLL(str(candidate), mode=getattr(ctypes, "RTLD_GLOBAL", 0))
+            return
+        except OSError:
+            continue
+
+
+def _load_library_at(path: Path, descriptor: Optional[RuntimeDescriptor] = None) -> ctypes.CDLL:
+    _preload_runtime_dependencies(descriptor)
+    return _configure_library(ctypes.CDLL(str(path)))
+
+
+def load_library(path: Optional[PathLikeArg] = None) -> ctypes.CDLL:
+    if path is not None or os.environ.get(LIB_ENV_VAR) or os.environ.get(LEGACY_LIB_ENV_VAR):
+        return _load_library_at(find_library_path(path))
+
+    errors: list[str] = []
+    selection = _runtime_selection()
+    source_root = _find_source_root()
+    if selection == "auto" and source_root is not None:
+        source_bridge = source_root / "zig-out" / "lib" / platform_library_name("kiwi_bridge")
+        if source_bridge.exists():
+            try:
+                return _load_library_at(source_bridge)
+            except OSError as exc:
+                errors.append(f"source: {exc}")
+
+    for descriptor in _selected_runtime_descriptors():
+        bridge = descriptor.bridge_library_path()
+        if bridge is None or not bridge.exists():
+            continue
+        try:
+            return _load_library_at(bridge, descriptor)
+        except OSError as exc:
+            errors.append(f"{descriptor.name}: {exc}")
+            if selection != "auto":
+                raise KiwiBridgeError(f"failed to load Kiwi runtime {descriptor.name!r}: {exc}") from exc
+
+    if selection != "auto" and source_root is not None:
+        source_bridge = source_root / "zig-out" / "lib" / platform_library_name("kiwi_bridge")
+        if source_bridge.exists():
+            try:
+                return _load_library_at(source_bridge)
+            except OSError as exc:
+                errors.append(f"source: {exc}")
+
+    if errors:
+        raise KiwiBridgeError("failed to load any Kiwi runtime: " + "; ".join(errors))
+    return _load_library_at(find_library_path())
+
+
+def _normalize_dims(dims: Optional[Sequence[int]], data) -> tuple[int, ...]:
+    if dims is not None:
+        return tuple(int(dim) for dim in dims)
+
+    shape = getattr(data, "shape", None)
+    if shape is None:
+        raise TypeError("dims are required unless data exposes a shape")
+    return tuple(int(dim) for dim in shape)
+
+
+def _product(dims: Sequence[int]) -> int:
+    total = 1
+    for dim in dims:
+        if dim < 0:
+            raise ValueError(f"negative dimension {dim}")
+        total *= dim
+    return total
+
+
+def _sequence_values(data, count: int):
+    if count == 1 and not isinstance(data, (list, tuple)):
+        return [data]
+    values = list(data)
+    if len(values) != count:
+        raise ValueError(f"expected {count} items, got {len(values)}")
+    return values
+
+
+def _prepare_array_data(data, dims: tuple[int, ...], c_type, numpy_dtype):
+    count = _product(dims)
+    if _np is not None:
+        arr = _np.asarray(data, dtype=numpy_dtype)
+        if arr.size != count:
+            raise ValueError(f"expected {count} items, got {arr.size}")
+        flat = _np.ascontiguousarray(arr.reshape(-1))
+        ptr = None if count == 0 else flat.ctypes.data_as(ctypes.POINTER(c_type))
+        return flat, ptr
+
+    values = _sequence_values(data, count)
+    buf = (c_type * count)(*values)
+    ptr = None if count == 0 else buf
+    return buf, ptr
+
+
+class KiwiSession:
+    def __init__(self, device: str = "auto", library_path: Optional[PathLikeArg] = None) -> None:
+        try:
+            device_value = DEVICE_VALUES[device]
+        except KeyError as exc:
+            raise ValueError(f"unknown Kiwi device {device!r}") from exc
+
+        self._lib = load_library(library_path)
+        self._handle = self._lib.kiwi_session_create(device_value)
+        if not self._handle:
+            raise KiwiBridgeError("failed to create Kiwi session")
+
+    def close(self) -> None:
+        if self._handle:
+            self._lib.kiwi_session_destroy(self._handle)
+            self._handle = None
+
+    def eval(self, source: str) -> KiwiEvalResult:
+        self._ensure_open()
+        encoded = source.encode("utf-8")
+        result = self._lib.kiwi_session_eval(self._handle, encoded, len(encoded))
+        try:
+            text = None
+            if result.text_ptr:
+                text = ctypes.string_at(result.text_ptr, result.text_len).decode("utf-8")
+            display_mime = None
+            if result.display_mime_ptr:
+                display_mime = ctypes.string_at(
+                    result.display_mime_ptr,
+                    result.display_mime_len,
+                ).decode("utf-8")
+            display_data = None
+            if result.display_data_ptr:
+                display_data = ctypes.string_at(
+                    result.display_data_ptr,
+                    result.display_data_len,
+                ).decode("utf-8")
+            return KiwiEvalResult(
+                status=STATUS_NAMES.get(result.status, "error"),
+                echoed=bool(result.echoed),
+                autograd_path=AUTOGRAD_PATH_NAMES.get(result.autograd_path, "none"),
+                text=text,
+                display_mime=display_mime,
+                display_data=display_data,
+            )
+        finally:
+            self._lib.kiwi_eval_result_free(result)
+
+    def set_global_float_array(self, name: str, data, dims: Optional[Sequence[int]] = None) -> None:
+        self._set_global_array(
+            "kiwi_session_set_global_float_array",
+            name,
+            data,
+            dims,
+            ctypes.c_float,
+            _np.float32 if _np is not None else None,
+        )
+
+    def set_global_mlx_float_array(self, name: str, data, dims: Optional[Sequence[int]] = None) -> None:
+        self._set_global_array(
+            "kiwi_session_set_global_mlx_float_array",
+            name,
+            data,
+            dims,
+            ctypes.c_float,
+            _np.float32 if _np is not None else None,
+        )
+
+    def set_global_int_array(self, name: str, data, dims: Optional[Sequence[int]] = None) -> None:
+        self._set_global_array(
+            "kiwi_session_set_global_int_array",
+            name,
+            data,
+            dims,
+            ctypes.c_int32,
+            _np.int32 if _np is not None else None,
+        )
+
+    def set_global_bool_array(self, name: str, data, dims: Optional[Sequence[int]] = None) -> None:
+        self._set_global_array(
+            "kiwi_session_set_global_bool_array",
+            name,
+            data,
+            dims,
+            ctypes.c_bool,
+            _np.bool_ if _np is not None else None,
+        )
+
+    def __enter__(self) -> "KiwiSession":
+        return self
+
+    def __exit__(self, exc_type, exc, tb) -> None:
+        self.close()
+
+    def _ensure_open(self) -> None:
+        if not self._handle:
+            raise KiwiBridgeError("Kiwi session is closed")
+
+    def _set_global_array(
+        self,
+        fn_name: str,
+        name: str,
+        data,
+        dims: Optional[Sequence[int]],
+        c_type,
+        numpy_dtype,
+    ) -> None:
+        self._ensure_open()
+        if not hasattr(self._lib, fn_name):
+            raise KiwiBridgeError(f"native Kiwi bridge does not expose {fn_name}")
+        dims_tuple = _normalize_dims(dims, data)
+        owner, data_ptr = _prepare_array_data(data, dims_tuple, c_type, numpy_dtype)
+        dims_owner = (ctypes.c_int32 * len(dims_tuple))(*dims_tuple)
+        dims_ptr = None if len(dims_tuple) == 0 else dims_owner
+        encoded = name.encode("utf-8")
+        status = getattr(self._lib, fn_name)(
+            self._handle,
+            encoded,
+            len(encoded),
+            data_ptr,
+            dims_ptr,
+            len(dims_tuple),
+        )
+        _ = owner
+        if status != 0:
+            raise KiwiBridgeError(f"failed to set global {name!r}: {STATUS_NAMES.get(status, 'error')}")
+
+
+__all__ = [
+    "AUTOGRAD_PATH_NAMES",
+    "KiwiBridgeError",
+    "KiwiEvalResult",
+    "KiwiSession",
+    "LEGACY_LIB_ENV_VAR",
+    "LIB_ENV_VAR",
+    "RUNTIME_ENV_VAR",
+    "RuntimeDescriptor",
+    "STATUS_NAMES",
+    "active_runtime_descriptor",
+    "build_bridge",
+    "default_library_candidates",
+    "default_cli_candidates",
+    "discover_runtime_descriptors",
+    "find_cli_path",
+    "find_library_path",
+    "implementation_root",
+    "load_library",
+    "platform_library_name",
+    "platform_executable_name",
+    "runtime_library_env_var",
+    "runtime_library_search_dir",
+]

kiwi_array/cli.py

@@ -0,0 +1,36 @@
+from __future__ import annotations
+
+import os
+import subprocess
+import sys
+
+from . import bridge
+
+
+def _launcher_environment() -> dict[str, str]:
+    env = os.environ.copy()
+    library_var = bridge.runtime_library_env_var()
+    if library_var is None:
+        return env
+    library_dir = bridge.runtime_library_search_dir()
+    existing = env.get(library_var)
+    env[library_var] = str(library_dir) if not existing else f"{library_dir}{os.pathsep}{existing}"
+    return env
+
+
+def main() -> int:
+    try:
+        cli = bridge.find_cli_path()
+    except FileNotFoundError as exc:
+        print(exc, file=sys.stderr)
+        return 1
+
+    argv = [str(cli), *sys.argv[1:]]
+    if hasattr(os, "execve"):
+        os.execve(str(cli), argv, _launcher_environment())
+    completed = subprocess.run(argv, env=_launcher_environment())
+    return int(completed.returncode)
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

kiwi_array/notebook.py

@@ -0,0 +1,268 @@
+from __future__ import annotations
+
+import json
+import os
+import platform
+import subprocess
+import sys
+import uuid
+from dataclasses import dataclass
+from typing import Iterator, Protocol
+
+from kiwi_array.bridge import KiwiEvalResult, KiwiSession, find_library_path, runtime_library_search_dir
+
+from . import __version__
+
+
+class SessionProtocol(Protocol):
+    def eval(self, source: str) -> KiwiEvalResult:
+        ...
+
+    def close(self) -> None:
+        ...
+
+
+@dataclass(frozen=True)
+class CellOutput:
+    line_no: int
+    text: str
+    autograd_path: str
+    display_mime: str | None = None
+    display_data: str | None = None
+
+
+@dataclass(frozen=True)
+class SmokeCheck:
+    name: str
+    ok: bool
+    detail: str
+
+
+class KiwiNotebookError(RuntimeError):
+    def __init__(self, line_no: int, status: str) -> None:
+        self.line_no = line_no
+        self.status = status
+        super().__init__(f"line {line_no}: !{status}")
+
+
+def iter_executable_lines(code: str) -> Iterator[tuple[int, str]]:
+    for line_no, raw_line in enumerate(code.splitlines(), start=1):
+        line = raw_line.rstrip("\r")
+        trimmed = line.strip(" \t")
+        if not trimmed or trimmed.startswith("/"):
+            continue
+        yield line_no, line
+
+
+def execute_cell(session: SessionProtocol, code: str) -> list[CellOutput]:
+    outputs: list[CellOutput] = []
+    for line_no, line in iter_executable_lines(code):
+        result = session.eval(line)
+        if result.status != "ok":
+            raise KiwiNotebookError(line_no, result.status)
+        if result.echoed and result.text is not None:
+            outputs.append(
+                CellOutput(
+                    line_no=line_no,
+                    text=result.text,
+                    autograd_path=result.autograd_path,
+                    display_mime=result.display_mime,
+                    display_data=result.display_data,
+                )
+            )
+    return outputs
+
+
+def display_bundle_for_output(output: CellOutput) -> dict[str, object]:
+    bundle: dict[str, object] = {"text/plain": output.text}
+    if output.display_mime is None or output.display_data is None:
+        return bundle
+    try:
+        payload = json.loads(output.display_data)
+    except json.JSONDecodeError:
+        return bundle
+    bundle[output.display_mime] = payload
+    if "vegalite" in output.display_mime:
+        bundle["text/html"] = vegalite_html(payload)
+    return bundle
+
+
+def vegalite_embed_options() -> dict[str, object]:
+    options: dict[str, object] = {"actions": False, "renderer": "canvas"}
+    theme = os.environ.get("KIWI_VEGALITE_THEME", "").strip()
+    if theme and theme.lower() not in {"default", "none"}:
+        options["theme"] = theme
+    return options
+
+
+def vegalite_html(spec: object) -> str:
+    element_id = f"kiwi-vegalite-{uuid.uuid4().hex}"
+    spec_json = json.dumps(spec, separators=(",", ":"))
+    options_json = json.dumps(vegalite_embed_options(), separators=(",", ":"))
+    return f"""
+<div id="{element_id}"></div>
+<script src="https://cdn.jsdelivr.net/npm/vega@5"></script>
+<script src="https://cdn.jsdelivr.net/npm/vega-lite@5"></script>
+<script src="https://cdn.jsdelivr.net/npm/vega-embed@6"></script>
+<script>
+(function() {{
+  const spec = {spec_json};
+  const render = function() {{
+    if (!window.vegaEmbed) {{
+      document.getElementById("{element_id}").textContent = {json.dumps("Vega-Embed failed to load.")};
+      return;
+    }}
+    window.vegaEmbed("#{element_id}", spec, {options_json});
+  }};
+  render();
+}})();
+</script>
+"""
+
+
+def _safe_value(label: str, value_fn) -> tuple[str, str]:
+    try:
+        value = str(value_fn())
+    except Exception as exc:  # pragma: no cover - defensive diagnostic path
+        value = f"!{type(exc).__name__}: {exc}"
+    return label, value
+
+
+def _nvidia_smi_summary() -> str:
+    if platform.system() != "Linux":
+        return "not linux"
+    try:
+        result = subprocess.run(
+            [
+                "nvidia-smi",
+                "--query-gpu=name,driver_version,memory.total",
+                "--format=csv,noheader,nounits",
+            ],
+            check=False,
+            capture_output=True,
+            text=True,
+            timeout=3,
+        )
+    except (OSError, subprocess.TimeoutExpired):
+        return "not found"
+    if result.returncode != 0:
+        message = (result.stderr or result.stdout).strip()
+        return message or f"exit {result.returncode}"
+    return result.stdout.strip() or "no gpu"
+
+
+def collect_info() -> list[tuple[str, str]]:
+    return [
+        ("kiwi-array", __version__),
+        ("python", sys.version.split()[0]),
+        ("platform", platform.platform()),
+        _safe_value("bridge", find_library_path),
+        _safe_value("runtime_libs", runtime_library_search_dir),
+        ("device_env", os.environ.get("KIWI_JUPYTER_DEVICE", "auto")),
+        ("vegalite_theme", os.environ.get("KIWI_VEGALITE_THEME", "default")),
+        ("nvidia_smi", _nvidia_smi_summary()),
+    ]
+
+
+def format_info_text(info: list[tuple[str, str]]) -> str:
+    width = max((len(label) for label, _ in info), default=0)
+    return "\n".join(f"{label.rjust(width)}: {value}" for label, value in info)
+
+
+def run_smoke(session: SessionProtocol, include_mlx: bool = False) -> list[SmokeCheck]:
+    checks: list[SmokeCheck] = []
+    for name, source, expected in [
+        ("scalar", "1+1", "2"),
+        ("vector", "+/1 2 3", "6"),
+        ("grad", "grad[{+/(x*x)}][1 2 3]", "2 4 6"),
+    ]:
+        result = session.eval(source)
+        ok = result.status == "ok" and result.text == expected
+        detail = result.text if result.status == "ok" and result.text is not None else f"!{result.status}"
+        checks.append(SmokeCheck(name=name, ok=ok, detail=detail))
+
+    if include_mlx:
+        setter = getattr(session, "set_global_mlx_float_array", None)
+        if setter is None:
+            checks.append(SmokeCheck(name="mlx_global", ok=False, detail="native bridge lacks MLX setter"))
+            return checks
+        try:
+            setter("kx", [1.0, 2.0, 3.0], (3,))
+            result = session.eval("grad[{+/(x*x)}][kx]")
+            ok = result.status == "ok" and result.text == "2 4 6"
+            detail = result.text if result.status == "ok" and result.text is not None else f"!{result.status}"
+            checks.append(SmokeCheck(name="mlx_global", ok=ok, detail=detail))
+        except Exception as exc:
+            checks.append(SmokeCheck(name="mlx_global", ok=False, detail=f"!{type(exc).__name__}: {exc}"))
+
+    return checks
+
+
+def format_smoke_text(checks: list[SmokeCheck]) -> str:
+    width = max((len(check.name) for check in checks), default=0)
+    lines = []
+    for check in checks:
+        status = "ok" if check.ok else "error"
+        lines.append(f"{check.name.rjust(width)}: {status} {check.detail}")
+    return "\n".join(lines)
+
+
+def load_ipython_extension(ipython) -> None:
+    from IPython.core.magic import Magics, line_cell_magic, line_magic, magics_class
+    from IPython.display import display
+
+    @magics_class
+    class KiwiMagics(Magics):
+        def __init__(self, shell) -> None:
+            super().__init__(shell)
+            self._session: KiwiSession | None = None
+
+        def _kiwi_session(self) -> KiwiSession:
+            if self._session is None:
+                device = os.environ.get("KIWI_JUPYTER_DEVICE", "auto")
+                library_path = os.environ.get("KIWI_BRIDGE_LIB")
+                self._session = KiwiSession(device=device, library_path=library_path)
+            return self._session
+
+        def _reset(self) -> None:
+            if self._session is not None:
+                self._session.close()
+                self._session = None
+
+        def _run(self, line: str, cell: str | None) -> None:
+            if cell is None and line.strip() in {"--reset", "-r"}:
+                self._reset()
+                return
+            code = line if cell is None else cell
+            outputs = execute_cell(self._kiwi_session(), code)
+            for output in outputs:
+                display(
+                    display_bundle_for_output(output),
+                    raw=True,
+                    metadata={"kiwi_autograd_path": output.autograd_path},
+                )
+
+        @line_cell_magic
+        def kiwi(self, line: str, cell: str | None = None) -> None:
+            self._run(line, cell)
+
+        @line_cell_magic
+        def k(self, line: str, cell: str | None = None) -> None:
+            self._run(line, cell)
+
+        @line_magic
+        def kinfo(self, line: str) -> None:
+            del line
+            print(format_info_text(collect_info()))
+
+        @line_magic
+        def ksmoke(self, line: str) -> None:
+            flags = set(line.split())
+            include_mlx = bool(flags & {"--mlx", "--gpu"})
+            print(format_smoke_text(run_smoke(self._kiwi_session(), include_mlx=include_mlx)))
+
+    ipython.register_magics(KiwiMagics)
+
+
+def unload_ipython_extension(ipython) -> None:
+    del ipython

kiwi_array-0.2.47.dist-info/METADATA

@@ -0,0 +1,65 @@
+Metadata-Version: 2.4
+Name: kiwi-array
+Version: 0.2.47
+Summary: Python loader, IPython magics, backend discovery, and CLI launcher for Kiwi.
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+Requires-Dist: kiwi-array-host==0.2.47
+Provides-Extra: host
+Requires-Dist: kiwi-array-host==0.2.47; extra == "host"
+Provides-Extra: cpu
+Requires-Dist: kiwi-array-cpu==0.2.47; extra == "cpu"
+Provides-Extra: metal
+Requires-Dist: kiwi-array-metal==0.2.47; extra == "metal"
+Provides-Extra: cuda12
+Requires-Dist: kiwi-array-cuda12==0.2.47; extra == "cuda12"
+Provides-Extra: notebook
+Requires-Dist: ipython>=8; extra == "notebook"
+Provides-Extra: jupyter
+Requires-Dist: kiwi-array-jupyter==0.2.47; extra == "jupyter"
+
+# kiwi-array
+
+`kiwi-array` is the Python loader package for Kiwi. It contains the
+`kiwi_array.bridge` ctypes wrapper, IPython magics, runtime backend discovery,
+and a small `kiwi` CLI launcher that execs the real native CLI from an
+installed backend payload.
+
+By default, installing `kiwi-array` also installs `kiwi-array-host`, the
+conservative host backend. Accelerated backends are explicit extras:
+`kiwi-array[cpu]`, `kiwi-array[metal]`, or `kiwi-array[cuda12]`.
+
+In IPython-compatible hosted notebooks, load the extension and use either the
+descriptive or short magic:
+
+```python
+%load_ext kiwi_array
+```
+
+```text
+%%k
+x:1 2 3
++/x
+```
+
+`%%kiwi` is registered as the explicit alias for `%%k`.
+
+Use the lightweight diagnostics when sharing hosted notebook setup cells:
+
+```python
+%kinfo
+%ksmoke
+```
+
+On a GPU runtime, `%ksmoke --gpu` also pushes a small MLX-backed vector through
+the native bridge and evaluates a gradient over it.
+
+Runtime payload packages such as `kiwi-array-host`, `kiwi-array-metal`, and
+`kiwi-array-cuda12` register themselves through the `kiwi_array.runtimes` entry
+point group. Set `KIWI_RUNTIME=host`, `KIWI_RUNTIME=cpu`,
+`KIWI_RUNTIME=metal`, `KIWI_RUNTIME=cuda12`, or leave `KIWI_RUNTIME=auto` to
+control runtime selection.
+
+Vega-Lite outputs use the notebook's default light rendering by default. Set
+`KIWI_VEGALITE_THEME=dark` before rendering if you want the HTML fallback to use
+Vega-Embed's dark theme.

kiwi_array-0.2.47.dist-info/RECORD

@@ -0,0 +1,9 @@
+kiwi_array/__init__.py,sha256=KS25NzHIwwg4tsOJ84QjhJM75tb_CxyN_yOigXRpNNg,174
+kiwi_array/bridge.py,sha256=CvOce3e65BanIBk_CH1aFnlNL1Yo5YOpWjrYgav8F3k,25420
+kiwi_array/cli.py,sha256=IsDcnvMOY1d_GTceFcrJwdx7vyR7yLMdEU3pOzQILJs,922
+kiwi_array/notebook.py,sha256=DfYRf1LALgBJV0XEBNXpWyi82bG8yeeVLKqce-R7cNM,8979
+kiwi_array-0.2.47.dist-info/METADATA,sha256=VQjuq1BfJ59x0RjC_b8_WJtz9wfP7i35ejrObXg_kz8,2154
+kiwi_array-0.2.47.dist-info/WHEEL,sha256=aeYiig01lYGDzBgS8HxWXOg3uV61G9ijOsup-k9o1sk,91
+kiwi_array-0.2.47.dist-info/entry_points.txt,sha256=rHIbSUqGhESM76fFGiKgf6InC-bDRNnJdJQe8uAEBAo,45
+kiwi_array-0.2.47.dist-info/top_level.txt,sha256=QZLBgjHZ20dRwlzhEjdi0HhTJaXovd4JoRSpsyu0dDk,11
+kiwi_array-0.2.47.dist-info/RECORD,,

kiwi_array-0.2.47.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (82.0.1)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

kiwi_array-0.2.47.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+kiwi = kiwi_array.cli:main

kiwi_array-0.2.47.dist-info/top_level.txt

@@ -0,0 +1 @@
+kiwi_array