scan added

keras/backend/jax/core.py

@@ -263,6 +263,11 @@ def slice_update(inputs, start_indices, updates):
     return jax.lax.dynamic_update_slice(inputs, updates, start_indices)
 
 
+def scan(f, init, xs, length=None, reverse=False, unroll=1):
+
+    return jax.lax.scan(f, init, xs, length, reverse, unroll)
+
+
 def while_loop(
     cond,
     body,

keras/backend/numpy/core.py

@@ -204,6 +204,22 @@ def while_loop(
     return loop_vars
 
 
+def scan(f, init, xs, length=None, reverse=False, unroll=1):
+
+    if xs is None:
+        xs = [None] * length
+    carry = init
+    ys = []
+    if reverse:
+        xs.reverse()
+    for x in xs:
+        carry, y = f(carry, x)
+        ys.append(y)
+    if reverse:
+        ys.reverse()
+    return carry, np.stack(ys)
+
+
 def fori_loop(lower, upper, body_fun, init_val):
     val = init_val
     for i in range(lower, upper):

keras/backend/tensorflow/core.py

@@ -194,6 +194,25 @@ def slice_update(inputs, start_indices, updates):
     return dynamic_update_slice(inputs, updates, start_indices)
 
 
+def scan(f, init, xs, length=None, reverse=False, unroll=1):
+
+    if xs is None:
+        xs = [None] * length
+    if reverse:
+        np.flip(xs)
+
+    init = (init, np.array(0, dtype=init.dtype))
+
+    carry, ys = tf.scan(f, xs, initializer=init)
+
+    ys = ys.numpy()
+    carry = carry.numpy()
+    if reverse:
+        np.flip(ys)
+
+    return carry[0], ys
+
+
 def while_loop(
     cond,
     body,

keras/backend/torch/core.py

@@ -335,6 +335,21 @@ def slice_update(inputs, start_indices, updates):
     return outputs
 
 
+def scan(f, init, xs, length=None, reverse=False, unroll=1):
+    if xs is None:
+        xs = [None] * length
+    carry = init
+    ys = []
+    if reverse:
+        xs.reverse()
+    for x in xs:
+        carry, y = f(carry, x)
+        ys.append(y)
+    if reverse:
+        ys.reverse()
+    return carry, np.stack(ys)
+
+
 def while_loop(
     cond,
     body,

keras/ops/core.py

@@ -168,6 +168,14 @@ def compute_output_spec(self, inputs, start_indices, updates):
         return KerasTensor(inputs.shape, dtype=inputs.dtype)
 
 
+class Scan(Operation):
+    def call(f, init, xs, length, reverse, unroll):
+        return backend.core.scan(f, init, xs, length, reverse, unroll)
+
+    def compute_output_spec(self, f, init, xs, length, reverse, unroll):
+        return KerasTensor(xs.shape, dtype=xs.dtype)
+
+
 @keras_export("keras.ops.slice_update")
 def slice_update(inputs, start_indices, updates):
     """Update an input by slicing in a tensor of updated values.
@@ -204,6 +212,46 @@ def slice_update(inputs, start_indices, updates):
     return backend.core.slice_update(inputs, start_indices, updates)
 
 
+@keras_export("keras.ops.scan")
+def scan(f, init, xs, length=None, reverse=False, unroll=1):
+    """Scan a function over leading array axes while carrying along state.
+
+    At a high level, this operation does
+    carry, y = f(carry, x) and adds the y value to an array which is
+    returned at the end along with the last carry.The x is taken from
+    the array of xs. The initial state of the carry can be set by using
+    the init argument.In the case of None argument for xs a None  array
+    will be initialized by using the length argument like this
+    xs = [None]*length.Example:
+
+    ```python
+    def f(carry, x):
+        x += 1
+        carry = x
+        return carry, x
+
+    inputs = keras.ops.scan(f,1,[0,1,2,3],length=None,reverse=False,unroll=1)
+    ```
+
+    Args:
+        f: The function that will be used to scan over the array xs.
+        init: The initial state of the carry argument for the scan function.
+        xs: The array that will be scanned over.
+        length: The length of the None xs array in the case of None xs argument.
+        reverse: If set to true the xs will be reversed at the start and
+        the ys array will be reversed at the end.
+        unroll: Optional positive int specifying, in the underlying operation of
+         the scan primitive, how many scan iterations to unroll within a single
+         iteration of a loop(Supported only on jax backend).
+
+    Returns:
+        A scanned array and a carry element.
+    """
+    if any_symbolic_tensors((init, xs)):
+        return SliceUpdate().symbolic_call(f, init, xs, length, reverse, unroll)
+    return backend.core.scan(f, init, xs, length, reverse, unroll)
+
+
 class WhileLoop(Operation):
     def __init__(self, cond, body, maximum_iterations):
         super().__init__()

keras/ops/core_test.py

@@ -1,3 +1,4 @@
+import jax
 import numpy as np
 import pytest
 
@@ -215,6 +216,48 @@ def test_slice_update(self):
         outputs = core.slice_update(inputs, start_indices, updates)
         self.assertAllClose(outputs[1:3, 1:3, 2:4, 2:4], np.zeros([2, 2, 2, 2]))
 
+    def test_scan(self):
+        def f(carry, x):
+            if type(carry) is list or type(carry) is tuple:
+                x += carry[0]
+                carry = carry[0]
+            else:
+
+                x += carry
+            return carry, x
+
+        init_carr = np.array(1)
+        xs = np.array([0, 1, 2, 3, 4, 5, 6])
+
+        carry_op, ys_op = core.scan(
+            f, init_carr, xs, length=len(xs), reverse=False
+        )
+        carry_jax, ys_jax = jax.lax.scan(
+            f, init_carr, xs, length=len(xs), reverse=False
+        )
+
+        ys_op = ys_op.tolist()
+        ys_jax = ys_jax.tolist()
+
+        self.assertEqual(carry_op, carry_jax)
+        self.assertListEqual(ys_op, ys_jax)
+
+        init_carr = np.array(1.1)
+        xs = np.array([0.1, 1.2, 2.3, 3.4, 4.5, 5.6, 6.7])
+
+        carry_op, ys_op = core.scan(
+            f, init_carr, xs, length=len(xs), reverse=True
+        )
+        carry_jax, ys_jax = jax.lax.scan(
+            f, init_carr, xs, length=len(xs), reverse=True
+        )
+
+        ys_op = ys_op.tolist()
+        ys_jax = ys_jax.tolist()
+
+        self.assertEqual(carry_op, carry_jax)
+        self.assertListEqual(ys_op, ys_jax)
+
     def test_while_loop(self):
         def cond(x, y):
             return x[0, 0] < 10

shell/format.sh

@@ -1,5 +1,5 @@
 #!/bin/bash
-set -Eeuo pipefail
+set -Eeu pipefail
 
 base_dir=$(dirname $(dirname $0))