klongpy 0.6.9__py3-none-any.whl → 0.7.1__py3-none-any.whl

klongpy/__init__.py

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

klongpy/adverbs.py

@@ -1,34 +1,9 @@
 from .core import *
-from .dyads import eval_dyad_add, eval_dyad_subtract, eval_dyad_multiply, eval_dyad_divide
 import functools
 import itertools
 
 
-def get_adverb_fn(klong, s, arity):
-    if s == "'":
-        return eval_adverb_each2 if arity == 2 else eval_adverb_each
-    elif s == '/':
-        return eval_adverb_over_neutral if arity == 2 else eval_adverb_over
-    elif s == '\\':
-        return eval_adverb_scan_over_neutral if arity == 2 else eval_adverb_scan_over
-    elif s == '\\~':
-        return (lambda f,a,b,k=klong: eval_adverb_scan_while(k,f,a,b)) if arity == 2 else eval_adverb_scan_converging
-    elif s == '\\*':
-        return eval_adverb_scan_iterating
-    elif s == ':\\':
-        return eval_adverb_each_left
-    elif s == ':\'':
-        return eval_adverb_each_pair
-    elif s == ':/':
-        return eval_adverb_each_right
-    elif s == ':*':
-        return eval_dyad_adverb_iterate
-    elif s == ':~':
-        return (lambda f,a,b,k=klong: eval_adverb_while(k,f,a,b)) if arity == 2 else eval_adverb_converge
-    raise RuntimeError(f"unknown adverb: {s}")
-
-
-def eval_adverb_converge(f, a, op):
+def eval_adverb_converge(f, a, op, backend):
     """
         f:~a                                                  [Converge]
 
@@ -54,10 +29,10 @@ def eval_adverb_converge(f, a, op):
     def _e(p,q):
         if not isinstance(p, type(q)):
             return False
-        if is_number(p):
-            return np.isclose(p,q)
-        elif np.isarray(p):
-            return kg_equal(p,q)
+        if backend.is_number(p):
+            return backend.np.isclose(p,q)
+        elif backend.is_array(p):
+            return backend.kg_equal(p, q)
         return p == q
     x = f(a)
     xx = f(x)
@@ -67,7 +42,7 @@ def eval_adverb_converge(f, a, op):
     return x
 
 
-def eval_adverb_each(f, a, op):
+def eval_adverb_each(f, a, op, backend):
     """
 
         f'a                                                       [Each]
@@ -92,17 +67,17 @@ def eval_adverb_each(f, a, op):
             return a
         has_str = False
         r = []
-        for x in str_to_chr_arr(a):
+        for x in backend.str_to_chr_arr(a):
             u = f(x)
             has_str |= isinstance(u,str)
             r.append(u)
-        return ''.join(r) if has_str else kg_asarray(r)
+        return ''.join(r) if has_str else backend.kg_asarray(r)
     if is_iterable(a):
         r = [f(x) for x in a]
-        return a if is_empty(a) else kg_asarray(r)
+        return a if is_empty(a) else backend.kg_asarray(r)
     elif is_dict(a):
-        r = [f(kg_asarray(x)) for x in a.items()]
-        return kg_asarray(r)
+        r = [f(backend.kg_asarray(x)) for x in a.items()]
+        return backend.kg_asarray(r)
     return f(a)
 
 
@@ -124,14 +99,14 @@ def eval_adverb_each2(f, a, b):
 
     """
     if is_empty(a) or is_empty(b):
-        return np.asarray([]) if is_list(a) or is_list(b) else ""
+        return bknp.asarray([]) if is_list(a) or is_list(b) else ""
     if is_atom(a) and is_atom(b):
         return f(a,b)
-    r = np.asarray([f(x,y) for x,y in zip(a,b)])
+    r = bknp.asarray([f(x,y) for x,y in zip(a,b)])
     return ''.join(r) if r.dtype == '<U1' else r
 
 
-def eval_adverb_each_left(f, a, b):
+def eval_adverb_each_left(f, a, b, backend):
     """
         a f:\b                                              [Each-Left]
         a f:/b                                              [Each-Right]
@@ -153,20 +128,19 @@ def eval_adverb_each_left(f, a, b):
         Examples: 1,:\[2 3 4]  -->  [[1 2] [1 3] [1 4]]
                   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])
+    b = backend.str_to_chr_arr(b) if isinstance(b,str) else b
+    return backend.kg_asarray([f(a,x) for x in b])
 
 
-def eval_adverb_each_right(f, a, b):
+def eval_adverb_each_right(f, a, b, backend):
     """
     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])
+    b = backend.str_to_chr_arr(b) if isinstance(b,str) else b
+    return backend.kg_asarray([f(x,a) for x in b])
 
 
-
-def eval_adverb_each_pair(f, a, op):
+def eval_adverb_each_pair(f, a, op, backend):
     """
 
         f:'a                                                 [Each-Pair]
@@ -185,8 +159,8 @@ def eval_adverb_each_pair(f, a, op):
     if is_atom(a) or (is_iterable(a) and len(a) == 1):
         return a
     j = isinstance(a, str)
-    a = str_to_chr_arr(a) if j else a
-    return kg_asarray([f(x,y) for x,y in zip(a[::],a[1::])])
+    a = backend.str_to_chr_arr(a) if j else a
+    return backend.kg_asarray([f(x,y) for x,y in zip(a[::],a[1::])])
 
 
 def eval_dyad_adverb_iterate(f, a, b):
@@ -208,7 +182,7 @@ def eval_dyad_adverb_iterate(f, a, b):
     return b
 
 
-def eval_adverb_over(f, a, op):
+def eval_adverb_over(f, a, op, backend):
     """
         f/a                                                       [Over]
 
@@ -227,23 +201,23 @@ 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 backend's ufunc reduce when available for better performance
+    np_backend = backend.np
     if isinstance(op, KGOp):
         if safe_eq(op.a,'+'):
-            return np.add.reduce(a)
+            return np_backend.add.reduce(a)
         elif safe_eq(op.a, '-'):
-            return np.subtract.reduce(a)
-        elif safe_eq(op.a, '*') and hasattr(np.multiply,'reduce'):
-            return np.multiply.reduce(a)
-        elif safe_eq(op.a, '%') and hasattr(np.divide,'reduce'):
-            return np.divide.reduce(a)
+            return np_backend.subtract.reduce(a)
+        elif safe_eq(op.a, '*') and hasattr(np_backend.multiply,'reduce'):
+            return np_backend.multiply.reduce(a)
+        elif safe_eq(op.a, '%') and hasattr(np_backend.divide,'reduce'):
+            return np_backend.divide.reduce(a)
         elif safe_eq(op.a, '&') and a.ndim == 1:
-            return np.min(a)
+            return np_backend.min(a)
         elif safe_eq(op.a, '|') and a.ndim == 1:
-            return np.max(a)
-        elif safe_eq(op.a, ',') and np.isarray(a) and a.dtype != 'O':
-            return a if a.ndim == 1 else np.concatenate(a, axis=0)
+            return np_backend.max(a)
+        elif safe_eq(op.a, ',') and np_backend.isarray(a) and a.dtype != 'O':
+            return a if a.ndim == 1 else np_backend.concatenate(a, axis=0)
     return functools.reduce(f, a)
 
 
@@ -280,7 +254,7 @@ def eval_adverb_over_neutral(f, a, b):
     return functools.reduce(f,b[1:],f(a,b[0]))
 
 
-def eval_adverb_scan_over_neutral(f, a, b):
+def eval_adverb_scan_over_neutral(f, a, b, backend):
     """
 
         f\a                                                  [Scan-Over]
@@ -309,31 +283,33 @@ def eval_adverb_scan_over_neutral(f, a, b):
         b = [b]
     b = [f(a,b[0]), *b[1:]]
     r = list(itertools.accumulate(b,f))
-    q = kg_asarray(r)
+    q = backend.kg_asarray(r)
     r = [a, *q]
-    return kg_asarray(r)
+    return backend.kg_asarray(r)
 
 
-def eval_adverb_scan_over(f, a, op):
+def eval_adverb_scan_over(f, a, op, backend):
     """
         see eval_adverb_scan_over_neutral
     """
     if is_atom(a):
         return a
-    # https://docs.cupy.dev/en/stable/reference/ufunc.html
-    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'):
-        return np.subtract.accumulate(a)
-    elif safe_eq(f, eval_dyad_multiply) and hasattr(np.multiple, 'accumulate'):
-        return np.multiple.accumulate(a)
-    elif safe_eq(f, eval_dyad_divide) and hasattr(np.divide, 'accumulate'):
-        return np.divide.accumulate(a)
+    # Use backend's ufunc accumulate when available for better performance
+    np_backend = backend.np
+    if isinstance(op, KGOp):
+        if safe_eq(op.a, '+') and hasattr(np_backend.add, 'accumulate'):
+            return np_backend.add.accumulate(a)
+        elif safe_eq(op.a, '-') and hasattr(np_backend.subtract, 'accumulate'):
+            return np_backend.subtract.accumulate(a)
+        elif safe_eq(op.a, '*') and hasattr(np_backend.multiply, 'accumulate'):
+            return np_backend.multiply.accumulate(a)
+        elif safe_eq(op.a, '%') and hasattr(np_backend.divide, 'accumulate'):
+            return np_backend.divide.accumulate(a)
     r = list(itertools.accumulate(a, f))
-    return kg_asarray(r)
+    return backend.kg_asarray(r)
 
 
-def eval_adverb_scan_converging(f, a, op):
+def eval_adverb_scan_converging(f, a, op, backend):
     """
 
         f\~a                                           [Scan-Converging]
@@ -356,15 +332,15 @@ def eval_adverb_scan_converging(f, a, op):
     x = a
     xx = f(a)
     r = [a, xx]
-    while not kg_equal(x,xx):
+    while not backend.kg_equal(x, xx):
         x = xx
         xx = f(x)
         r.append(xx)
     r.pop()
-    return kg_asarray(r)
+    return backend.kg_asarray(r)
 
 
-def eval_adverb_scan_while(klong, f, a, b):
+def eval_adverb_scan_while(klong, f, a, b, backend):
     """
 
         a f\~b                                              [Scan-While]
@@ -389,10 +365,10 @@ def eval_adverb_scan_while(klong, f, a, b):
         b = f(b)
         r.append(b)
     r.pop()
-    return kg_asarray(r)
+    return backend.kg_asarray(r)
 
 
-def eval_adverb_scan_iterating(f, a, b):
+def eval_adverb_scan_iterating(f, a, b, backend):
     """
 
         a f\*b                                          [Scan-Iterating]
@@ -410,7 +386,7 @@ def eval_adverb_scan_iterating(f, a, b):
         b = f(b)
         r.append(b)
         a = a - 1
-    return kg_asarray(r)
+    return backend.kg_asarray(r)
 
 
 def eval_adverb_while(klong, f, a, b):
@@ -429,3 +405,29 @@ def eval_adverb_while(klong, f, a, b):
     while klong.eval(KGCall(a, b, arity=1)):
         b = f(b)
     return b
+
+
+def get_adverb_fn(klong, s, arity):
+    backend = klong._backend
+
+    if s == "'":
+        return eval_adverb_each2 if arity == 2 else lambda f,a,op: eval_adverb_each(f,a,op,backend)
+    elif s == '/':
+        return eval_adverb_over_neutral if arity == 2 else lambda f,a,op: eval_adverb_over(f,a,op,backend)
+    elif s == '\\':
+        return (lambda f,a,b: eval_adverb_scan_over_neutral(f,a,b,backend)) if arity == 2 else lambda f,a,op: eval_adverb_scan_over(f,a,op,backend)
+    elif s == '\\~':
+        return (lambda f,a,b: eval_adverb_scan_while(klong,f,a,b,backend)) if arity == 2 else lambda f,a,op: eval_adverb_scan_converging(f,a,op,backend)
+    elif s == '\\*':
+        return lambda f,a,b: eval_adverb_scan_iterating(f,a,b,backend)
+    elif s == ':\\':
+        return lambda f,a,b: eval_adverb_each_left(f,a,b,backend)
+    elif s == ':\'':
+        return lambda f,a,op: eval_adverb_each_pair(f,a,op,backend)
+    elif s == ':/':
+        return lambda f,a,b: eval_adverb_each_right(f,a,b,backend)
+    elif s == ':*':
+        return eval_dyad_adverb_iterate
+    elif s == ':~':
+        return (lambda f,a,b: eval_adverb_while(klong,f,a,b)) if arity == 2 else lambda f,a,op: eval_adverb_converge(f,a,op,backend)
+    raise RuntimeError(f"unknown adverb: {s}")

klongpy/autograd.py

@@ -0,0 +1,299 @@
+import numpy as np
+from .core import KGLambda, KGCall, KGSym, KGFn
+
+
+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 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