Add jvp op (#21720)

* add jvp op

* add jvp op

* bug fix

* add symbolic call

* fix doc.

Author: pass-lin

keras/api/_tf_keras/keras/ops/__init__.py

@@ -37,6 +37,7 @@
 from keras.src.ops.linalg import eig as eig
 from keras.src.ops.linalg import eigh as eigh
 from keras.src.ops.linalg import inv as inv
+from keras.src.ops.linalg import jvp as jvp
 from keras.src.ops.linalg import lstsq as lstsq
 from keras.src.ops.linalg import lu_factor as lu_factor
 from keras.src.ops.linalg import norm as norm

keras/api/_tf_keras/keras/ops/linalg/__init__.py

@@ -10,6 +10,7 @@
 from keras.src.ops.linalg import eig as eig
 from keras.src.ops.linalg import eigh as eigh
 from keras.src.ops.linalg import inv as inv
+from keras.src.ops.linalg import jvp as jvp
 from keras.src.ops.linalg import lstsq as lstsq
 from keras.src.ops.linalg import lu_factor as lu_factor
 from keras.src.ops.linalg import norm as norm

keras/api/ops/__init__.py

@@ -37,6 +37,7 @@
 from keras.src.ops.linalg import eig as eig
 from keras.src.ops.linalg import eigh as eigh
 from keras.src.ops.linalg import inv as inv
+from keras.src.ops.linalg import jvp as jvp
 from keras.src.ops.linalg import lstsq as lstsq
 from keras.src.ops.linalg import lu_factor as lu_factor
 from keras.src.ops.linalg import norm as norm

keras/api/ops/linalg/__init__.py

@@ -10,6 +10,7 @@
 from keras.src.ops.linalg import eig as eig
 from keras.src.ops.linalg import eigh as eigh
 from keras.src.ops.linalg import inv as inv
+from keras.src.ops.linalg import jvp as jvp
 from keras.src.ops.linalg import lstsq as lstsq
 from keras.src.ops.linalg import lu_factor as lu_factor
 from keras.src.ops.linalg import norm as norm

keras/src/backend/jax/linalg.py

@@ -97,3 +97,7 @@ def lstsq(a, b, rcond=None):
     a = convert_to_tensor(a)
     b = convert_to_tensor(b)
     return jnp.linalg.lstsq(a, b, rcond=rcond)[0]
+
+
+def jvp(fun, primals, tangents, has_aux=False):
+    return jax.jvp(fun, primals, tangents, has_aux=has_aux)

keras/src/backend/numpy/linalg.py

@@ -96,3 +96,7 @@ def lstsq(a, b, rcond=None):
     a = convert_to_tensor(a)
     b = convert_to_tensor(b)
     return np.linalg.lstsq(a, b, rcond=rcond)[0]
+
+
+def jvp(fun, primals, tangents, has_aux=False):
+    raise NotImplementedError("JVP is not supported by the Numpy backend.")

keras/src/backend/openvino/linalg.py

@@ -56,3 +56,7 @@ def svd(x, full_matrices=True, compute_uv=True):
 
 def lstsq(a, b, rcond=None):
     raise NotImplementedError("`lstsq` is not supported with openvino backend")
+
+
+def jvp(fun, primals, tangents, has_aux=False):
+    raise NotImplementedError("`jvp` is not supported with openvino backend")

keras/src/backend/tensorflow/linalg.py

@@ -244,3 +244,27 @@ def lstsq(a, b, rcond=None):
     if b_orig_ndim == 1:
         x = tf.reshape(x, [-1])
     return x
+
+
+def jvp(fun, primals, tangents, has_aux=False):
+    primal_flat = tf.nest.flatten(primals)
+    tangent_flat = tf.nest.flatten(tangents)
+
+    tangent_flat = [
+        tf.cast(t, p.dtype) for t, p in zip(tangent_flat, primal_flat)
+    ]
+
+    with tf.autodiff.ForwardAccumulator(primal_flat, tangent_flat) as acc:
+        if has_aux:
+            primals_out, aux = fun(*primals)
+        else:
+            primals_out = fun(*primals)
+
+        primals_out_flat = tf.nest.flatten(primals_out)
+        tangents_out_flat = [acc.jvp(po) for po in primals_out_flat]
+
+    tangents_out = tf.nest.pack_sequence_as(primals_out, tangents_out_flat)
+
+    if has_aux:
+        return primals_out, tangents_out, aux
+    return primals_out, tangents_out

keras/src/backend/torch/linalg.py

@@ -80,3 +80,7 @@ def lstsq(a, b, rcond=None):
     a = convert_to_tensor(a)
     b = convert_to_tensor(b)
     return torch.linalg.lstsq(a, b, rcond=rcond)[0]
+
+
+def jvp(fun, primals, tangents, has_aux=False):
+    return torch.func.jvp(fun, primals, tangents, has_aux=has_aux)

keras/src/ops/linalg.py

@@ -1,4 +1,5 @@
 from keras.src import backend
+from keras.src import tree
 from keras.src.api_export import keras_export
 from keras.src.backend import KerasTensor
 from keras.src.backend import any_symbolic_tensors
@@ -732,3 +733,95 @@ def _assert_a_b_compat(a, b):
                 "Expected `a.shape[-1] == b.shape[-1]`. "
                 f"Received: a.shape={a.shape}, b.shape={b.shape}"
             )
+
+
+class JVP(Operation):
+    def __init__(self, has_aux=False, *, name=None):
+        super().__init__(name=name)
+        self.has_aux = has_aux
+
+    def call(self, fun, primals, tangents):
+        """Computes the JVP of `fun` at `primals` along `tangents`.
+
+        Args:
+            fun: A callable that takes tensors (or nested structures) as input
+                 and returns a tensor (or nested structure) as output.
+            primals: Input tensors (or nested structures) at which the Jacobian
+                     of `fun` is evaluated.
+            tangents: Tensors (or nested structures) representing the direction
+                      vectors for the JVP. Must have the same structure as
+                      `primals`.
+
+        Returns:
+            If `has_aux` is False:
+                A tuple (primals_out, tangents_out) where:
+                - primals_out: Output of `fun(*primals)`
+                - tangents_out: JVP of `fun` at `primals` along `tangents`
+            If `has_aux` is True:
+                A tuple (primals_out, tangents_out, aux) where:
+                - aux: Auxiliary data returned by `fun`
+        """
+        return backend.linalg.jvp(fun, primals, tangents, has_aux=self.has_aux)
+
+    def compute_output_spec(self, fun, primals, tangents):
+        # Infer primal output spec
+        if self.has_aux:
+            primals_out_spec, aux_spec = backend.compute_output_spec(
+                fun, *primals
+            )
+        else:
+            primals_out_spec = backend.compute_output_spec(fun, *primals)
+
+        # Tangents output should match primals output in structure and shape
+        tangents_out_spec = tree.map_structure(
+            lambda x: KerasTensor(x.shape, x.dtype), primals_out_spec
+        )
+
+        if self.has_aux:
+            return primals_out_spec, tangents_out_spec, aux_spec
+        return primals_out_spec, tangents_out_spec
+
+
+@keras_export(["keras.ops.jvp", "keras.ops.linalg.jvp"])
+def jvp(fun, primals, tangents, has_aux=False):
+    """Computes a (forward-mode) Jacobian-vector product of `fun`.
+    Args:
+        fun: Function to be differentiated. Its arguments should be arrays,
+            scalars, or standard Python containers of arrays or scalars. It
+            should return an array, scalar, or standard Python container of
+            arrays or scalars.
+        primals: The primal values at which the Jacobian of `fun` should be
+                evaluated. Should be either a tuple or a list of arguments,
+                and its length should be equal to the number of positional
+                parameters of `fun`.
+        tangents: The tangent vector for which the Jacobian-vector product
+                should be evaluated. Should be either a tuple or a list of
+                tangents, with the same tree structure and array shapes as
+                `primals`.
+        has_aux: Optional, bool. Indicates whether `fun` returns a pair where
+                the first element is considered the output of the mathematical
+                function to be differentiated and the second element is
+                auxiliary data. Default is False.
+
+    Returns:
+        If `has_aux` is False, returns a (`primals_out`, `tangents_out`) pair,
+        where `primals_out` is `fun(*primals)`, and `tangents_out` is the
+        Jacobian-vector product of `fun` evaluated at `primals` with
+        `tangents`. The `tangents_out` value has the same Python tree
+        structure and shapes as `primals_out`.
+
+        If `has_aux` is True, returns a (`primals_out`, `tangents_out`, `aux`)
+        tuple where `aux` is the auxiliary data returned by `fun`.
+
+    Example:
+    >>> from keras import ops
+    >>> a1, a2 = ops.convert_to_tensor(0.1), ops.convert_to_tensor(0.2)
+    >>> primals, tangents = ops.jvp(ops.sin, (a1,), (a2,))
+    >>> primals
+    0.09983342
+    >>> tangents
+    0.19900084
+    """
+    if any_symbolic_tensors((primals, tangents)):
+        return JVP(has_aux=has_aux).symbolic_call(fun, primals, tangents)
+    return backend.linalg.jvp(fun, primals, tangents, has_aux=has_aux)