LPU with VJP/JVP #1

Activity

1. heptaflar

   added

   enhancementNew feature or requestNew feature or request

   and removed

   enhancementNew feature or requestNew feature or request

   on Apr 23, 2026
2. heptaflar

   added the

   Feature

   issue type on Apr 23, 2026
3. heptaflar

   changed the title [-]Implement LPU with VJP/JVP for JAX-based F/M[/-] [+]LPU with VJP/JVP for JAX-based F/M[/+] on Apr 23, 2026
4. heptaflar commented on Oct 4, 2026

   

   heptaflar

   on Oct 4, 2026

   Last edited by heptaflar

   CollaboratorAuthor

   More actions

   Google AI generated code:

   from typing import Callable
   import jax
   import jax.numpy as jnp
   from jaxtyping import Array
   
   def make_lpu_functional(f: Callable[[Array], Array], d: int, diag: bool = False) -> Callable[[Array, Array], Array]:
       """
       Erzeugt das Gesetz der Fehlerfortpflanzung (LPU) rein funktional über JVP/VJP.
       Erfüllt das Principle of Least Knowledge: Es wird kein Jacobi-Tensor allokiert.
   
       :param f: Die primale Funktion f(x) -> y.
       :param d: Die Anzahl der inneren Tensor-Dimensionen (Stufe der Eingabe).
       :param diag: True, wenn nur die Varianzen der Ausgabe berechnet werden sollen.
       :returns: Eine Funktion lpu(x, u) -> Ausgangs-Unsicherheitstensor.
       """
   
       def lpu(x: Array, u: Array) -> Array:
           """
           :param x: Primaler Input-Tensor. Shape: (*In_Shape)
           :param u: Unsicherheits-Tensor.
                     Falls unkorreliert (u.ndim == d): Shape (*In_Shape)
                     Falls voll-korreliert (u.ndim == 2*d): Shape (*In_Shape, *In_Shape)
           """
           in_shape = x.shape
           
           # --- FALL 1: Fehler in x sind strikt unkorreliert (u.ndim == d) ---
           if u.ndim == d:
               # Wir bauen die Basis der Einheitsrichtungen rein strukturell auf
               n_elements = x.size
               eye_flat = jnp.eye(n_elements)
               eye_tensor = eye_flat.reshape(in_shape + in_shape)
               
               # Skalierung mit den Standardabweichungen (Wurzel aus Varianz)
               std_dev = jnp.sqrt(u)
               std_expanded = std_dev[(None,) * d + (...,)]
               V_tensors = eye_tensor * std_expanded  # Shape: (*In_Shape_links, *In_Shape_rechts)
   
               def ein_jvp(v_richtung):
                   _, jvp_erg = jax.jvp(f, (x,), (v_richtung,))
                   return jvp_erg  # Shape: (*Out_Shape)
   
               # d-mal vmap über die linken In_Shape-Achsen
               v_func = ein_jvp
               for _ in range(d):
                   v_func = jax.vmap(v_func, in_axes=0, out_axes=0)
               
               jvp_all = v_func(V_tensors)  # Shape: (*In_Shape, *Out_Shape)
               
               # Fehlerfortpflanzung: Summe der quadrierten JVPs über die Eingangsachsen
               return jnp.sum(jvp_all**2, axis=tuple(range(d)))
   
           # --- FALL 2: Volle Eingangs-Kovarianzmatrix (u.ndim == 2*d) ---
           else:
               # Schritt 1: W = G * U  via JVP (Vorwärtsmodus)
               def ein_jvp(v_richtung):
                   _, jvp_erg = jax.jvp(f, (x,), (v_richtung,))
                   return jvp_erg
   
               # vmap über die RECHTEN d Achsen von U (die Richtungsachsen)
               w_func = ein_jvp
               for _ in range(d):
                   w_func = jax.vmap(w_func, in_axes=-1, out_axes=-1)
                   
               W = w_func(u)  # Shape: (*In_Shape, *Out_Shape)
   
               # Schritt 2: V = W * G^T via VJP (Rückwärtsmodus)
               y, vjp_fun = jax.vjp(f, x)
               out_rank = len(y.shape)
               
               # Fall 2a: Nur die Ausgangs-Varianzen gewünscht (diag=True)
               if diag:
                   # Wir mappen das VJP über die hinteren Out_Shape-Achsen von W
                   v_func = vjp_fun
                   for _ in range(out_rank):
                       v_func = jax.vmap(v_func, in_axes=-1, out_axes=-1)
                   
                   G_W = v_func(W)  # Shape: (*In_Shape, *In_Shape)
                   
                   # Symmetrische Reduktion auf die Ausgangsvarianzen
                   return jnp.sum(G_W, axis=tuple(range(d)))
   
               # Fall 2b: Volle Ausgangs-Kovarianzmatrix gewünscht
               else:
                   # Wir mappen das VJP über die vorderen In_Shape-Achsen von W
                   v_func = vjp_fun
                   for _ in range(d):
                       v_func = jax.vmap(v_func, in_axes=0, out_axes=0)
                   
                   return v_func(W)  # Shape: (*Out_Shape, *Out_Shape)
   
       return lpu
   
   
   # Universelle, JIT-kompilierte Batch-Verarbeitung über Achse 0
   @jax.jit(static_argnums=(0, 3, 4))
   def lpu_batched(f: Callable, d: int, batch_x: Array, batch_u: Array, diag: bool = False) -> Array:
       """
       Verarbeitet eine Batch von Eingaben parallel über Achse 0 analog zu Ihrer Vorlage.
       """
       return jax.vmap(make_lpu_functional(f, d, diag), in_axes=(0, 0))(batch_x, batch_u)
5. heptaflar

   changed the title [-]LPU with VJP/JVP for JAX-based F/M[/-] [+]LPU with VJP/JVP[/+] on Oct 4, 2026
6. heptaflar

   added

   enhancementNew feature or requestNew feature or request

   on Oct 4, 2026
7. heptaflar

   self-assigned this

   on Oct 4, 2026
8. heptaflar

   removed

   enhancementNew feature or requestNew feature or request

   on Oct 4, 2026
9. heptaflar commented on Oct 4, 2026

   

   heptaflar

   on Oct 4, 2026

   CollaboratorAuthor

   More actions

   # 1. Ihre bewährte Matrix-Variante (umbenannt für interne Eindeutigkeit)
   def make_lpu_matrix(d: int, diag: bool = False):
       def _lpu(g: Array, u: Array) -> Array:
           dims = tuple(range(-d, 0))
           gu = jnp.tensordot(g, u, (dims, dims)) if u.ndim != d else g * u
           return jnp.tensordot(gu, g, (dims, dims)) if not diag else jnp.sum(gu * g, dims)
       return _lpu
   
   @jax.jit(static_argnums=(0, 3))
   def lpu_matrix_batched(d: int, g: Array, u: Array, diag: bool = False) -> Array:
       return jax.vmap(make_lpu_matrix(d, diag), in_axes=(0, 0))(g, u)
   
   
   # 2. Die funktionale JVP/VJP-Variante (ohne Flachklopfen)
   def make_lpu_functional(f: Callable[[Array], Array], d: int, diag: bool = False):
       def _lpu(x: Array, u: Array) -> Array:
           in_shape = x.shape
           if u.ndim == d:
               n_elements = x.size
               eye_tensor = jnp.eye(n_elements).reshape(in_shape + in_shape)
               std_expanded = jnp.sqrt(u)[(None,) * d + (...,)]
               V_tensors = eye_tensor * std_expanded
   
               def ein_jvp(v):
                   return jax.jvp(f, (x,), (v,))[1]
   
               v_func = ein_jvp
               for _ in range(d):
                   v_func = jax.vmap(v_func, in_axes=0, out_axes=0)
               return jnp.sum(v_func(V_tensors)**2, axis=tuple(range(d)))
           else:
               def ein_jvp(v):
                   return jax.jvp(f, (x,), (v,))[
               w_func = ein_jvp
               for _ in range(d):
                   w_func = jax.vmap(w_func, in_axes=-1, out_axes=-1)
               W = w_func(u)
   
               y, vjp_fun = jax.vjp(f, x)
               if diag:
                   v_func = vjp_fun
                   for _ in range(len(y.shape)):
                       v_func = jax.vmap(v_func, in_axes=-1, out_axes=-1)
                   return jnp.sum(v_func(W), axis=tuple(range(d)))
               else:
                   v_func = vjp_fun
                   for _ in range(d):
                       v_func = jax.vmap(v_func, in_axes=0, out_axes=0)
                   return v_func(W)
       return _lpu
   
   @jax.jit(static_argnums=(1, 4), skipped_argnums=(0,))
   def lpu_batched_jvp(f: Callable, d: int, batch_x: Array, batch_u: Array, diag: bool = False) -> Array:
       return jax.vmap(make_lpu_functional(f, d, diag), in_axes=(0, 0))(batch_x, batch_u)
10. heptaflar commented on Oct 4, 2026

    

    heptaflar

    on Oct 4, 2026

    CollaboratorAuthor

    More actions

    class ToF(F):
        r"""
        Adapts a pure function
    
        .. math::
            f: \(\mathbb{R}^{m} \to \mathbb{R}^{n}, \quad         (x) \mapsto\) f(x)
    
        where :math:`m, n` are shapes (natural numbers or tuples
        of natural numbers) to the function interface ``F``.
        """
    
        def __init__(
            self, f: Callable[[Array], Array], rev: bool = True, jit: bool = True
        ):
            """
            Creates a new instance of this class.
    
            :param f: The function :math:`f`.
            :param rev: Use reverse mode for the Jacobian.
            :param jit: Switches JIT compilation on and off (for debugging).
            """
            self._f = jax.jit(f) if jit else f
            self._rev = rev
            self._jit = jit
    
        def eval(self, x: np.ndarray) -> np.ndarray:
            x_ = jnp.asarray(x)
            y_ = vec(self._f, x_) if self._jit else vec_no_jit(self._f, x_)
            return np.asarray(y_)
    
        def jac(self, x: np.ndarray) -> np.ndarray:
            x_ = jnp.asarray(x)
            g_ = (
                jac(self._f, x_, self._rev)
                if self._jit
                else jac_no_jit(self._f, x_, self._rev)
            )
            return np.asarray(g_)
    
        def lpu(
            self, x: np.ndarray, u: np.ndarray, diag: bool = False, jvp: bool = False
        ) -> np.ndarray:
            """
            The law of propagation of uncertainty.
    
            :param x: Primal input array (including batch dimension).
            :param u: Uncertainty/Covariance array.
            :param diag: To return only variance elements.
            :param jvp: If True, uses the functional JVP/VJP path without explicit Jacobians.
            """
            x_ = jnp.asarray(x)
            u_ = jnp.asarray(u)
            d = x_.ndim - 1
    
            if jvp:
                # Nutzt den funktionalen Least-Knowledge-Pfad direkt mit self._f
                # (Das JIT-Verhalten wird über die globale lpu-Schnittstelle oder self._f gesteuert)
                v_ = lpu_batched_jvp(self._f, d, x_, u_, diag)
            else:
                # Ihre bisherige matrix-basierte Variante
                g_ = (
                    jac(self._f, x_, self._rev)
                    if self._jit
                    else jac_no_jit(self._f, x_, self._rev)
                )
                v_ = lpu_matrix_batched(d, g_, u_, diag)
                
            return np.asarray(v_)
    
        @property
        def f(self) -> Callable[[Array], Array]:
            return self._f