klongpy 0.6.8__py3-none-any.whl → 0.7.0__py3-none-any.whl

klongpy/__init__.py

@@ -1,2 +1,20 @@
 from .interpreter import KlongInterpreter, KlongException
-__all__ = ["KlongInterpreter", "KlongException"]
+from .backend import TorchUnsupportedDtypeError
+from .backends import (
+    get_backend,
+    register_backend,
+    list_backends,
+    BackendProvider,
+    UnsupportedDtypeError,
+)
+
+__all__ = [
+    "KlongInterpreter",
+    "KlongException",
+    "TorchUnsupportedDtypeError",
+    "UnsupportedDtypeError",
+    "get_backend",
+    "register_backend",
+    "list_backends",
+    "BackendProvider",
+]

klongpy/adverbs.py

@@ -1,5 +1,6 @@
 from .core import *
 from .dyads import eval_dyad_add, eval_dyad_subtract, eval_dyad_multiply, eval_dyad_divide
+from .backend import kg_asarray, is_number, str_to_chr_arr, kg_equal
 import functools
 import itertools
 
@@ -154,7 +155,7 @@ def eval_adverb_each_left(f, a, b):
                   1,:/[2 3 4]  -->  [[2 1] [3 1] [4 1]]
     """
     b = str_to_chr_arr(b) if isinstance(b,str) else b
-    return np.asarray([f(a,x) for x in b])
+    return kg_asarray([f(a,x) for x in b])
 
 
 def eval_adverb_each_right(f, a, b):
@@ -162,7 +163,7 @@ def eval_adverb_each_right(f, a, b):
     see: eval_dyad_adverb_each_left
     """
     b = str_to_chr_arr(b) if isinstance(b,str) else b
-    return np.asarray([f(x,a) for x in b])
+    return kg_asarray([f(x,a) for x in b])
 
 
 
@@ -227,8 +228,7 @@ def eval_adverb_over(f, a, op):
         return a
     if len(a) == 1:
         return a[0]
-    # https://docs.cupy.dev/en/stable/reference/ufunc.html
-    # TODO: can we use NumPy reduce when CuPy backend primary?
+    # Use NumPy/PyTorch ufunc reduce when available for better performance
     if isinstance(op, KGOp):
         if safe_eq(op.a,'+'):
             return np.add.reduce(a)
@@ -320,7 +320,7 @@ def eval_adverb_scan_over(f, a, op):
     """
     if is_atom(a):
         return a
-    # https://docs.cupy.dev/en/stable/reference/ufunc.html
+    # Use NumPy/PyTorch ufunc accumulate when available for better performance
     if safe_eq(f, eval_dyad_add) and hasattr(np.add, 'accumulate'):
         return np.add.accumulate(a)
     elif safe_eq(f, eval_dyad_subtract) and hasattr(np.subtract, 'accumulate'):

klongpy/autograd.py

@@ -0,0 +1,308 @@
+import numpy as np
+from .core import KGLambda, KGCall, KGSym, KGFn
+from .backend import get_default_backend
+
+
+class AutogradError(Exception):
+    """Base class for autograd-related errors."""
+    pass
+
+
+class AutogradChainBrokenError(AutogradError):
+    """Raised when the gradient computation chain is broken."""
+
+    def __init__(self, context, expected, actual, suggestion=None):
+        self.context = context
+        self.expected = expected
+        self.actual = actual
+        self.suggestion = suggestion
+        msg = f"Autograd chain broken at {context}: expected {expected}, got {actual}."
+        if suggestion:
+            msg += f" {suggestion}"
+        super().__init__(msg)
+
+
+class NonScalarLossError(AutogradError):
+    """Raised when the loss function returns a non-scalar value."""
+
+    def __init__(self, shape):
+        self.shape = shape
+        super().__init__(
+            f"Loss function must return a scalar, got shape {shape}. "
+            "Use sum (+/) or mean (%#) to reduce to a scalar."
+        )
+
+
+def _get_float_dtype(backend):
+    """Get the appropriate float dtype for the current backend."""
+    # MPS doesn't support float64
+    if hasattr(backend, 'supports_float64') and not backend.supports_float64():
+        return np.float32
+    return np.float64
+
+
+def _scalar_value(x, backend):
+    """Extract scalar value from various array/tensor types.
+
+    Raises:
+        NonScalarLossError: If x is not a scalar value.
+    """
+    x = backend.to_numpy(x) if backend.is_backend_array(x) else x
+    if isinstance(x, np.ndarray):
+        if x.ndim == 0:
+            return float(x.item())
+        elif x.size == 1:
+            return float(x.flat[0])
+        else:
+            raise NonScalarLossError(tuple(x.shape))
+    return float(x)
+
+
+def _to_func_input(x, backend, require_grad=False):
+    """Convert numpy array to appropriate input type for function call.
+
+    Args:
+        x: Input array (numpy)
+        backend: Backend provider
+        require_grad: If True and backend supports autograd, create grad tensor.
+                      For numeric gradient, this should be False.
+    """
+    if require_grad and backend.supports_autograd():
+        return backend.create_grad_tensor(x)
+    return x
+
+
+def _invoke_fn(klong, fn, args):
+    """Invoke a Klong function with the given arguments.
+
+    Handles all function types uniformly:
+    - KGSym, KGLambda: wrap in KGCall with args
+    - KGFn, KGCall: extract inner function, wrap in KGCall with args
+    - callable: call directly with args
+    """
+    if callable(fn) and not isinstance(fn, (KGSym, KGLambda, KGFn)):
+        return fn(*args)
+    inner = fn.a if isinstance(fn, KGFn) else fn
+    return klong.call(KGCall(inner, list(args), len(args)))
+
+
+def numeric_grad(func, x, backend, eps=None):
+    """Compute numeric gradient of scalar-valued function."""
+    # Get appropriate float dtype
+    float_dtype = _get_float_dtype(backend)
+
+    # Use larger epsilon for float32 to maintain precision
+    if eps is None:
+        eps = 1e-4 if float_dtype == np.float32 else 1e-6
+
+    # Convert backend tensors to numpy for gradient computation
+    if backend.is_backend_array(x):
+        x = backend.to_numpy(x)
+    x = np.asarray(x, dtype=float_dtype)
+
+    grad = np.zeros_like(x, dtype=float_dtype)
+    it = np.nditer(x, flags=['multi_index'], op_flags=['readwrite'])
+    while not it.finished:
+        idx = it.multi_index
+        orig = float(x[idx])
+        x[idx] = orig + eps
+        f_pos = _scalar_value(func(_to_func_input(x.copy(), backend)), backend)
+        x[idx] = orig - eps
+        f_neg = _scalar_value(func(_to_func_input(x.copy(), backend)), backend)
+        grad[idx] = (f_pos - f_neg) / (2 * eps)
+        x[idx] = orig
+        it.iternext()
+    return grad
+
+
+def grad_of_fn(klong, fn, x):
+    """
+    Return gradient of Klong or Python function ``fn`` at ``x``.
+
+    Uses PyTorch autograd when available (USE_TORCH=1), otherwise
+    falls back to numeric differentiation.
+    """
+    backend = klong._backend
+    call_fn = lambda v: _invoke_fn(klong, fn, [v])
+
+    if backend.supports_autograd():
+        return backend.compute_autograd(call_fn, x)
+    else:
+        return numeric_grad(call_fn, x, backend)
+
+
+def torch_autograd(func, x):
+    """Compute gradient using PyTorch autograd (requires torch backend)."""
+    backend = get_default_backend()
+    if not backend.supports_autograd():
+        raise RuntimeError("PyTorch autograd requires torch backend (USE_TORCH=1)")
+    return backend.compute_autograd(func, x)
+
+
+def numeric_jacobian(func, x, backend, eps=None):
+    """
+    Compute Jacobian matrix of func at point x using finite differences.
+
+    For f: R^n -> R^m, returns m x n matrix where J[i,j] = df_i/dx_j.
+
+    Args:
+        func: Callable that takes an array and returns an array
+        x: Input point (array)
+        backend: Backend provider
+        eps: Step size for finite differences (default: 1e-6 or 1e-4 for float32)
+
+    Returns:
+        Jacobian matrix as numpy array
+    """
+    float_dtype = _get_float_dtype(backend)
+    if eps is None:
+        eps = 1e-4 if float_dtype == np.float32 else 1e-6
+
+    # Convert to numpy
+    if backend.is_backend_array(x):
+        x = backend.to_numpy(x)
+    x = np.asarray(x, dtype=float_dtype).flatten()
+
+    # Evaluate function at x to get output shape
+    f0 = func(_to_func_input(x.copy(), backend))
+    if backend.is_backend_array(f0):
+        f0 = backend.to_numpy(f0)
+    f0 = np.asarray(f0, dtype=float_dtype).flatten()
+
+    n = len(x)  # Input dimension
+    m = len(f0)  # Output dimension
+    jacobian = np.zeros((m, n), dtype=float_dtype)
+
+    for j in range(n):
+        x_plus = x.copy()
+        x_plus[j] += eps
+        x_minus = x.copy()
+        x_minus[j] -= eps
+
+        f_plus = func(_to_func_input(x_plus, backend))
+        f_minus = func(_to_func_input(x_minus, backend))
+
+        if backend.is_backend_array(f_plus):
+            f_plus = backend.to_numpy(f_plus)
+        if backend.is_backend_array(f_minus):
+            f_minus = backend.to_numpy(f_minus)
+
+        f_plus = np.asarray(f_plus, dtype=float_dtype).flatten()
+        f_minus = np.asarray(f_minus, dtype=float_dtype).flatten()
+
+        jacobian[:, j] = (f_plus - f_minus) / (2 * eps)
+
+    return jacobian
+
+
+def jacobian_of_fn(klong, fn, x):
+    """
+    Compute Jacobian matrix of Klong function fn at point x.
+
+    For f: R^n -> R^m, returns m x n matrix where J[i,j] = df_i/dx_j.
+
+    Args:
+        klong: KlongInterpreter instance
+        fn: Function (KGSym, KGLambda, KGFn, KGCall, or callable)
+        x: Input point
+
+    Returns:
+        Jacobian matrix
+    """
+    backend = klong._backend
+    call_fn = lambda v: _invoke_fn(klong, fn, [v])
+
+    if backend.supports_autograd():
+        try:
+            return backend.compute_jacobian(call_fn, x)
+        except Exception:
+            # Fall back to numeric if torch jacobian fails
+            return numeric_jacobian(call_fn, x, backend=backend)
+    else:
+        return numeric_jacobian(call_fn, x, backend=backend)
+
+
+def multi_jacobian_of_fn(klong, fn, param_syms):
+    """
+    Compute Jacobians for multiple parameters in one call.
+
+    Args:
+        klong: KlongInterpreter instance
+        fn: Function (KGSym, KGLambda, KGFn, KGCall, or callable)
+             Should be a niladic function that references the parameters
+        param_syms: List of KGSym parameter symbols to differentiate with respect to
+
+    Returns:
+        List of Jacobian matrices, one per parameter
+    """
+    backend = klong._backend
+    param_values = [klong[sym] for sym in param_syms]
+    call_fn = lambda: _invoke_fn(klong, fn, [])
+
+    jacobians = []
+    for sym, val in zip(param_syms, param_values):
+        original = klong[sym]
+
+        def single_param_fn(v, s=sym, orig=original):
+            """Wrapper that sets param to v, calls fn, restores param."""
+            klong[s] = v
+            try:
+                return call_fn()
+            finally:
+                klong[s] = orig
+
+        if backend.supports_autograd():
+            try:
+                jac = backend.compute_jacobian(single_param_fn, val)
+            except Exception:
+                jac = numeric_jacobian(single_param_fn, val, backend=backend)
+        else:
+            jac = numeric_jacobian(single_param_fn, val, backend=backend)
+
+        # Restore original value after jacobian computation
+        klong[sym] = original
+        jacobians.append(jac)
+
+    return jacobians
+
+
+def multi_grad_of_fn(klong, fn, param_syms):
+    """
+    Compute gradients for multiple parameters in one call.
+
+    Args:
+        klong: KlongInterpreter instance
+        fn: Loss function (KGSym, KGLambda, KGFn, KGCall, or callable)
+             Should be a niladic function that references the parameters
+        param_syms: List of KGSym parameter symbols to differentiate with respect to
+
+    Returns:
+        List of gradients, one per parameter
+    """
+    backend = klong._backend
+    # Access context directly to avoid KGFnWrapper wrapping
+    param_values = [klong._context[sym] for sym in param_syms]
+
+    def call_fn_with_tensors(tensors):
+        """Call the loss function with tensor values temporarily bound to symbols."""
+        originals = {sym: klong._context[sym] for sym in param_syms}
+        try:
+            for sym, tensor in zip(param_syms, tensors):
+                klong[sym] = tensor
+            return _invoke_fn(klong, fn, [])
+        finally:
+            for sym, orig in originals.items():
+                klong[sym] = orig
+
+    if backend.supports_autograd():
+        return backend.compute_multi_autograd(call_fn_with_tensors, param_values)
+    else:
+        # Fallback: compute numeric gradients one at a time
+        grads = []
+        for i, sym in enumerate(param_syms):
+            def single_param_fn(v, idx=i):
+                vals = list(param_values)
+                vals[idx] = v
+                return call_fn_with_tensors(vals)
+            grads.append(numeric_grad(single_param_fn, param_values[i], backend))
+        return grads

klongpy/backend.py

@@ -1,103 +1,171 @@
-import os
-import warnings
-
-# Attempt to import CuPy. If not available, set use_gpu to False.
-use_gpu = bool(os.environ.get('USE_GPU') == '1')
-if use_gpu:
-    try:
-        import cupy as np
-        use_gpu = True
-    except ImportError:
-        import numpy as np
-        use_gpu = False
-else:
-    import numpy as np
-
-
-def is_supported_type(x):
-    """
-    CuPy does not support strings or jagged arrays.
-    Note: add any other unsupported types here.
-    """
-    if isinstance(x, str) or is_jagged_array(x):
-        return False
-    return True
+"""
+Backend compatibility module for KlongPy.
+
+This module provides backward compatibility with the old global backend system.
+New code should use the backends package directly:
+
+    from klongpy.backends import get_backend, BackendProvider
+
+For per-interpreter backends, use:
+
+    klong = KlongInterpreter(backend='torch')
+"""
+import numpy as real_np
+
+from .backends import (
+    get_backend,
+    register_backend,
+    list_backends,
+    BackendProvider,
+    UnsupportedDtypeError,
+    TorchUnsupportedDtypeError,
+    NumpyBackendProvider,
+    TorchBackendProvider,
+    KGChar,
+    is_jagged_array,
+    is_supported_type,
+)
+
+
+# Global backend state for backward compatibility
+# This is used by modules that import `np` and `use_torch` directly
+_default_backend = get_backend()
+
+# Backward compatibility: expose np and use_torch at module level
+np = _default_backend.np
+use_torch = _default_backend.name == 'torch'
+
+
+def get_default_backend():
+    """Get the default backend provider."""
+    return _default_backend
 
 
-def is_jagged_array(x):
+def to_numpy(x):
+    """Convert tensor/array to numpy, handling device transfers and 0-dim arrays."""
+    result = _default_backend.to_numpy(x)
+    if isinstance(result, real_np.ndarray) and result.ndim == 0:
+        return result.item()
+    return result
+
+
+def to_display(x):
+    """Convert value to display-friendly format."""
+    return _default_backend.to_display(x)
+
+
+def array_size(a):
     """
-    Check if an array is jagged.
+    Get the total number of elements in an array/tensor.
+
+    Works with both numpy arrays and torch tensors.
     """
-    if isinstance(x, list):
-        # If the lengths of sublists vary, it's a jagged array.
-        return len(set(map(len, x))) > 1
-    return False
-
-if use_gpu:
-    import cupy
-    import numpy
-
-    class CuPyReductionKernelWrapper:
-        def __init__(self, fn, reduce_fn_1, reduce_fn_2):
-            self.fn = fn
-            self.reduce_fn_1 = reduce_fn_1
-            self.reduce_fn_2 = reduce_fn_2
-
-        def __call__(self, *args, **kwargs):
-            return self.fn(*args, **kwargs)
-
-        def reduce(self, x):
-            return self.reduce_fn_1(x) if x.ndim == 1 else self.reduce_fn_2(x[0], x[1])
-
-    add_reduce_2 = cupy.ElementwiseKernel(
-            'T x, T y',
-            'T z',
-            'z = (x + y)',
-            'add_reduce_2')
-    np.add = CuPyReductionKernelWrapper(cupy.add, cupy.sum, add_reduce_2)
-
-    def subtract_reduce_1(x):
-        return 2*x[0] - cupy.sum(x)
-
-    subtract_reduce_2 = cupy.ElementwiseKernel(
-            'T x, T y',
-            'T z',
-            'z = (x - y)',
-            'subtract_reduce_2')
-    np.subtract = CuPyReductionKernelWrapper(cupy.subtract, subtract_reduce_1, subtract_reduce_2)
-
-    multiply_reduce_1 = cupy.ReductionKernel(
-                'T x',
-                'T y',
-                'x',
-                'a * b',
-                'y = a',
-                '1',
-                'multiply_reduce_1'
-             )
-    multiply_reduce_2 = cupy.ElementwiseKernel(
-            'T x, T y',
-            'T z',
-            'z = (x * y)',
-            'multiply_reduce_2')
-    np.multiply = CuPyReductionKernelWrapper(cupy.multiply, multiply_reduce_1, multiply_reduce_2)
-
-    def divide_reduce_1(x):
-        raise NotImplementedError()
-
-    divide_reduce_2 = cupy.ElementwiseKernel(
-            'T x, T y',
-            'T z',
-            'z = (x / y)',
-            'divide_reduce_2')
-    np.divide = CuPyReductionKernelWrapper(cupy.divide, divide_reduce_1, divide_reduce_2)
-
-    np.isarray = lambda x: isinstance(x, (numpy.ndarray, cupy.ndarray))
-
-#    np.hstack = lambda x: cupy.hstack(x) if use_gpu and is_supported_type(x) else numpy.hstack(x)
-else:
-    np.seterr(divide='ignore')
-    warnings.filterwarnings("error", category=np.VisibleDeprecationWarning)
-    np.isarray = lambda x: isinstance(x, np.ndarray)
-
-np
+    return _default_backend.array_size(a)
+
+
+def safe_equal(x, y):
+    """Compare two values for equality, handling backend-specific array types."""
+    return _default_backend.safe_equal(x, y)
+
+
+def detach_if_needed(x):
+    """Detach array from computation graph if needed."""
+    return _default_backend.detach_if_needed(x)
+
+
+def to_int_array(a):
+    """Convert array to integer type."""
+    return _default_backend.to_int_array(a)
+
+
+def power(a, b):
+    """Compute a^b, handling gradient tracking if applicable."""
+    return _default_backend.power(a, b)
+
+
+def has_gradient(x):
+    """Check if x is tracking gradients (for autograd)."""
+    return _default_backend.has_gradient(x)
+
+
+def kg_asarray(a):
+    """Convert input to array using the default backend's kg_asarray method."""
+    return _default_backend.kg_asarray(a)
+
+
+def is_integer(x):
+    """Check if x is an integer type using the default backend."""
+    return _default_backend.is_integer(x)
+
+
+def is_float(x):
+    """Check if x is a float type using the default backend."""
+    return _default_backend.is_float(x)
+
+
+def is_number(a):
+    """Check if a is a number (integer or float) using the default backend."""
+    return _default_backend.is_number(a)
+
+
+def get_dtype_kind(arr):
+    """Get the dtype 'kind' character for an array using the default backend."""
+    return _default_backend.get_dtype_kind(arr)
+
+
+def str_to_chr_arr(s):
+    """Convert string to character array using the default backend."""
+    return _default_backend.str_to_chr_arr(s)
+
+
+def kg_argsort(a, descending=False):
+    """Argsort array using the default backend."""
+    from .core import kg_argsort as core_kg_argsort
+    return core_kg_argsort(a, _default_backend, descending=descending)
+
+
+def vec_fn(a, f):
+    """Apply a function f to an array a, with support for nested arrays."""
+    from .core import vec_fn as core_vec_fn
+    return core_vec_fn(a, f, _default_backend)
+
+
+def kg_equal(a, b):
+    """Compare two values or arrays for equality using the default backend."""
+    from .core import kg_equal as core_kg_equal
+    return core_kg_equal(a, b, _default_backend)
+
+
+__all__ = [
+    'np',
+    'use_torch',
+    'get_backend',
+    'get_default_backend',
+    'register_backend',
+    'list_backends',
+    'BackendProvider',
+    'UnsupportedDtypeError',
+    'TorchUnsupportedDtypeError',
+    'NumpyBackendProvider',
+    'TorchBackendProvider',
+    'KGChar',
+    'is_supported_type',
+    'is_jagged_array',
+    'to_numpy',
+    'to_display',
+    'array_size',
+    'safe_equal',
+    'detach_if_needed',
+    'to_int_array',
+    'power',
+    'has_gradient',
+    'kg_asarray',
+    'is_integer',
+    'is_float',
+    'is_number',
+    'get_dtype_kind',
+    'str_to_chr_arr',
+    'kg_argsort',
+    'vec_fn',
+    'kg_equal',
+]