Add files using upload-large-folder tool

SwarmUI/dlbackend/ComfyUI/venv/lib/python3.10/site-packages/keras/src/ops/__init__.py

@@ -0,0 +1,16 @@
+# from keras.src.ops.numpy import Matmul, matmul
+# from keras.src.ops.numpy import Add, add
+# from keras.src.ops.numpy import Multiply, multiply
+
+from keras.src.backend import cast
+from keras.src.backend import cond
+from keras.src.backend import is_tensor
+from keras.src.backend import name_scope
+from keras.src.backend import random
+from keras.src.ops import image
+from keras.src.ops import operation_utils
+from keras.src.ops.core import *  # noqa: F403
+from keras.src.ops.linalg import *  # noqa: F403
+from keras.src.ops.math import *  # noqa: F403
+from keras.src.ops.nn import *  # noqa: F403
+from keras.src.ops.numpy import *  # noqa: F403

SwarmUI/dlbackend/ComfyUI/venv/lib/python3.10/site-packages/keras/src/ops/core.py

@@ -0,0 +1,1167 @@
+import ml_dtypes
+import numpy as np
+
+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
+from keras.src.backend.common.backend_utils import slice_along_axis
+from keras.src.ops.operation import Operation
+from keras.src.utils import traceback_utils
+
+
+class Map(Operation):
+    def __init__(self):
+        super().__init__()
+
+    def call(self, f, xs):
+        return backend.core.map(f, xs)
+
+    def compute_output_spec(self, f, xs):
+        x = xs[0]
+        n = xs.shape[0]
+        y = backend.compute_output_spec(f, x)
+
+        def append_batch_axis(x):
+            return KerasTensor(
+                shape=(n,) + x.shape, dtype=x.dtype, sparse=x.sparse
+            )
+
+        y = tree.map_structure(append_batch_axis, y)
+        return y
+
+
+@keras_export("keras.ops.map")
+def map(f, xs):
+    """Map a function over leading array axes.
+
+    Like Python’s builtin map, except inputs and outputs are in the form of
+    stacked arrays. Consider using the `vectorized_map()` transform instead,
+    unless you need to apply a function element by element for reduced memory
+    usage or heterogeneous computation with other control flow primitives.
+
+    When `xs` is an array type, the semantics of `map()` are given by this
+    Python implementation:
+
+    ```python
+    def map(f, xs):
+        return np.stack([f(x) for x in xs])
+    ```
+
+    Args:
+        f: Callable defines the function to apply element-wise over the first
+            axis or axes of `xs`.
+        xs: Values over which to map along the leading axis.
+
+    Returns:
+        Mapped values.
+
+    Examples:
+
+    >>> f = lambda x: x**2
+    >>> xs = keras.ops.arange(10)
+    >>> ys = keras.ops.map(f, xs)
+    >>> ys
+    [0, 1, 4, 9, 16, 25, 36, 49, 64, 81]
+
+    >>> f = lambda x: {"y1": x**2, "y2": x * 10}  # Can have nested outputs
+    >>> ys = keras.ops.map(f, xs)
+    >>> ys["y1"]
+    [0, 1, 4, 9, 16, 25, 36, 49, 64, 81]
+    >>> ys["y2"]
+    [0, 10, 20, 30, 40, 50, 60, 70, 80, 90]
+    """
+    if any_symbolic_tensors((xs,)):
+        return Map().symbolic_call(f, xs)
+    return backend.core.map(f, xs)
+
+
+class Scan(Operation):
+    def __init__(self, reverse=False, unroll=1):
+        super().__init__()
+        self.reverse = reverse
+        self.unroll = unroll
+
+    def call(self, f, init, xs, length):
+        return backend.core.scan(
+            f, init, xs, length, reverse=self.reverse, unroll=self.unroll
+        )
+
+    def compute_output_spec(self, f, init, xs, length):
+        if xs is None:
+            n = int(length)
+            x = None
+        else:
+            n = (
+                int(length)
+                if length is not None
+                else tree.flatten(xs)[0].shape[0]
+            )
+            x = xs[0]
+
+        carry, y = backend.compute_output_spec(f, init, x)
+        y = KerasTensor(shape=(n,) + y.shape, dtype=y.dtype, sparse=y.sparse)
+        return carry, y
+
+
+@keras_export("keras.ops.scan")
+def scan(f, init, xs=None, length=None, reverse=False, unroll=1):
+    """Scan a function over leading array axes while carrying along state.
+
+    When the type of `xs` is an array type or `None`, and the type of `ys` is an
+    array type, the semantics of `scan()` are given roughly by this Python
+    implementation:
+
+    ```python
+    def scan(f, init, xs, length=None):
+        if xs is None:
+            xs = [None] * length
+        carry = init
+        ys = []
+        for x in xs:
+            carry, y = f(carry, x)
+            ys.append(y)
+        return carry, np.stack(ys)
+    ```
+
+    The loop-carried value `carry` (`init`) must hold a fixed shape and dtype
+    across all iterations.
+
+    In TensorFlow, `y` must match `carry` in shape and dtype. This is not
+    required in other backends.
+
+    Args:
+        f: Callable defines the logic for each loop iteration. This accepts two
+            arguments where the first is a value of the loop carry and the
+            second is a slice of `xs` along its leading axis.
+            This callable returns a pair where the first represents a new value
+            for the loop carry and the second represents a slice of the output.
+        init: The initial loop carry value. This can be a scalar, tensor, or any
+            nested structure. It must match the structure of the first element
+            returned by `f`.
+        xs: Optional value to scan along its leading axis. This can be a tensor
+            or any nested structure. If `xs` is not provided, you must specify
+            `length` to define the number of loop iterations.
+            Defaults to `None`.
+        length: Optional integer specifying the number of loop iterations.
+            If `length` is not provided, it defaults to the sizes of leading
+            axis of the arrays in `xs`. Defaults to `None`.
+        reverse: Optional boolean specifying whether to run the scan iteration
+            forward or in reverse, equivalent to reversing the leading axes of
+            the arrays in both `xs` and in `ys`.
+        unroll: Optional positive integer or boolean specifying how many scan
+            iterations to unroll within a single iteration of a loop. If an
+            integer is provided, it determines how many unrolled loop iterations
+            to run within a single rolled iteration of the loop. If a boolean is
+            provided, it will determine if the loop is completely unrolled
+            (`unroll=True`) or left completely unrolled (`unroll=False`).
+            Note that unrolling is only supported by JAX and TensorFlow
+            backends.
+
+    Returns:
+        A pair where the first element represents the final loop carry value and
+        the second element represents the stacked outputs of `f` when scanned
+        over the leading axis of the inputs.
+
+    Examples:
+
+    >>> sum_fn = lambda c, x: (c + x, c + x)
+    >>> init = keras.ops.array(0)
+    >>> xs = keras.ops.array([1, 2, 3, 4, 5])
+    >>> carry, result = keras.ops.scan(sum_fn, init, xs)
+    >>> carry
+    15
+    >>> result
+    [1, 3, 6, 10, 15]
+    """
+    if any_symbolic_tensors((init, xs)):
+        return Scan(reverse=reverse, unroll=unroll).symbolic_call(
+            f, init, xs, length
+        )
+    return backend.core.scan(
+        f, init, xs, length, reverse=reverse, unroll=unroll
+    )
+
+
+class AssociativeScan(Operation):
+    def __init__(self, reverse=False):
+        super().__init__()
+        self.reverse = reverse
+
+    def call(self, f, elems, axis=0):
+        return backend.core.associative_scan(
+            f, elems, reverse=self.reverse, axis=axis
+        )
+
+    def compute_output_spec(self, f, elems, axis):
+        elems_flat = tree.flatten(elems)
+        lens = [elem.shape[axis] for elem in elems_flat]
+        if len(set(lens)) != 1:
+            raise ValueError(
+                "Array inputs to associative_scan must have the same "
+                "first dimension. (saw: {})".format(
+                    [elem.shape for elem in elems_flat]
+                )
+            )
+
+        x = tree.pack_sequence_as(
+            elems, [slice_along_axis(x, 0, 1, axis=axis) for x in elems_flat]
+        )
+        y_spec = backend.compute_output_spec(f, x, x)
+
+        def _restore_shape(x):
+            return KerasTensor(
+                shape=elems_flat[0].shape, dtype=x.dtype, sparse=x.sparse
+            )
+
+        y_spec = tree.map_structure(_restore_shape, y_spec)
+        return y_spec
+
+
+@keras_export("keras.ops.associative_scan")
+def associative_scan(f, elems, reverse=False, axis=0):
+    """Performs a scan with an associative binary operation, in parallel.
+
+    This operation his similar to `scan`, with the key difference that
+    `associative_scan` is a parallel implementation with
+    potentially significant performance benefits, especially when jit compiled.
+    The catch is that it can only be used when `f` is a binary associative
+    operation (i.e. it must verify `f(a, f(b, c)) == f(f(a, b), c)`).
+
+    For an introduction to associative scans, refer to this paper:
+    Blelloch, Guy E. 1990.
+    [Prefix Sums and Their Applications](
+        https://www.cs.cmu.edu/~guyb/papers/Ble93.pdf).
+
+    Args:
+        f: A Python callable implementing an associative binary operation with
+            signature `r = f(a, b)`. Function `f` must be associative, i.e.,
+            it must satisfy the equation
+            `f(a, f(b, c)) == f(f(a, b), c)`.
+            The inputs and result are (possibly nested Python tree structures
+            of) array(s) matching `elems`. Each array has a dimension in place
+            of the `axis` dimension. `f` should be applied elementwise over
+            the `axis` dimension.
+            The result `r` has the same shape (and structure) as the
+            two inputs `a` and `b`.
+        elems: A (possibly nested Python tree structure of) array(s), each with
+            an `axis` dimension of size `num_elems`.
+        reverse: A boolean stating if the scan should be reversed with respect
+            to the `axis` dimension.
+        axis: an integer identifying the axis over which the scan should occur.
+
+    Returns:
+        A (possibly nested Python tree structure of) array(s) of the same shape
+        and structure as `elems`, in which the `k`'th element of `axis` is
+        the result of recursively applying `f` to combine the first `k`
+        elements of `elems` along `axis`. For example, given
+        `elems = [a, b, c, ...]`, the result would be
+        `[a, f(a, b), f(f(a, b), c), ...]`.
+
+    Examples:
+
+    >>> sum_fn = lambda x, y: x + y
+    >>> xs = keras.ops.arange(5)
+    >>> ys = keras.ops.associative_scan(sum_fn, xs, axis=0)
+    >>> ys
+    [0, 1, 3, 6, 10]
+
+    >>> sum_fn = lambda x, y: [x[0] + y[0], x[1] + y[1], x[2] + y[2]]
+    >>> xs = [keras.ops.array([[1, 2]]) for _ in range(3)]
+    >>> ys = keras.ops.associative_scan(sum_fn, xs, axis=0)
+    >>> ys
+    [[1, 3], [1, 3], [1, 3]]
+    """
+    if any_symbolic_tensors((elems,)):
+        return AssociativeScan(reverse=reverse).symbolic_call(f, elems, axis)
+    return backend.core.associative_scan(f, elems, reverse=reverse, axis=axis)
+
+
+class Scatter(Operation):
+    def call(self, indices, values, shape):
+        return backend.core.scatter(indices, values, shape)
+
+    def compute_output_spec(self, indices, values, shape):
+        return KerasTensor(shape, dtype=values.dtype)
+
+
+@keras_export("keras.ops.scatter")
+def scatter(indices, values, shape):
+    """Returns a tensor of shape `shape` where `indices` are set to `values`.
+
+    At a high level, this operation does `zeros[indices] = updates` and
+    returns the output. It is equivalent to:
+
+    ```python
+    zeros = keras.ops.zeros(shape)
+    output = keras.ops.scatter_update(zeros, indices, values)
+    ```
+
+    Args:
+        indices: A tensor or list/tuple specifying
+            indices for the values in `values`.
+        values: A tensor, the values to be set at `indices`.
+        shape: Shape of the output tensor.
+
+    Example:
+
+    >>> indices = [[0, 1], [1, 1]]
+    >>> values = np.array([1., 1.])
+    >>> keras.ops.scatter(indices, values, shape=(2, 2))
+    array([[0., 1.],
+           [0., 1.]])
+    """
+    if any_symbolic_tensors((indices, values, shape)):
+        return Scatter().symbolic_call(indices, values, shape)
+    return backend.core.scatter(indices, values, shape)
+
+
+class ScatterUpdate(Operation):
+    def call(self, inputs, indices, updates):
+        return backend.core.scatter_update(inputs, indices, updates)
+
+    def compute_output_spec(self, inputs, indices, updates):
+        return KerasTensor(inputs.shape, dtype=inputs.dtype)
+
+
+@keras_export("keras.ops.scatter_update")
+def scatter_update(inputs, indices, updates):
+    """Update inputs via updates at scattered (sparse) indices.
+
+    At a high level, this operation does `inputs[indices] = updates`.
+    Assume `inputs` is a tensor of shape `(D0, D1, ..., Dn)`, there are 2 main
+    usages of `scatter_update`.
+
+    1. `indices` is a 2D tensor of shape `(num_updates, n)`, where `num_updates`
+        is the number of updates to perform, and `updates` is a 1D tensor of
+        shape `(num_updates,)`. For example, if `inputs` is `zeros((4, 4, 4))`,
+        and we want to update `inputs[1, 2, 3]` and `inputs[0, 1, 3]` as 1, then
+        we can use:
+
+    ```python
+    inputs = np.zeros((4, 4, 4))
+    indices = [[1, 2, 3], [0, 1, 3]]
+    updates = np.array([1., 1.])
+    inputs = keras.ops.scatter_update(inputs, indices, updates)
+    ```
+
+    2 `indices` is a 2D tensor of shape `(num_updates, k)`, where `num_updates`
+        is the number of updates to perform, and `k` (`k < n`) is the size of
+        each index in `indices`. `updates` is a `n - k`-D tensor of shape
+        `(num_updates, inputs.shape[k:])`. For example, if
+        `inputs = np.zeros((4, 4, 4))`, and we want to update `inputs[1, 2, :]`
+        and `inputs[2, 3, :]` as `[1, 1, 1, 1]`, then `indices` would have shape
+        `(num_updates, 2)` (`k = 2`), and `updates` would have shape
+        `(num_updates, 4)` (`inputs.shape[2:] = 4`). See the code below:
+
+    ```python
+    inputs = np.zeros((4, 4, 4))
+    indices = [[1, 2], [2, 3]]
+    updates = np.array([[1., 1., 1, 1,], [1., 1., 1, 1,])
+    inputs = keras.ops.scatter_update(inputs, indices, updates)
+    ```
+
+    Args:
+        inputs: A tensor, the tensor to be updated.
+        indices: A tensor or list/tuple of shape `(N, inputs.ndim)`, specifying
+            indices to update. `N` is the number of indices to update, must be
+            equal to the first dimension of `updates`.
+        updates: A tensor, the new values to be put to `inputs` at `indices`.
+
+    Returns:
+        A tensor, has the same shape and dtype as `inputs`.
+    """
+    if any_symbolic_tensors((inputs, indices, updates)):
+        return ScatterUpdate().symbolic_call(inputs, indices, updates)
+    return backend.core.scatter_update(inputs, indices, updates)
+
+
+class Slice(Operation):
+    def call(self, inputs, start_indices, shape):
+        return backend.core.slice(inputs, start_indices, shape)
+
+    def compute_output_spec(self, inputs, start_indices, shape):
+        return KerasTensor(shape, dtype=inputs.dtype)
+
+
+@keras_export("keras.ops.slice")
+def slice(inputs, start_indices, shape):
+    """Return a slice of an input tensor.
+
+    At a high level, this operation is an explicit replacement for array slicing
+    e.g. `inputs[start_indices: start_indices + shape]`.
+    Unlike slicing via brackets, this operation will accept tensor start
+    indices on all backends, which is useful when indices dynamically computed
+    via other tensor operations.
+
+    ```python
+    inputs = np.zeros((5, 5))
+    start_indices = np.array([3, 3])
+    shape = np.array([2, 2])
+    inputs = keras.ops.slice(inputs, start_indices, shape)
+    ```
+
+    Args:
+        inputs: A tensor, the tensor to be updated.
+        start_indices: A list/tuple of shape `(inputs.ndim,)`, specifying
+            the starting indices for updating.
+        shape: The full shape of the returned slice.
+
+    Returns:
+        A tensor, has the same shape and dtype as `inputs`.
+    """
+    if any_symbolic_tensors((inputs, start_indices, shape)):
+        return Slice().symbolic_call(inputs, start_indices, shape)
+    return backend.core.slice(inputs, start_indices, shape)
+
+
+class SliceUpdate(Operation):
+    def call(self, inputs, start_indices, updates):
+        return backend.core.slice_update(inputs, start_indices, updates)
+
+    def compute_output_spec(self, inputs, start_indices, updates):
+        return KerasTensor(inputs.shape, dtype=inputs.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.
+
+    At a high level, this operation does
+    `inputs[start_indices: start_indices + updates.shape] = updates`.
+    Assume inputs is a tensor of shape `(D0, D1, ..., Dn)`,
+    `start_indices` must be a list/tuple of n integers, specifying the starting
+    indices. `updates` must have the same rank as `inputs`, and the size of each
+    dim must not exceed `Di - start_indices[i]`. For example, if we have 2D
+    inputs `inputs = np.zeros((5, 5))`, and we want to update the intersection
+    of last 2 rows and last 2 columns as 1, i.e.,
+    `inputs[3:, 3:] = np.ones((2, 2))`, then we can use the code below:
+
+    ```python
+    inputs = np.zeros((5, 5))
+    start_indices = [3, 3]
+    updates = np.ones((2, 2))
+    inputs = keras.ops.slice_update(inputs, start_indices, updates)
+    ```
+
+    Args:
+        inputs: A tensor, the tensor to be updated.
+        start_indices: A list/tuple of shape `(inputs.ndim,)`, specifying
+            the starting indices for updating.
+        updates: A tensor, the new values to be put to `inputs` at `indices`.
+            `updates` must have the same rank as `inputs`.
+
+    Returns:
+        A tensor, has the same shape and dtype as `inputs`.
+    """
+    if any_symbolic_tensors((inputs, start_indices, updates)):
+        return SliceUpdate().symbolic_call(inputs, start_indices, updates)
+    return backend.core.slice_update(inputs, start_indices, updates)
+
+
+class Switch(Operation):
+    def call(self, index, branches, *operands):
+        return backend.core.switch(index, branches, *operands)
+
+    def compute_output_spec(self, index, branches, *operands):
+        # We use first branch for output_spec
+        spec = backend.compute_output_spec(branches[0], *operands)
+        return spec
+
+
+@keras_export("keras.ops.switch")
+def switch(index, branches, *operands):
+    """Apply exactly one of the `branches` given by `index`.
+
+    If `index` is out of bounds, it is clamped to within bounds.
+
+    The semantics of `switch` are given roughly by this Python implementation:
+
+    ```python
+    def switch(index, branches, *operands):
+        index = clamp(0, index, len(branches) - 1)
+        return branches[index](*operands)
+    ```
+
+    Args:
+        index: An integer scalar indicating which branch function to apply.
+        branches: A sequence of functions to be applied based on `index`.
+        operands: Inputs to whichever branch is applied.
+
+    Returns:
+        The outputs of `branch(*operands)` for the branch that was selected
+        based on `index`.
+
+    Examples:
+
+    >>> add_fn = lambda x, y: x + y
+    >>> subtract_fn = lambda x, y: x - y
+    >>> x = keras.ops.array(2.0)
+    >>> y = keras.ops.array(0.5)
+    >>> branches = [add_fn, subtract_fn]
+    >>> keras.ops.switch(0, branches, x, y)
+    2.5
+
+    >>> keras.ops.switch(1, branches, x, y)
+    1.5
+    """
+    if any_symbolic_tensors(operands):
+        return Switch().symbolic_call(index, branches, *operands)
+    return backend.core.switch(index, branches, *operands)
+
+
+class WhileLoop(Operation):
+    def __init__(self, cond, body, maximum_iterations):
+        super().__init__()
+        self.cond = cond
+        self.body = body
+        self.maximum_iterations = maximum_iterations
+
+    def call(self, loop_vars):
+        return backend.core.while_loop(
+            self.cond,
+            self.body,
+            loop_vars,
+            maximum_iterations=self.maximum_iterations,
+        )
+
+    def compute_output_spec(self, loop_vars):
+        return [KerasTensor(v.shape, dtype=v.dtype) for v in loop_vars]
+
+
+@keras_export("keras.ops.while_loop")
+def while_loop(
+    cond,
+    body,
+    loop_vars,
+    maximum_iterations=None,
+):
+    """While loop implementation.
+
+    Args:
+        cond: A callable that represents the termination condition of the loop.
+            Must accept a `loop_vars` like structure as an argument. If
+            `loop_vars` is a tuple or list, each element of `loop_vars` will be
+            passed positionally to the callable.
+        body: A callable that represents the loop body. Must accept a
+            `loop_vars` like structure as an argument, and return update value
+            with the same structure. If `loop_vars` is a tuple or list, each
+            element of `loop_vars` will be passed positionally to the callable.
+        loop_vars: An arbitrary nested structure of tensor state to persist
+            across loop iterations.
+        maximum_iterations: Optional maximum number of iterations of the while
+            loop to run. If provided, the `cond` output is AND-ed with an
+            additional condition ensuring the number of iterations executed is
+            no greater than `maximum_iterations`.
+
+    Returns:
+        A list/tuple of tensors, has the same shape and dtype as `inputs`.
+
+    Examples:
+
+    >>> i = 0
+    >>> cond = lambda i: i < 10
+    >>> body = lambda i: i + 1
+    >>> keras.ops.while_loop(cond, body, i)
+    10
+
+    >>> x, y = 0, 1
+    >>> cond = lambda x, y: x < 10
+    >>> body = lambda x, y: (x + 1, y + 1)
+    >>> keras.ops.while_loop(cond, body, (x, y))
+    10, 11
+    """
+    return backend.core.while_loop(
+        cond,
+        body,
+        loop_vars,
+        maximum_iterations=maximum_iterations,
+    )
+
+
+class StopGradient(Operation):
+    def __init__(self):
+        super().__init__()
+
+    def call(self, variable):
+        return backend.core.stop_gradient(variable)
+
+    def compute_output_spec(self, variable):
+        return KerasTensor(variable.shape, dtype=variable.dtype)
+
+
+@keras_export("keras.ops.stop_gradient")
+def stop_gradient(variable):
+    """Stops gradient computation.
+
+    Args:
+        variable: A tensor variable for which the gradient
+            computation is to be disabled.
+
+    Returns:
+        The variable with gradient computation disabled.
+
+    Examples:
+
+    >>> var = keras.backend.convert_to_tensor(
+    ...     [1., 2., 3.],
+    ...     dtype="float32"
+    ... )
+    >>> var = keras.ops.stop_gradient(var)
+    """
+    if any_symbolic_tensors((variable,)):
+        return StopGradient().symbolic_call(variable)
+    return backend.core.stop_gradient(variable)
+
+
+class ForiLoop(Operation):
+    def __init__(self, lower, upper, body_fun):
+        super().__init__()
+        self.lower = lower
+        self.upper = upper
+        self.body_fun = body_fun
+
+    def call(self, init_val):
+        return backend.core.fori_loop(
+            self.lower,
+            self.upper,
+            self.body_fun,
+            init_val,
+        )
+
+    def compute_output_spec(self, init_val):
+        return KerasTensor(init_val.shape, dtype=init_val.dtype)
+
+
+@keras_export("keras.ops.fori_loop")
+def fori_loop(lower, upper, body_fun, init_val):
+    """For loop implementation.
+
+    Args:
+        lower: The initial value of the loop variable.
+        upper: The upper bound of the loop variable.
+        body_fun: A callable that represents the loop body. Must take two
+            arguments: the loop variable and the loop state. The loop state
+            should be updated and returned by this function.
+        init_val: The initial value of the loop state.
+
+    Returns:
+        The final state after the loop.
+
+    Example:
+
+    >>> lower = 0
+    >>> upper = 10
+    >>> body_fun = lambda i, s: (i + 1, s + i)
+    >>> init_val = 0
+    >>> keras.ops.fori_loop(lower, upper, body_fun, init_val)
+    45
+    """
+    if any_symbolic_tensors((lower, upper, init_val)):
+        return ForiLoop(lower, upper, body_fun).symbolic_call(init_val)
+    return backend.core.fori_loop(lower, upper, body_fun, init_val)
+
+
+class Unstack(Operation):
+    def __init__(self, num=None, axis=0):
+        super().__init__()
+        self.num = num
+        self.axis = axis
+
+    def call(self, x):
+        return backend.core.unstack(x, self.num, self.axis)
+
+    def compute_output_spec(self, x):
+        axis = self.axis
+        if axis < 0:
+            axis = len(x.shape) + axis
+        output_shapes = x.shape[:axis] + x.shape[axis + 1 :]
+        num = self.num
+        if num is None:
+            num = x.shape[axis]
+        if num is None:
+            raise ValueError(
+                "Cannot infer argument `num` from shape "
+                f"{x.shape}. Either provide a tensor with a "
+                "concrete shape in the `axis` dimension or "
+                "explicitly pass the `num` argument."
+            )
+        output = [
+            KerasTensor(shape=output_shapes, dtype=x.dtype) for _ in range(num)
+        ]
+        return output
+
+
+@keras_export("keras.ops.unstack")
+def unstack(x, num=None, axis=0):
+    """Unpacks the given dimension of a rank-R tensor into rank-(R-1) tensors.
+
+    Args:
+        x: The input tensor.
+        num: The length of the dimension axis. Automatically inferred
+            if `None`.
+        axis: The axis along which to unpack.
+
+    Returns:
+        A list of tensors unpacked along the given axis.
+
+    Example:
+
+    >>> x = keras.ops.array([[1, 2], [3, 4]])
+    >>> keras.ops.unstack(x, axis=0)
+    [array([1, 2]), array([3, 4])]
+    """
+    if any_symbolic_tensors((x,)):
+        return Unstack(num, axis).symbolic_call(x)
+    return backend.core.unstack(x, num=num, axis=axis)
+
+
+@keras_export("keras.ops.shape")
+def shape(x):
+    """Gets the shape of the tensor input.
+
+    Note: On the TensorFlow backend, when `x` is a `tf.Tensor` with dynamic
+    shape, dimensions which are dynamic in the context of a compiled function
+    will have a `tf.Tensor` value instead of a static integer value.
+
+    Args:
+        x: A tensor. This function will try to access the `shape` attribute of
+            the input tensor.
+
+    Returns:
+        A tuple of integers or None values, indicating the shape of the input
+            tensor.
+
+    Example:
+
+    >>> x = keras.ops.zeros((8, 12))
+    >>> keras.ops.shape(x)
+    (8, 12)
+    """
+    if any_symbolic_tensors((x,)):
+        return x.shape
+    return backend.core.shape(x)
+
+
+@keras_export("keras.ops.dtype")
+def dtype(x):
+    """Return the dtype of the tensor input as a standardized string.
+
+    Note that due to the standardization, the dtype will not compare equal
+    to the backend-specific version of the dtype.
+
+    Args:
+        x: A tensor. This function will try to access the `dtype` attribute of
+            the input tensor.
+
+    Returns:
+        A string indicating the dtype of the input tensor, e.g. `"float32"`.
+
+    Example:
+
+    >>> x = keras.ops.zeros((8, 12))
+    >>> keras.ops.dtype(x)
+    'float32'
+
+    """
+    return backend.standardize_dtype(x.dtype)
+
+
+class Cast(Operation):
+    def __init__(self, dtype):
+        super().__init__()
+        self.dtype = backend.standardize_dtype(dtype)
+
+    def call(self, x):
+        return backend.core.cast(x, self.dtype)
+
+    def compute_output_spec(self, x):
+        return backend.KerasTensor(shape=x.shape, dtype=self.dtype)
+
+
+@keras_export("keras.ops.cast")
+def cast(x, dtype):
+    """Cast a tensor to the desired dtype.
+
+    Args:
+        x: A tensor or variable.
+        dtype: The target type.
+
+    Returns:
+        A tensor of the specified `dtype`.
+
+    Example:
+
+    >>> x = keras.ops.arange(4)
+    >>> x = keras.ops.cast(x, dtype="float16")
+    """
+    dtype = backend.standardize_dtype(dtype)
+
+    if any_symbolic_tensors((x,)):
+        return Cast(dtype=dtype)(x)
+    return backend.core.cast(x, dtype)
+
+
+class SaturateCast(Operation):
+    def __init__(self, dtype):
+        super().__init__()
+        self.dtype = backend.standardize_dtype(dtype)
+
+    def call(self, x):
+        return _saturate_cast(x, self.dtype)
+
+    def compute_output_spec(self, x):
+        return backend.KerasTensor(shape=x.shape, dtype=self.dtype)
+
+
+@keras_export("keras.ops.saturate_cast")
+def saturate_cast(x, dtype):
+    """Performs a safe saturating cast to the desired dtype.
+
+    Saturating cast prevents data type overflow when casting to `dtype` with
+    smaller values range. E.g.
+    `ops.cast(ops.cast([-1, 256], "float32"), "uint8")` returns `[255, 0]`,
+    but `ops.saturate_cast(ops.cast([-1, 256], "float32"), "uint8")` returns
+    `[0, 255]`.
+
+    Args:
+        x: A tensor or variable.
+        dtype: The target type.
+
+    Returns:
+        A safely casted tensor of the specified `dtype`.
+
+    Example:
+
+    Image resizing with bicubic interpolation may produce values outside
+    original range.
+    >>> image2x2 = np.array([0, 1, 254, 255], dtype="uint8").reshape(1, 2, 2, 1)
+    >>> image4x4 = tf.image.resize(image2x2, (4, 4), method="bicubic")
+    >>> print(image4x4.numpy().squeeze())
+    >>> # [[-22.500004 -22.204624 -21.618908 -21.32353 ]
+    >>> #  [ 52.526054  52.82143   53.407146  53.70253 ]
+    >>> #  [201.29752  201.59288  202.17859  202.47395 ]
+    >>> #  [276.32355  276.61893  277.20465  277.50006 ]]
+
+    Casting this resized image back to `uint8` will cause overflow.
+    >>> image4x4_casted = ops.cast(image4x4, "uint8")
+    >>> print(image4x4_casted.numpy().squeeze())
+    >>> # [[234 234 235 235]
+    >>> #  [ 52  52  53  53]
+    >>> #  [201 201 202 202]
+    >>> #  [ 20  20  21  21]]
+
+    Saturate casting to `uint8` will clip values to `uint8` range before
+    casting and will not cause overflow.
+    >>> image4x4_saturate_casted = ops.saturate_cast(image4x4, "uint8")
+    >>> print(image4x4_saturate_casted.numpy().squeeze())
+    >>> # [[  0   0   0   0]
+    >>> #  [ 52  52  53  53]
+    >>> #  [201 201 202 202]
+    >>> #  [255 255 255 255]]
+
+    """
+    dtype = backend.standardize_dtype(dtype)
+
+    if any_symbolic_tensors((x,)):
+        return SaturateCast(dtype=dtype)(x)
+    return _saturate_cast(x, dtype)
+
+
+def _saturate_cast(x, dtype, backend_module=None):
+    backend_module = backend_module or backend
+
+    def get_dtype_min_max(dtype):
+        if "bool" == dtype:
+            dtype_min = 0
+            dtype_max = 1
+        elif "int" in dtype:
+            dtype_min = ml_dtypes.iinfo(dtype).min
+            dtype_max = ml_dtypes.iinfo(dtype).max
+        else:
+            dtype_min = ml_dtypes.finfo(dtype).min
+            dtype_max = ml_dtypes.finfo(dtype).max
+        return dtype_min, dtype_max
+
+    dtype = backend.standardize_dtype(dtype)
+    in_dtype = backend.standardize_dtype(x.dtype)
+    in_min, in_max = get_dtype_min_max(in_dtype)
+    out_min, out_max = get_dtype_min_max(dtype)
+
+    # The output min/max may not actually be representable in the
+    # in_dtype (e.g. casting float32 to uint32).  This can lead to undefined
+    # behavior when trying to cast a value outside the valid range of the
+    # target type. We work around this by nudging the min/max to fall within
+    # the valid output range. The catch is that we may actually saturate
+    # to a value less than the true saturation limit, but this is the best we
+    # can do in order to avoid UB without backend op.
+    min_limit = np.maximum(in_min, out_min).astype(in_dtype)
+    if min_limit < out_min:
+        min_limit = np.nextafter(min_limit, 0, dtype=in_dtype)
+    max_limit = np.minimum(in_max, out_max).astype(in_dtype)
+    if max_limit > out_max:
+        max_limit = np.nextafter(max_limit, 0, dtype=in_dtype)
+
+    # Unconditionally apply `clip` to fix `inf` behavior.
+    x = backend_module.numpy.clip(x, min_limit, max_limit)
+
+    return backend_module.cast(x, dtype)
+
+
+class ConvertToTensor(Operation):
+    def __init__(self, dtype, sparse):
+        super().__init__()
+        self.dtype = backend.standardize_dtype(dtype)
+        self.sparse = sparse
+
+    def call(self, x):
+        return backend.core.convert_to_tensor(
+            x, dtype=self.dtype, sparse=self.sparse
+        )
+
+    def compute_output_spec(self, x):
+        dtype = x.dtype if self.dtype is None else self.dtype
+        sparse = (
+            False if self.sparse is not None and not self.sparse else x.sparse
+        )
+        return backend.KerasTensor(shape=x.shape, dtype=dtype, sparse=sparse)
+
+
+@keras_export("keras.ops.convert_to_tensor")
+def convert_to_tensor(x, dtype=None, sparse=None):
+    """Convert a NumPy array to a tensor.
+
+    Args:
+        x: A NumPy array, Python array (can be nested) or a backend tensor.
+        dtype: The target type. If `None`, the type of `x` is used.
+        sparse: Whether to keep sparse tensors. `False` will cause sparse
+            tensors to be densified. The default value of `None` means that
+            sparse tensors are kept only if the backend supports them.
+
+    Returns:
+        A backend tensor of the specified `dtype` and sparseness.
+
+    Example:
+
+    >>> x = np.array([1, 2, 3])
+    >>> y = keras.ops.convert_to_tensor(x)
+    """
+    if any_symbolic_tensors((x,)):
+        return ConvertToTensor(dtype=dtype, sparse=sparse)(x)
+    return backend.core.convert_to_tensor(x, dtype=dtype, sparse=sparse)
+
+
+@keras_export("keras.ops.convert_to_numpy")
+def convert_to_numpy(x):
+    """Convert a tensor to a NumPy array.
+
+    Args:
+        x: A tensor.
+
+    Returns:
+        A NumPy array.
+    """
+    if any_symbolic_tensors((x,)):
+        # This will raise a `ValueError` defined in the `KerasTensor` class.
+        # We trigger it rather than duplicate it here.
+        return np.array(x)
+    return backend.convert_to_numpy(x)
+
+
+class Cond(Operation):
+    @traceback_utils.filter_traceback
+    def __call__(self, *args, **kwargs):
+        def call_fn(*args, **kwargs):
+            if any_symbolic_tensors(args, kwargs):
+                return self.symbolic_call(*args, **kwargs)
+            else:
+                return self.call(*args, **kwargs)
+
+        if traceback_utils.is_traceback_filtering_enabled():
+            # Wrap self.call to provide helpful info in case of exception
+            call_fn = traceback_utils.inject_argument_info_in_traceback(
+                call_fn,
+                object_name=(f"{self.__class__.__name__}.call()"),
+            )
+            return call_fn(*args, **kwargs)
+
+        # Plain flow.
+        return call_fn(*args, **kwargs)
+
+    def call(self, pred, true_fn, false_fn):
+        return backend.core.cond(pred, true_fn, false_fn)
+
+    def compute_output_spec(self, pred, true_fn, false_fn):
+        true_fn_spec = backend.compute_output_spec(true_fn)
+        false_fn_spec = backend.compute_output_spec(false_fn)
+        if not self._check_output_spec(true_fn_spec, false_fn_spec):
+            raise ValueError(
+                "`true_fn` and `false_fn` should return outputs "
+                "of the same kind (struct, dtype and shape). "
+                f"Got {true_fn_spec} and {false_fn_spec} instead."
+            )
+        return true_fn_spec
+
+    def _check_output_spec(self, true_fn_spec, false_fn_spec):
+        try:
+            tree.assert_same_structure(true_fn_spec, false_fn_spec)
+        except:
+            return False
+
+        def check_leaf(t_spec, f_spec):
+            if t_spec is None or f_spec is None:
+                return t_spec is None and f_spec is None
+            return t_spec.shape == f_spec.shape and t_spec.dtype == f_spec.dtype
+
+        same = tree.map_structure(check_leaf, true_fn_spec, false_fn_spec)
+        return all(tree.flatten(same))
+
+
+@keras_export("keras.ops.cond")
+def cond(pred, true_fn, false_fn):
+    """Conditionally applies `true_fn` or `false_fn`.
+
+    Args:
+        pred: Boolean scalar type
+        true_fn: Callable returning the output for the `pred == True` case.
+        false_fn: Callable returning the output for the `pred == False` case.
+
+    Returns:
+        The output of either `true_fn` or `false_fn` depending on pred.
+    """
+    return Cond()(pred, true_fn, false_fn)
+
+
+# TODO: also create an Op subclass VectorizedMap.
+@keras_export("keras.ops.vectorized_map")
+def vectorized_map(function, elements):
+    """Parallel map of `function` on axis 0 of tensor(s) `elements`.
+
+    Schematically, `vectorized_map` implements the following,
+    in the case of a single tensor input `elements`:
+
+    ```python
+    def vectorized_map(function, elements)
+        outputs = []
+        for e in elements:
+            outputs.append(function(e))
+        return stack(outputs)
+    ```
+
+    In the case of an iterable of tensors `elements`,
+    it implements the following:
+
+    ```python
+    def vectorized_map(function, elements)
+        batch_size = elements[0].shape[0]
+        outputs = []
+        for index in range(batch_size):
+            outputs.append(function([e[index] for e in elements]))
+        return np.stack(outputs)
+    ```
+
+    In this case, `function` is expected to take as input
+    a single list of tensor arguments.
+    """
+    return backend.core.vectorized_map(function, elements)
+
+
+@keras_export("keras.ops.is_tensor")
+def is_tensor(x):
+    """Check whether the given object is a tensor.
+
+    Note: This checks for backend specific tensors so passing a TensorFlow
+    tensor would return `False` if your backend is PyTorch or JAX.
+
+    Args:
+        x: A variable.
+
+    Returns:
+        `True` if `x` is a tensor, otherwise `False`.
+    """
+    return backend.core.is_tensor(x)
+
+
+@keras_export("keras.ops.custom_gradient")
+def custom_gradient(f):
+    """Decorator to define a function with a custom gradient.
+
+    This decorator allows fine grained control over the gradients of a sequence
+    for operations. This may be useful for multiple reasons, including providing
+    a more efficient or numerically stable gradient for a sequence of
+    operations.
+
+    Args:
+        f: Function `f(*args)` that returns a tuple
+            `(output, grad_fn)`, where:
+            - `args` is a sequence of (nested structures of) tensor inputs to
+                the function.
+            - `output` is a (nested structure of) tensor outputs of applying
+                operations in `forward_fn` to `args`.
+            - `grad_fn` is a function with the signature `grad_fn(*args,
+                upstream)` which returns a tuple of tensors the same size as
+                (flattened) `args`: the derivatives of tensors in `output` with
+                respect to the tensors in `args`. `upstream` is a tensor or
+                sequence of tensors holding the initial value gradients for each
+                tensor in `output`.
+
+    Returns:
+        A function `h(*args)` which returns the same value as
+        `f(*args)[0]` and whose gradient is determined by
+        `f(*args)[1]`.
+
+
+    Examples:
+
+    1. Backend-agnostic example.
+
+    ```python
+    @ops.custom_gradient
+    def log1pexp(x):
+        e = ops.exp(x)
+
+        def grad(*args, upstream=None):
+            if upstream is None:
+                (upstream,) = args
+            return ops.multiply(upstream, 1.0 - 1.0 / ops.add(1, e))
+
+        return ops.log(1 + e), grad
+    ```
+
+    Note that the grad function that returns gradient computation
+    requires `args` as well as an `upstream` keyword argument, depending
+    on the backend being set. With the JAX and TensorFlow backends,
+    it requires only one argument, whereas it might use the `upstream`
+    argument in the case of the PyTorch backend.
+
+    When working with TensorFlow/JAX backend, `grad(upstream)`
+    is sufficient. With PyTorch, the `grad` function requires
+    `*args` as well as `upstream`, e.g. `def grad(*args, upstream)`.
+    Follow the previous example to use `@ops.custom_gradient` in
+    a way that is compatible with all backends.
+
+    2. Here's JAX & TensorFlow-specific example:
+
+    ```python
+    @ops.custom_gradient
+    def log1pexp(x):
+        e = ops.exp(x)
+        def grad(upstream):
+            return ops.multiply(upstream, 1.0 - 1.0 / ops.add(1, e))
+        return ops.log(1 + e), grad
+    ```
+
+    3. Lastly, here's a PyTorch-specific example,
+    using `*args` & `upstream`:
+
+    ```python
+    @ops.custom_gradient
+    def log1pexp(x):
+        e = ops.exp(x)
+        def grad(*args, upstream):
+            return ops.multiply(upstream, 1.0 - 1.0 / ops.add(1, e))
+        return ops.log(1 + e), grad
+    ```
+    """
+    return backend.core.custom_gradient(f)

SwarmUI/dlbackend/ComfyUI/venv/lib/python3.10/site-packages/keras/src/ops/function.py

@@ -0,0 +1,423 @@
+import collections
+
+from keras.src import tree
+from keras.src.api_export import keras_export
+from keras.src.backend import KerasTensor
+from keras.src.backend.config import backend
+from keras.src.ops.operation import Operation
+
+
+@keras_export("keras.Function")
+class Function(Operation):
+    """Class that encapsulates a computation graph of Keras operations.
+
+    You can use a `Function` to capture the computation graph linking
+    some input tensors to some output tensors, and reapply the same
+    computation on new inputs.
+
+    A `Function` is similar to a Functional Model, with the difference
+    that it is stateless (it does not track state variables)
+    and does not implement the `Layer` API.
+
+    Example:
+
+    ```python
+    input_1 = keras.KerasTensor(shape=(None, 2, 3))
+    input_2 = keras.KerasTensor(shape=(None, 2, 3))
+    x = input_1 + input_2
+    output = keras.ops.sigmoid(x)
+    fn = keras.Function(inputs=[input_1, input_2], outputs=output)
+
+    input_1_val = np.random.random((4, 2, 3))
+    input_2_val = np.random.random((4, 2, 3))
+    output_val = fn([input_1_val, input_2_val])
+    ```
+
+    Args:
+        inputs: `KerasTensor` instance or nested structured of
+            `KerasTensor` instances.
+        outputs: `KerasTensor` instance or nested structured of
+            `KerasTensor` instances. They should be computable
+            given only the values of `inputs`.
+        name: String. The name of the function.
+    """
+
+    def __init__(self, inputs, outputs, name=None):
+        super().__init__(name=name)
+
+        if backend() == "tensorflow":
+            # Temporary work around for
+            # https://github.com/keras-team/keras/issues/931
+            # This stop tensorflow from wrapping tf.function output in a
+            # _DictWrapper object.
+            _self_setattr_tracking = getattr(
+                self, "_self_setattr_tracking", True
+            )
+            self._self_setattr_tracking = False
+        self._inputs_struct = tree.map_structure(lambda x: x, inputs)
+        self._outputs_struct = tree.map_structure(lambda x: x, outputs)
+        self._inputs = tree.flatten(inputs)
+        self._outputs = tree.flatten(outputs)
+        if not self._inputs:
+            raise ValueError(
+                "`inputs` argument cannot be empty. Received:\n"
+                f"inputs={inputs}\n"
+                f"outputs={outputs}"
+            )
+        if not self._outputs:
+            raise ValueError(
+                "`outputs` argument cannot be empty. Received:\n"
+                f"inputs={inputs}\n"
+                f"outputs={outputs}"
+            )
+
+        if backend() == "tensorflow":
+            self._self_setattr_tracking = _self_setattr_tracking
+
+        (nodes, nodes_by_depth, operations, operations_by_depth) = map_graph(
+            self._inputs, self._outputs
+        )
+        self._nodes = nodes
+        self._nodes_by_depth = nodes_by_depth
+        self._operations = operations
+        self._operations_by_depth = operations_by_depth
+
+    @property
+    def operations(self):
+        return self._operations[:]
+
+    @property
+    def inputs(self):
+        """Flat list of the symbolic inputs of the Function."""
+        return self._inputs
+
+    @property
+    def outputs(self):
+        """Flat list of the symbolic outputs of the Function."""
+        return self._outputs
+
+    def compute_output_spec(self, inputs):
+        self._assert_input_compatibility(inputs)
+        # Check if input shapes are identical to ref input shapes,
+        # if so take a shortcut.
+        shortcut = True
+        for x, x_ref in zip(tree.flatten(inputs), self._inputs):
+            if x.shape != x_ref.shape:
+                shortcut = False
+                break
+        if shortcut:
+            return tree.map_structure(
+                lambda x: KerasTensor(shape=x.shape, dtype=x.dtype),
+                self._outputs_struct,
+            )
+        # No luck; take the long road through the graph.
+        # Original Keras used a cache to avoid recomputing all this
+        # when known input shapes where seen again. Perhaps a good
+        # idea to bring that back.
+        return self._run_through_graph(
+            inputs, operation_fn=lambda op: op.compute_output_spec
+        )
+
+    def compute_output_shape(self, input_shape):
+        # Wrap `input_shape` into the structure of KerasTensor to utilize
+        # `compute_output_spec`.
+        input_shape_struct = tree.map_shape_structure(
+            lambda x: KerasTensor(shape=x), input_shape
+        )
+        # Ensure that dtype and sparse settings are the same as self._inputs,
+        # because we only care about the shape in this function.
+        for x, x_ref in zip(tree.flatten(input_shape_struct), self._inputs):
+            x._dtype = x_ref.dtype
+            x._sparse = x_ref.sparse
+        output_spec = self.compute_output_spec(input_shape_struct)
+        return tree.map_structure(lambda x: x.shape, output_spec)
+
+    def call(self, inputs):
+        """Computes output tensors for new inputs."""
+        self._assert_input_compatibility(inputs)
+        return self._run_through_graph(inputs, operation_fn=lambda op: op)
+
+    def _run_through_graph(self, inputs, operation_fn, call_fn=None):
+        """Execute the graph.
+
+        At each node we compute outputs via
+        `operation_fn(node.operation)(*args, **kwargs)`.
+        """
+        inputs = tree.flatten(inputs)
+
+        # Dictionary mapping reference tensors to computed tensors.
+        tensor_dict = {}
+        for x, y in zip(self.inputs, inputs):
+            tensor_dict[id(x)] = y
+
+        nodes_by_depth = self._nodes_by_depth
+        depth_keys = list(nodes_by_depth.keys())
+        depth_keys.sort(reverse=True)
+
+        for depth in depth_keys:
+            nodes = nodes_by_depth[depth]
+            for node in nodes:
+                if not node.operation or node.is_input:
+                    continue  # Input tensors already exist.
+
+                if any(id(x) not in tensor_dict for x in node.input_tensors):
+                    continue  # Node is not computable, try skipping.
+
+                args, kwargs = node.arguments.fill_in(tensor_dict)
+                op = operation_fn(node.operation)
+                if call_fn is not None:
+                    outputs = call_fn(op, *args, **kwargs)
+                else:
+                    outputs = op(*args, **kwargs)
+
+                # Update tensor_dict.
+                for x, y in zip(node.outputs, tree.flatten(outputs)):
+                    tensor_dict[id(x)] = y
+
+        output_tensors = []
+        for x in self.outputs:
+            output_tensors.append(tensor_dict[id(x)])
+
+        return tree.pack_sequence_as(self._outputs_struct, output_tensors)
+
+    def _assert_input_compatibility(self, inputs):
+        try:
+            tree.assert_same_structure(inputs, self._inputs_struct)
+        except ValueError:
+            raise ValueError(
+                "Function was called with an invalid input structure. "
+                f"Expected input structure: {self._inputs_struct}\n"
+                f"Received input structure: {inputs}"
+            )
+        for x, x_ref in zip(tree.flatten(inputs), self._inputs):
+            if len(x.shape) != len(x_ref.shape):
+                raise ValueError(
+                    f"{self.__class__.__name__} was passed "
+                    f"incompatible inputs. For input '{x_ref.name}', "
+                    f"expected shape {x_ref.shape}, but received "
+                    f"instead a tensor with shape {x.shape}."
+                )
+            for dim, ref_dim in zip(x.shape, x_ref.shape):
+                if ref_dim is not None and dim is not None:
+                    if dim != ref_dim:
+                        raise ValueError(
+                            f"{self.__class__.__name__} was passed "
+                            f"incompatible inputs. For input '{x_ref.name}', "
+                            f"expected shape {x_ref.shape}, but received "
+                            f"instead a tensor with shape {x.shape}."
+                        )
+
+
+def make_node_key(op, node_index):
+    return str(id(op)) + "_ib-" + str(node_index)
+
+
+def map_graph(inputs, outputs):
+    """Validates a graph's topology and gather its operations and nodes.
+
+    Args:
+        inputs: List of input tensors.
+        outputs: List of outputs tensors.
+
+    Returns:
+        A tuple `(nodes, nodes_by_depth, operations, operations_by_depth)`.
+        - nodes: set of Node instances
+        - nodes_by_depth: dict mapping ints (depth) to lists of node instances.
+        - operations: list of Operation instances.
+        - operations_by_depth: dict mapping ints (depth) to lists of Operation
+            instances.
+    """
+    # "depth" is number of operations between output Node and the Node.
+    # Nodes are ordered from inputs -> outputs.
+    nodes_in_decreasing_depth, operation_indices = _build_map(inputs, outputs)
+    network_nodes = {
+        make_node_key(node.operation, node.operation._inbound_nodes.index(node))
+        for node in nodes_in_decreasing_depth
+    }
+
+    nodes_depths = {}  # dict {node: depth value}
+    operations_depths = {}  # dict {operation: depth value}
+
+    for node in reversed(nodes_in_decreasing_depth):
+        # If the depth is not set, the node has no outbound nodes (depth 0).
+        depth = nodes_depths.setdefault(node, 0)
+
+        # Update the depth of the corresponding operation
+        previous_depth = operations_depths.get(node.operation, 0)
+        # If we've seen this operation before at a higher depth,
+        # we should use that depth instead of the node depth.
+        # This is necessary for shared operations that have inputs at different
+        # depth levels in the graph.
+        depth = max(depth, previous_depth)
+        operations_depths[node.operation] = depth
+        nodes_depths[node] = depth
+
+        # Update the depth of inbound nodes.
+        # The "depth" of a node is the max of the depths
+        # of all nodes it is connected to + 1.
+        for node_dep in node.parent_nodes:
+            previous_depth = nodes_depths.get(node_dep, 0)
+            nodes_depths[node_dep] = max(depth + 1, previous_depth)
+
+    # Handle inputs that are not connected to outputs.
+    # We do not error out here because the inputs may be used to compute losses
+    # and metrics.
+    for input_t in inputs:
+        input_operation = input_t._keras_history[0]
+        if input_operation and input_operation not in operations_depths:
+            operations_depths[input_operation] = 0
+            operation_indices[input_operation] = -1
+            nodes_depths[input_operation._inbound_nodes[0]] = 0
+            network_nodes.add(make_node_key(input_operation, 0))
+
+    # Build a dict {depth: list of nodes with this depth}
+    nodes_by_depth = collections.defaultdict(list)
+    for node, depth in nodes_depths.items():
+        nodes_by_depth[depth].append(node)
+
+    # Build a dict {depth: list of operations with this depth}
+    operations_by_depth = collections.defaultdict(list)
+    for operation, depth in operations_depths.items():
+        operations_by_depth[depth].append(operation)
+
+    # Get sorted list of operation depths.
+    depth_keys = list(operations_by_depth.keys())
+    depth_keys.sort(reverse=True)
+
+    # Set self.operations ordered by depth.
+    operations = []
+    for depth in depth_keys:
+        operations_for_depth = operations_by_depth[depth]
+        # Network.operations needs to have a deterministic order:
+        # here we order them by traversal order.
+        operations_for_depth.sort(key=lambda x: operation_indices[x])
+        operations.extend(operations_for_depth)
+
+    # Get sorted list of node depths.
+    depth_keys = list(nodes_by_depth.keys())
+    depth_keys.sort(reverse=True)
+
+    # Check that all tensors required are computable.
+    # computable_tensors: all tensors in the graph
+    # that can be computed from the inputs provided.
+    computable_tensors = set()
+    for x in inputs:
+        computable_tensors.add(x)
+
+    operations_with_complete_input = []  # To provide a better error msg.
+    for depth in depth_keys:
+        for node in nodes_by_depth[depth]:
+            for x in tree.flatten(node.input_tensors):
+                if x not in computable_tensors:
+                    operation = node.operation
+                    raise ValueError(
+                        "Graph disconnected: cannot find parent for "
+                        f"tensor {x} at operation '{operation}'. "
+                        "The following previous operations were accessed "
+                        f"without issue: {operations_with_complete_input}"
+                    )
+                operations_with_complete_input.append(node.operation.name)
+
+            for x in tree.flatten(node.outputs):
+                computable_tensors.add(x)
+
+    # Ensure name unicity, which will be crucial for serialization
+    # (since serialized nodes refer to operations by their name).
+    all_names = [operation.name for operation in operations]
+    for name in all_names:
+        if all_names.count(name) != 1:
+            raise ValueError(
+                f'The name "{name}" is used {all_names.count(name)} '
+                "times in the model. All operation names should be unique."
+            )
+    return network_nodes, nodes_by_depth, operations, operations_by_depth
+
+
+def _build_map(inputs, outputs):
+    """Topologically sort nodes in order from inputs to outputs.
+
+    It uses a depth-first search to topologically sort nodes that appear in the
+    _keras_history connectivity metadata of `outputs`.
+
+    Args:
+        outputs: the output tensors whose _keras_history metadata should be
+                walked. This may be an arbitrary nested structure.
+
+    Returns:
+        A tuple like (ordered_nodes, operation_to_first_traversal_index)
+        ordered_nodes: list of nodes appearing in the keras history,
+            topologically sorted from original inputs to the `outputs`.
+            (If outputs have different sets of ancestors, the inputs to one
+            output may appear after a different output).
+        operation_to_first_traversal_index:
+            A dict mapping operation to the traversal index in the DFS where it
+            is seen. Note: if a operation is shared by several nodes, the dict
+            will onlystore the index corresponding to the *first* time the
+            operation seen.
+    """
+    finished_nodes = set()
+    nodes_in_progress = set()
+    nodes_in_decreasing_depth = []  # nodes from inputs -> outputs.
+    operation_indices = {}  # operation -> in traversal order.
+    for output in tree.flatten(outputs):
+        _build_map_helper(
+            inputs,
+            output,
+            finished_nodes,
+            nodes_in_progress,
+            nodes_in_decreasing_depth,
+            operation_indices,
+        )
+    return nodes_in_decreasing_depth, operation_indices
+
+
+def _build_map_helper(
+    inputs,
+    tensor,
+    finished_nodes,
+    nodes_in_progress,
+    nodes_in_decreasing_depth,
+    operation_indices,
+):
+    """Recursive helper for `_build_map`."""
+    (
+        operation,
+        node_index,
+        _,
+    ) = tensor._keras_history
+    if not operation:
+        return
+
+    node = operation._inbound_nodes[node_index]
+
+    # Don't repeat work for shared subgraphs
+    if node in finished_nodes:
+        return
+
+    # Prevent cycles.
+    if node in nodes_in_progress:
+        raise ValueError(
+            f"Tensor {tensor} from operation '{operation.name}' is part of a "
+            "cycle."
+        )
+
+    # Store the traversal order for operation sorting.
+    if operation not in operation_indices:
+        operation_indices[operation] = len(operation_indices)
+
+    # Propagate to all previous tensors connected to this node.
+    nodes_in_progress.add(node)
+    if not node.is_input and tensor not in tree.flatten(inputs):
+        for tensor in node.input_tensors:
+            _build_map_helper(
+                inputs,
+                tensor,
+                finished_nodes,
+                nodes_in_progress,
+                nodes_in_decreasing_depth,
+                operation_indices,
+            )
+
+    finished_nodes.add(node)
+    nodes_in_progress.remove(node)
+    nodes_in_decreasing_depth.append(node)

SwarmUI/dlbackend/ComfyUI/venv/lib/python3.10/site-packages/keras/src/ops/image.py

@@ -0,0 +1,1235 @@
+from keras.src import backend
+from keras.src import ops
+from keras.src.api_export import keras_export
+from keras.src.backend import KerasTensor
+from keras.src.backend import any_symbolic_tensors
+from keras.src.ops.operation import Operation
+from keras.src.ops.operation_utils import compute_conv_output_shape
+
+
+class RGBToGrayscale(Operation):
+    def __init__(self, data_format=None):
+        super().__init__()
+        self.data_format = backend.standardize_data_format(data_format)
+
+    def call(self, images):
+        return backend.image.rgb_to_grayscale(
+            images, data_format=self.data_format
+        )
+
+    def compute_output_spec(self, images):
+        images_shape = list(images.shape)
+        if len(images_shape) not in (3, 4):
+            raise ValueError(
+                "Invalid images rank: expected rank 3 (single image) "
+                "or rank 4 (batch of images). "
+                f"Received: images.shape={images_shape}"
+            )
+        if self.data_format == "channels_last":
+            images_shape[-1] = 1
+        else:
+            images_shape[-3] = 1
+        return KerasTensor(shape=images_shape, dtype=images.dtype)
+
+
+@keras_export("keras.ops.image.rgb_to_grayscale")
+def rgb_to_grayscale(images, data_format=None):
+    """Convert RGB images to grayscale.
+
+    This function converts RGB images to grayscale images. It supports both
+    3D and 4D tensors.
+
+    Args:
+        images: Input image or batch of images. Must be 3D or 4D.
+        data_format: A string specifying the data format of the input tensor.
+            It can be either `"channels_last"` or `"channels_first"`.
+            `"channels_last"` corresponds to inputs with shape
+            `(batch, height, width, channels)`, while `"channels_first"`
+            corresponds to inputs with shape `(batch, channels, height, width)`.
+            If not specified, the value will default to
+            `keras.config.image_data_format`.
+
+    Returns:
+        Grayscale image or batch of grayscale images.
+
+    Examples:
+
+    >>> import numpy as np
+    >>> from keras import ops
+    >>> x = np.random.random((2, 4, 4, 3))
+    >>> y = ops.image.rgb_to_grayscale(x)
+    >>> y.shape
+    (2, 4, 4, 1)
+
+    >>> x = np.random.random((4, 4, 3)) # Single RGB image
+    >>> y = ops.image.rgb_to_grayscale(x)
+    >>> y.shape
+    (4, 4, 1)
+
+    >>> x = np.random.random((2, 3, 4, 4))
+    >>> y = ops.image.rgb_to_grayscale(x, data_format="channels_first")
+    >>> y.shape
+    (2, 1, 4, 4)
+    """
+    if any_symbolic_tensors((images,)):
+        return RGBToGrayscale(data_format=data_format).symbolic_call(images)
+    return backend.image.rgb_to_grayscale(images, data_format=data_format)
+
+
+class RGBToHSV(Operation):
+    def __init__(self, data_format=None):
+        super().__init__()
+        self.data_format = backend.standardize_data_format(data_format)
+
+    def call(self, images):
+        return backend.image.rgb_to_hsv(images, data_format=self.data_format)
+
+    def compute_output_spec(self, images):
+        images_shape = list(images.shape)
+        dtype = images.dtype
+        if len(images_shape) not in (3, 4):
+            raise ValueError(
+                "Invalid images rank: expected rank 3 (single image) "
+                "or rank 4 (batch of images). "
+                f"Received: images.shape={images_shape}"
+            )
+        if not backend.is_float_dtype(dtype):
+            raise ValueError(
+                "Invalid images dtype: expected float dtype. "
+                f"Received: images.dtype={dtype}"
+            )
+        return KerasTensor(shape=images_shape, dtype=images.dtype)
+
+
+@keras_export("keras.ops.image.rgb_to_hsv")
+def rgb_to_hsv(images, data_format=None):
+    """Convert RGB images to HSV.
+
+    `images` must be of float dtype, and the output is only well defined if the
+    values in `images` are in `[0, 1]`.
+
+    All HSV values are in `[0, 1]`. A hue of `0` corresponds to pure red, `1/3`
+    is pure green, and `2/3` is pure blue.
+
+    Args:
+        images: Input image or batch of images. Must be 3D or 4D.
+        data_format: A string specifying the data format of the input tensor.
+            It can be either `"channels_last"` or `"channels_first"`.
+            `"channels_last"` corresponds to inputs with shape
+            `(batch, height, width, channels)`, while `"channels_first"`
+            corresponds to inputs with shape `(batch, channels, height, width)`.
+            If not specified, the value will default to
+            `keras.config.image_data_format`.
+
+    Returns:
+        HSV image or batch of HSV images.
+
+    Examples:
+
+    >>> import numpy as np
+    >>> from keras import ops
+    >>> x = np.random.random((2, 4, 4, 3))
+    >>> y = ops.image.rgb_to_hsv(x)
+    >>> y.shape
+    (2, 4, 4, 3)
+
+    >>> x = np.random.random((4, 4, 3)) # Single RGB image
+    >>> y = ops.image.rgb_to_hsv(x)
+    >>> y.shape
+    (4, 4, 3)
+
+    >>> x = np.random.random((2, 3, 4, 4))
+    >>> y = ops.image.rgb_to_hsv(x, data_format="channels_first")
+    >>> y.shape
+    (2, 3, 4, 4)
+    """
+    if any_symbolic_tensors((images,)):
+        return RGBToHSV(data_format=data_format).symbolic_call(images)
+    return backend.image.rgb_to_hsv(images, data_format=data_format)
+
+
+class HSVToRGB(Operation):
+    def __init__(self, data_format=None):
+        super().__init__()
+        self.data_format = backend.standardize_data_format(data_format)
+
+    def call(self, images):
+        return backend.image.hsv_to_rgb(images, data_format=self.data_format)
+
+    def compute_output_spec(self, images):
+        images_shape = list(images.shape)
+        dtype = images.dtype
+        if len(images_shape) not in (3, 4):
+            raise ValueError(
+                "Invalid images rank: expected rank 3 (single image) "
+                "or rank 4 (batch of images). "
+                f"Received: images.shape={images_shape}"
+            )
+        if not backend.is_float_dtype(dtype):
+            raise ValueError(
+                "Invalid images dtype: expected float dtype. "
+                f"Received: images.dtype={dtype}"
+            )
+        return KerasTensor(shape=images_shape, dtype=images.dtype)
+
+
+@keras_export("keras.ops.image.hsv_to_rgb")
+def hsv_to_rgb(images, data_format=None):
+    """Convert HSV images to RGB.
+
+    `images` must be of float dtype, and the output is only well defined if the
+    values in `images` are in `[0, 1]`.
+
+    Args:
+        images: Input image or batch of images. Must be 3D or 4D.
+        data_format: A string specifying the data format of the input tensor.
+            It can be either `"channels_last"` or `"channels_first"`.
+            `"channels_last"` corresponds to inputs with shape
+            `(batch, height, width, channels)`, while `"channels_first"`
+            corresponds to inputs with shape `(batch, channels, height, width)`.
+            If not specified, the value will default to
+            `keras.config.image_data_format`.
+
+    Returns:
+        RGB image or batch of RGB images.
+
+    Examples:
+
+    >>> import numpy as np
+    >>> from keras import ops
+    >>> x = np.random.random((2, 4, 4, 3))
+    >>> y = ops.image.hsv_to_rgb(x)
+    >>> y.shape
+    (2, 4, 4, 3)
+
+    >>> x = np.random.random((4, 4, 3)) # Single HSV image
+    >>> y = ops.image.hsv_to_rgb(x)
+    >>> y.shape
+    (4, 4, 3)
+
+    >>> x = np.random.random((2, 3, 4, 4))
+    >>> y = ops.image.hsv_to_rgb(x, data_format="channels_first")
+    >>> y.shape
+    (2, 3, 4, 4)
+    """
+    if any_symbolic_tensors((images,)):
+        return HSVToRGB(data_format=data_format).symbolic_call(images)
+    return backend.image.hsv_to_rgb(images, data_format=data_format)
+
+
+class Resize(Operation):
+    def __init__(
+        self,
+        size,
+        interpolation="bilinear",
+        antialias=False,
+        crop_to_aspect_ratio=False,
+        pad_to_aspect_ratio=False,
+        fill_mode="constant",
+        fill_value=0.0,
+        data_format=None,
+    ):
+        super().__init__()
+        self.size = tuple(size)
+        self.interpolation = interpolation
+        self.antialias = antialias
+        self.crop_to_aspect_ratio = crop_to_aspect_ratio
+        self.pad_to_aspect_ratio = pad_to_aspect_ratio
+        self.fill_mode = fill_mode
+        self.fill_value = fill_value
+        self.data_format = backend.standardize_data_format(data_format)
+
+    def call(self, images):
+        return _resize(
+            images,
+            self.size,
+            interpolation=self.interpolation,
+            antialias=self.antialias,
+            data_format=self.data_format,
+            crop_to_aspect_ratio=self.crop_to_aspect_ratio,
+            pad_to_aspect_ratio=self.pad_to_aspect_ratio,
+            fill_mode=self.fill_mode,
+            fill_value=self.fill_value,
+        )
+
+    def compute_output_spec(self, images):
+        images_shape = list(images.shape)
+        if len(images_shape) not in (3, 4):
+            raise ValueError(
+                "Invalid images rank: expected rank 3 (single image) "
+                "or rank 4 (batch of images). Received input with shape: "
+                f"images.shape={images.shape}"
+            )
+        if self.data_format == "channels_last":
+            height_axis, width_axis = -3, -2
+        else:
+            height_axis, width_axis = -2, -1
+        images_shape[height_axis] = self.size[0]
+        images_shape[width_axis] = self.size[1]
+        return KerasTensor(shape=images_shape, dtype=images.dtype)
+
+
+@keras_export("keras.ops.image.resize")
+def resize(
+    images,
+    size,
+    interpolation="bilinear",
+    antialias=False,
+    crop_to_aspect_ratio=False,
+    pad_to_aspect_ratio=False,
+    fill_mode="constant",
+    fill_value=0.0,
+    data_format=None,
+):
+    """Resize images to size using the specified interpolation method.
+
+    Args:
+        images: Input image or batch of images. Must be 3D or 4D.
+        size: Size of output image in `(height, width)` format.
+        interpolation: Interpolation method. Available methods are `"nearest"`,
+            `"bilinear"`, and `"bicubic"`. Defaults to `"bilinear"`.
+        antialias: Whether to use an antialiasing filter when downsampling an
+            image. Defaults to `False`.
+        crop_to_aspect_ratio: If `True`, resize the images without aspect
+            ratio distortion. When the original aspect ratio differs
+            from the target aspect ratio, the output image will be
+            cropped so as to return the
+            largest possible window in the image (of size `(height, width)`)
+            that matches the target aspect ratio. By default
+            (`crop_to_aspect_ratio=False`), aspect ratio may not be preserved.
+        pad_to_aspect_ratio: If `True`, pad the images without aspect
+            ratio distortion. When the original aspect ratio differs
+            from the target aspect ratio, the output image will be
+            evenly padded on the short side.
+        fill_mode: When using `pad_to_aspect_ratio=True`, padded areas
+            are filled according to the given mode. Only `"constant"` is
+            supported at this time
+            (fill with constant value, equal to `fill_value`).
+        fill_value: Float. Padding value to use when `pad_to_aspect_ratio=True`.
+        data_format: A string specifying the data format of the input tensor.
+            It can be either `"channels_last"` or `"channels_first"`.
+            `"channels_last"` corresponds to inputs with shape
+            `(batch, height, width, channels)`, while `"channels_first"`
+            corresponds to inputs with shape `(batch, channels, height, width)`.
+            If not specified, the value will default to
+            `keras.config.image_data_format`.
+
+    Returns:
+        Resized image or batch of images.
+
+    Examples:
+
+    >>> x = np.random.random((2, 4, 4, 3)) # batch of 2 RGB images
+    >>> y = keras.ops.image.resize(x, (2, 2))
+    >>> y.shape
+    (2, 2, 2, 3)
+
+    >>> x = np.random.random((4, 4, 3)) # single RGB image
+    >>> y = keras.ops.image.resize(x, (2, 2))
+    >>> y.shape
+    (2, 2, 3)
+
+    >>> x = np.random.random((2, 3, 4, 4)) # batch of 2 RGB images
+    >>> y = keras.ops.image.resize(x, (2, 2),
+    ...     data_format="channels_first")
+    >>> y.shape
+    (2, 3, 2, 2)
+    """
+    if len(size) != 2:
+        raise ValueError(
+            "Expected `size` to be a tuple of 2 integers. "
+            f"Received: size={size}"
+        )
+    if len(images.shape) < 3 or len(images.shape) > 4:
+        raise ValueError(
+            "Invalid images rank: expected rank 3 (single image) "
+            "or rank 4 (batch of images). Received input with shape: "
+            f"images.shape={images.shape}"
+        )
+    if pad_to_aspect_ratio and crop_to_aspect_ratio:
+        raise ValueError(
+            "Only one of `pad_to_aspect_ratio` & `crop_to_aspect_ratio` "
+            "can be `True`."
+        )
+    if any_symbolic_tensors((images,)):
+        return Resize(
+            size,
+            interpolation=interpolation,
+            antialias=antialias,
+            data_format=data_format,
+            crop_to_aspect_ratio=crop_to_aspect_ratio,
+            pad_to_aspect_ratio=pad_to_aspect_ratio,
+            fill_mode=fill_mode,
+            fill_value=fill_value,
+        ).symbolic_call(images)
+    return _resize(
+        images,
+        size,
+        interpolation=interpolation,
+        antialias=antialias,
+        crop_to_aspect_ratio=crop_to_aspect_ratio,
+        data_format=data_format,
+        pad_to_aspect_ratio=pad_to_aspect_ratio,
+        fill_mode=fill_mode,
+        fill_value=fill_value,
+    )
+
+
+def _resize(
+    images,
+    size,
+    interpolation="bilinear",
+    antialias=False,
+    crop_to_aspect_ratio=False,
+    pad_to_aspect_ratio=False,
+    fill_mode="constant",
+    fill_value=0.0,
+    data_format=None,
+):
+    resized = backend.image.resize(
+        images,
+        size,
+        interpolation=interpolation,
+        antialias=antialias,
+        crop_to_aspect_ratio=crop_to_aspect_ratio,
+        data_format=data_format,
+        pad_to_aspect_ratio=pad_to_aspect_ratio,
+        fill_mode=fill_mode,
+        fill_value=fill_value,
+    )
+    if resized.dtype == images.dtype:
+        # Only `torch` backend will cast result to original dtype with
+        # correct rounding and without dtype overflow
+        return resized
+    if backend.is_int_dtype(images.dtype):
+        resized = ops.round(resized)
+    return ops.saturate_cast(resized, images.dtype)
+
+
+class AffineTransform(Operation):
+    def __init__(
+        self,
+        interpolation="bilinear",
+        fill_mode="constant",
+        fill_value=0,
+        data_format=None,
+    ):
+        super().__init__()
+        self.interpolation = interpolation
+        self.fill_mode = fill_mode
+        self.fill_value = fill_value
+        self.data_format = backend.standardize_data_format(data_format)
+
+    def call(self, images, transform):
+        return backend.image.affine_transform(
+            images,
+            transform,
+            interpolation=self.interpolation,
+            fill_mode=self.fill_mode,
+            fill_value=self.fill_value,
+            data_format=self.data_format,
+        )
+
+    def compute_output_spec(self, images, transform):
+        if len(images.shape) not in (3, 4):
+            raise ValueError(
+                "Invalid images rank: expected rank 3 (single image) "
+                "or rank 4 (batch of images). Received input with shape: "
+                f"images.shape={images.shape}"
+            )
+        if len(transform.shape) not in (1, 2):
+            raise ValueError(
+                "Invalid transform rank: expected rank 1 (single transform) "
+                "or rank 2 (batch of transforms). Received input with shape: "
+                f"transform.shape={transform.shape}"
+            )
+        return KerasTensor(images.shape, dtype=images.dtype)
+
+
+@keras_export("keras.ops.image.affine_transform")
+def affine_transform(
+    images,
+    transform,
+    interpolation="bilinear",
+    fill_mode="constant",
+    fill_value=0,
+    data_format=None,
+):
+    """Applies the given transform(s) to the image(s).
+
+    Args:
+        images: Input image or batch of images. Must be 3D or 4D.
+        transform: Projective transform matrix/matrices. A vector of length 8 or
+            tensor of size N x 8. If one row of transform is
+            `[a0, a1, a2, b0, b1, b2, c0, c1]`, then it maps the output point
+            `(x, y)` to a transformed input point
+            `(x', y') = ((a0 x + a1 y + a2) / k, (b0 x + b1 y + b2) / k)`,
+            where `k = c0 x + c1 y + 1`. The transform is inverted compared to
+            the transform mapping input points to output points. Note that
+            gradients are not backpropagated into transformation parameters.
+            Note that `c0` and `c1` are only effective when using TensorFlow
+            backend and will be considered as `0` when using other backends.
+        interpolation: Interpolation method. Available methods are `"nearest"`,
+            and `"bilinear"`. Defaults to `"bilinear"`.
+        fill_mode: Points outside the boundaries of the input are filled
+            according to the given mode. Available methods are `"constant"`,
+            `"nearest"`, `"wrap"` and `"reflect"`. Defaults to `"constant"`.
+            - `"reflect"`: `(d c b a | a b c d | d c b a)`
+                The input is extended by reflecting about the edge of the last
+                pixel.
+            - `"constant"`: `(k k k k | a b c d | k k k k)`
+                The input is extended by filling all values beyond
+                the edge with the same constant value k specified by
+                `fill_value`.
+            - `"wrap"`: `(a b c d | a b c d | a b c d)`
+                The input is extended by wrapping around to the opposite edge.
+            - `"nearest"`: `(a a a a | a b c d | d d d d)`
+                The input is extended by the nearest pixel.
+        fill_value: Value used for points outside the boundaries of the input if
+            `fill_mode="constant"`. Defaults to `0`.
+        data_format: A string specifying the data format of the input tensor.
+            It can be either `"channels_last"` or `"channels_first"`.
+            `"channels_last"` corresponds to inputs with shape
+            `(batch, height, width, channels)`, while `"channels_first"`
+            corresponds to inputs with shape `(batch, channels, height, width)`.
+            If not specified, the value will default to
+            `keras.config.image_data_format`.
+
+    Returns:
+        Applied affine transform image or batch of images.
+
+    Examples:
+
+    >>> x = np.random.random((2, 64, 80, 3)) # batch of 2 RGB images
+    >>> transform = np.array(
+    ...     [
+    ...         [1.5, 0, -20, 0, 1.5, -16, 0, 0],  # zoom
+    ...         [1, 0, -20, 0, 1, -16, 0, 0],  # translation
+    ...     ]
+    ... )
+    >>> y = keras.ops.image.affine_transform(x, transform)
+    >>> y.shape
+    (2, 64, 80, 3)
+
+    >>> x = np.random.random((64, 80, 3)) # single RGB image
+    >>> transform = np.array([1.0, 0.5, -20, 0.5, 1.0, -16, 0, 0])  # shear
+    >>> y = keras.ops.image.affine_transform(x, transform)
+    >>> y.shape
+    (64, 80, 3)
+
+    >>> x = np.random.random((2, 3, 64, 80)) # batch of 2 RGB images
+    >>> transform = np.array(
+    ...     [
+    ...         [1.5, 0, -20, 0, 1.5, -16, 0, 0],  # zoom
+    ...         [1, 0, -20, 0, 1, -16, 0, 0],  # translation
+    ...     ]
+    ... )
+    >>> y = keras.ops.image.affine_transform(x, transform,
+    ...     data_format="channels_first")
+    >>> y.shape
+    (2, 3, 64, 80)
+    """
+    if any_symbolic_tensors((images, transform)):
+        return AffineTransform(
+            interpolation=interpolation,
+            fill_mode=fill_mode,
+            fill_value=fill_value,
+            data_format=data_format,
+        ).symbolic_call(images, transform)
+    return backend.image.affine_transform(
+        images,
+        transform,
+        interpolation=interpolation,
+        fill_mode=fill_mode,
+        fill_value=fill_value,
+        data_format=data_format,
+    )
+
+
+class ExtractPatches(Operation):
+    def __init__(
+        self,
+        size,
+        strides=None,
+        dilation_rate=1,
+        padding="valid",
+        data_format=None,
+    ):
+        super().__init__()
+        if isinstance(size, int):
+            size = (size, size)
+        self.size = size
+        self.strides = strides
+        self.dilation_rate = dilation_rate
+        self.padding = padding
+        self.data_format = backend.standardize_data_format(data_format)
+
+    def call(self, images):
+        return _extract_patches(
+            images=images,
+            size=self.size,
+            strides=self.strides,
+            dilation_rate=self.dilation_rate,
+            padding=self.padding,
+            data_format=self.data_format,
+        )
+
+    def compute_output_spec(self, images):
+        images_shape = list(images.shape)
+        original_ndim = len(images_shape)
+        if not self.strides:
+            strides = (self.size[0], self.size[1])
+        if self.data_format == "channels_last":
+            channels_in = images_shape[-1]
+        else:
+            channels_in = images_shape[-3]
+        if original_ndim == 3:
+            images_shape = [1] + images_shape
+        filters = self.size[0] * self.size[1] * channels_in
+        kernel_size = (self.size[0], self.size[1])
+        out_shape = compute_conv_output_shape(
+            images_shape,
+            filters,
+            kernel_size,
+            strides=strides,
+            padding=self.padding,
+            data_format=self.data_format,
+            dilation_rate=self.dilation_rate,
+        )
+        if original_ndim == 3:
+            out_shape = out_shape[1:]
+        return KerasTensor(shape=out_shape, dtype=images.dtype)
+
+
+@keras_export("keras.ops.image.extract_patches")
+def extract_patches(
+    images,
+    size,
+    strides=None,
+    dilation_rate=1,
+    padding="valid",
+    data_format=None,
+):
+    """Extracts patches from the image(s).
+
+    Args:
+        images: Input image or batch of images. Must be 3D or 4D.
+        size: Patch size int or tuple (patch_height, patch_width)
+        strides: strides along height and width. If not specified, or
+            if `None`, it defaults to the same value as `size`.
+        dilation_rate: This is the input stride, specifying how far two
+            consecutive patch samples are in the input. For value other than 1,
+            strides must be 1. NOTE: `strides > 1` is not supported in
+            conjunction with `dilation_rate > 1`
+        padding: The type of padding algorithm to use: `"same"` or `"valid"`.
+        data_format: A string specifying the data format of the input tensor.
+            It can be either `"channels_last"` or `"channels_first"`.
+            `"channels_last"` corresponds to inputs with shape
+            `(batch, height, width, channels)`, while `"channels_first"`
+            corresponds to inputs with shape `(batch, channels, height, width)`.
+            If not specified, the value will default to
+            `keras.config.image_data_format`.
+
+    Returns:
+        Extracted patches 3D (if not batched) or 4D (if batched)
+
+    Examples:
+
+    >>> image = np.random.random(
+    ...     (2, 20, 20, 3)
+    ... ).astype("float32") # batch of 2 RGB images
+    >>> patches = keras.ops.image.extract_patches(image, (5, 5))
+    >>> patches.shape
+    (2, 4, 4, 75)
+    >>> image = np.random.random((20, 20, 3)).astype("float32") # 1 RGB image
+    >>> patches = keras.ops.image.extract_patches(image, (3, 3), (1, 1))
+    >>> patches.shape
+    (18, 18, 27)
+    """
+    if any_symbolic_tensors((images,)):
+        return ExtractPatches(
+            size=size,
+            strides=strides,
+            dilation_rate=dilation_rate,
+            padding=padding,
+            data_format=data_format,
+        ).symbolic_call(images)
+
+    return _extract_patches(
+        images, size, strides, dilation_rate, padding, data_format=data_format
+    )
+
+
+def _extract_patches(
+    images,
+    size,
+    strides=None,
+    dilation_rate=1,
+    padding="valid",
+    data_format=None,
+):
+    if isinstance(size, int):
+        patch_h = patch_w = size
+    elif len(size) == 2:
+        patch_h, patch_w = size[0], size[1]
+    else:
+        raise TypeError(
+            "Invalid `size` argument. Expected an "
+            f"int or a tuple of length 2. Received: size={size}"
+        )
+    data_format = backend.standardize_data_format(data_format)
+    if data_format == "channels_last":
+        channels_in = images.shape[-1]
+    elif data_format == "channels_first":
+        channels_in = images.shape[-3]
+    if not strides:
+        strides = size
+    out_dim = patch_h * patch_w * channels_in
+    kernel = backend.numpy.eye(out_dim, dtype=images.dtype)
+    kernel = backend.numpy.reshape(
+        kernel, (patch_h, patch_w, channels_in, out_dim)
+    )
+    _unbatched = False
+    if len(images.shape) == 3:
+        _unbatched = True
+        images = backend.numpy.expand_dims(images, axis=0)
+    patches = backend.nn.conv(
+        inputs=images,
+        kernel=kernel,
+        strides=strides,
+        padding=padding,
+        data_format=data_format,
+        dilation_rate=dilation_rate,
+    )
+    if _unbatched:
+        patches = backend.numpy.squeeze(patches, axis=0)
+    return patches
+
+
+class MapCoordinates(Operation):
+    def __init__(self, order, fill_mode="constant", fill_value=0):
+        super().__init__()
+        self.order = order
+        self.fill_mode = fill_mode
+        self.fill_value = fill_value
+
+    def call(self, inputs, coordinates):
+        return backend.image.map_coordinates(
+            inputs,
+            coordinates,
+            order=self.order,
+            fill_mode=self.fill_mode,
+            fill_value=self.fill_value,
+        )
+
+    def compute_output_spec(self, inputs, coordinates):
+        if coordinates.shape[0] != len(inputs.shape):
+            raise ValueError(
+                "First dim of `coordinates` must be the same as the rank of "
+                "`inputs`. "
+                f"Received inputs with shape: {inputs.shape} and coordinate "
+                f"leading dim of {coordinates.shape[0]}"
+            )
+        if len(coordinates.shape) < 2:
+            raise ValueError(
+                "Invalid coordinates rank: expected at least rank 2."
+                f" Received input with shape: {coordinates.shape}"
+            )
+        return KerasTensor(coordinates.shape[1:], dtype=inputs.dtype)
+
+
+@keras_export("keras.ops.image.map_coordinates")
+def map_coordinates(
+    inputs, coordinates, order, fill_mode="constant", fill_value=0
+):
+    """Map the input array to new coordinates by interpolation.
+
+    Note that interpolation near boundaries differs from the scipy function,
+    because we fixed an outstanding bug
+    [scipy/issues/2640](https://github.com/scipy/scipy/issues/2640).
+
+    Args:
+        inputs: The input array.
+        coordinates: The coordinates at which inputs is evaluated.
+        order: The order of the spline interpolation. The order must be `0` or
+            `1`. `0` indicates the nearest neighbor and `1` indicates the linear
+            interpolation.
+        fill_mode: Points outside the boundaries of the inputs are filled
+            according to the given mode. Available methods are `"constant"`,
+            `"nearest"`, `"wrap"` and `"mirror"` and `"reflect"`. Defaults to
+            `"constant"`.
+            - `"constant"`: `(k k k k | a b c d | k k k k)`
+                The inputs is extended by filling all values beyond
+                the edge with the same constant value k specified by
+                `fill_value`.
+            - `"nearest"`: `(a a a a | a b c d | d d d d)`
+                The inputs is extended by the nearest pixel.
+            - `"wrap"`: `(a b c d | a b c d | a b c d)`
+                The inputs is extended by wrapping around to the opposite edge.
+            - `"mirror"`: `(c d c b | a b c d | c b a b)`
+                The inputs is extended by mirroring about the edge.
+            - `"reflect"`: `(d c b a | a b c d | d c b a)`
+                The inputs is extended by reflecting about the edge of the last
+                pixel.
+        fill_value: Value used for points outside the boundaries of the inputs
+            if `fill_mode="constant"`. Defaults to `0`.
+
+    Returns:
+        Output input or batch of inputs.
+
+    """
+    if any_symbolic_tensors((inputs, coordinates)):
+        return MapCoordinates(
+            order,
+            fill_mode,
+            fill_value,
+        ).symbolic_call(inputs, coordinates)
+    return backend.image.map_coordinates(
+        inputs,
+        coordinates,
+        order,
+        fill_mode,
+        fill_value,
+    )
+
+
+class PadImages(Operation):
+    def __init__(
+        self,
+        top_padding=None,
+        left_padding=None,
+        bottom_padding=None,
+        right_padding=None,
+        target_height=None,
+        target_width=None,
+        data_format=None,
+    ):
+        super().__init__()
+        self.top_padding = top_padding
+        self.left_padding = left_padding
+        self.bottom_padding = bottom_padding
+        self.right_padding = right_padding
+        self.target_height = target_height
+        self.target_width = target_width
+        self.data_format = backend.standardize_data_format(data_format)
+
+    def call(self, images):
+        return _pad_images(
+            images,
+            self.top_padding,
+            self.left_padding,
+            self.bottom_padding,
+            self.right_padding,
+            self.target_height,
+            self.target_width,
+            self.data_format,
+        )
+
+    def compute_output_spec(self, images):
+        images_shape = list(images.shape)
+
+        if self.data_format == "channels_last":
+            height_axis, width_axis = -3, -2
+            height, width = images_shape[height_axis], images_shape[width_axis]
+        else:
+            height_axis, width_axis = -2, -1
+            height, width = images_shape[height_axis], images_shape[width_axis]
+
+        target_height = self.target_height
+        if target_height is None and height is not None:
+            target_height = self.top_padding + height + self.bottom_padding
+        target_width = self.target_width
+        if target_width is None and width is not None:
+            target_width = self.left_padding + width + self.right_padding
+
+        images_shape[height_axis] = target_height
+        images_shape[width_axis] = target_width
+        return KerasTensor(shape=images_shape, dtype=images.dtype)
+
+
+@keras_export("keras.ops.image.pad_images")
+def pad_images(
+    images,
+    top_padding=None,
+    left_padding=None,
+    bottom_padding=None,
+    right_padding=None,
+    target_height=None,
+    target_width=None,
+    data_format=None,
+):
+    """Pad `images` with zeros to the specified `height` and `width`.
+
+    Args:
+        images: Input image or batch of images. Must be 3D or 4D.
+        top_padding: Number of rows of zeros to add on top.
+        left_padding: Number of columns of zeros to add on the left.
+        bottom_padding: Number of rows of zeros to add at the bottom.
+        right_padding: Number of columns of zeros to add on the right.
+        target_height: Height of output images.
+        target_width: Width of output images.
+        data_format: A string specifying the data format of the input tensor.
+            It can be either `"channels_last"` or `"channels_first"`.
+            `"channels_last"` corresponds to inputs with shape
+            `(batch, height, width, channels)`, while `"channels_first"`
+            corresponds to inputs with shape `(batch, channels, height, width)`.
+            If not specified, the value will default to
+            `keras.config.image_data_format`.
+
+    Returns:
+        Padded image or batch of images.
+
+    Example:
+
+    >>> images = np.random.random((15, 25, 3))
+    >>> padded_images = keras.ops.image.pad_images(
+    ...     images, 2, 3, target_height=20, target_width=30
+    ... )
+    >>> padded_images.shape
+    (20, 30, 3)
+
+    >>> batch_images = np.random.random((2, 15, 25, 3))
+    >>> padded_batch = keras.ops.image.pad_images(
+    ...     batch_images, 2, 3, target_height=20, target_width=30
+    ... )
+    >>> padded_batch.shape
+    (2, 20, 30, 3)"""
+
+    if any_symbolic_tensors((images,)):
+        return PadImages(
+            top_padding,
+            left_padding,
+            bottom_padding,
+            right_padding,
+            target_height,
+            target_width,
+            data_format,
+        ).symbolic_call(images)
+
+    return _pad_images(
+        images,
+        top_padding,
+        left_padding,
+        bottom_padding,
+        right_padding,
+        target_height,
+        target_width,
+        data_format,
+    )
+
+
+def _pad_images(
+    images,
+    top_padding,
+    left_padding,
+    bottom_padding,
+    right_padding,
+    target_height,
+    target_width,
+    data_format=None,
+):
+    data_format = backend.standardize_data_format(data_format)
+    images = backend.convert_to_tensor(images)
+    images_shape = ops.shape(images)
+
+    # Check
+    if len(images_shape) not in (3, 4):
+        raise ValueError(
+            f"Invalid shape for argument `images`: "
+            "it must have rank 3 or 4. "
+            f"Received: images.shape={images_shape}"
+        )
+    if [top_padding, bottom_padding, target_height].count(None) != 1:
+        raise ValueError(
+            "Must specify exactly two of "
+            "top_padding, bottom_padding, target_height. "
+            f"Received: top_padding={top_padding}, "
+            f"bottom_padding={bottom_padding}, "
+            f"target_height={target_height}"
+        )
+    if [left_padding, right_padding, target_width].count(None) != 1:
+        raise ValueError(
+            "Must specify exactly two of "
+            "left_padding, right_padding, target_width. "
+            f"Received: left_padding={left_padding}, "
+            f"right_padding={right_padding}, "
+            f"target_width={target_width}"
+        )
+
+    is_batch = False if len(images_shape) == 3 else True
+    if data_format == "channels_last":
+        height, width = images_shape[-3], images_shape[-2]
+    else:
+        height, width = images_shape[-2], images_shape[-1]
+
+    # Infer padding
+    if top_padding is None:
+        top_padding = target_height - bottom_padding - height
+    if bottom_padding is None:
+        bottom_padding = target_height - top_padding - height
+    if left_padding is None:
+        left_padding = target_width - right_padding - width
+    if right_padding is None:
+        right_padding = target_width - left_padding - width
+
+    if top_padding < 0:
+        raise ValueError(
+            f"top_padding must be >= 0. Received: top_padding={top_padding}"
+        )
+    if left_padding < 0:
+        raise ValueError(
+            "left_padding must be >= 0. "
+            f"Received: left_padding={left_padding}"
+        )
+    if right_padding < 0:
+        raise ValueError(
+            "right_padding must be >= 0. "
+            f"Received: right_padding={right_padding}"
+        )
+    if bottom_padding < 0:
+        raise ValueError(
+            "bottom_padding must be >= 0. "
+            f"Received: bottom_padding={bottom_padding}"
+        )
+
+    # Compute pad_width
+    pad_width = [[top_padding, bottom_padding], [left_padding, right_padding]]
+    if data_format == "channels_last":
+        pad_width = pad_width + [[0, 0]]
+    else:
+        pad_width = [[0, 0]] + pad_width
+    if is_batch:
+        pad_width = [[0, 0]] + pad_width
+
+    padded_images = backend.numpy.pad(images, pad_width)
+    return padded_images
+
+
+class CropImages(Operation):
+    def __init__(
+        self,
+        top_cropping,
+        left_cropping,
+        bottom_cropping,
+        right_cropping,
+        target_height,
+        target_width,
+        data_format=None,
+    ):
+        super().__init__()
+        self.top_cropping = top_cropping
+        self.bottom_cropping = bottom_cropping
+        self.left_cropping = left_cropping
+        self.right_cropping = right_cropping
+        self.target_height = target_height
+        self.target_width = target_width
+        self.data_format = backend.standardize_data_format(data_format)
+
+    def call(self, images):
+        return _crop_images(
+            images,
+            self.top_cropping,
+            self.left_cropping,
+            self.bottom_cropping,
+            self.right_cropping,
+            self.target_height,
+            self.target_width,
+            self.data_format,
+        )
+
+    def compute_output_spec(self, images):
+        images_shape = list(images.shape)
+
+        if self.data_format == "channels_last":
+            height_axis, width_axis = -3, -2
+        else:
+            height_axis, width_axis = -2, -1
+        height, width = images_shape[height_axis], images_shape[width_axis]
+
+        if height is None and self.target_height is None:
+            raise ValueError(
+                "When the height of the images is unknown, `target_height` "
+                "must be specified."
+                f"Received images.shape={images_shape} and "
+                f"target_height={self.target_height}"
+            )
+        if width is None and self.target_width is None:
+            raise ValueError(
+                "When the width of the images is unknown, `target_width` "
+                "must be specified."
+                f"Received images.shape={images_shape} and "
+                f"target_width={self.target_width}"
+            )
+
+        target_height = self.target_height
+        if target_height is None:
+            target_height = height - self.top_cropping - self.bottom_cropping
+        target_width = self.target_width
+        if target_width is None:
+            target_width = width - self.left_cropping - self.right_cropping
+
+        images_shape[height_axis] = target_height
+        images_shape[width_axis] = target_width
+        return KerasTensor(shape=images_shape, dtype=images.dtype)
+
+
+@keras_export("keras.ops.image.crop_images")
+def crop_images(
+    images,
+    top_cropping=None,
+    left_cropping=None,
+    bottom_cropping=None,
+    right_cropping=None,
+    target_height=None,
+    target_width=None,
+    data_format=None,
+):
+    """Crop `images` to a specified `height` and `width`.
+
+    Args:
+        images: Input image or batch of images. Must be 3D or 4D.
+        top_cropping: Number of columns to crop from the top.
+        left_cropping: Number of columns to crop from the left.
+        bottom_cropping: Number of columns to crop from the bottom.
+        right_cropping: Number of columns to crop from the right.
+        target_height: Height of the output images.
+        target_width: Width of the output images.
+        data_format: A string specifying the data format of the input tensor.
+            It can be either `"channels_last"` or `"channels_first"`.
+            `"channels_last"` corresponds to inputs with shape
+            `(batch, height, width, channels)`, while `"channels_first"`
+            corresponds to inputs with shape `(batch, channels, height, width)`.
+            If not specified, the value will default to
+            `keras.config.image_data_format`.
+
+    Returns:
+        Cropped image or batch of images.
+
+    Example:
+
+    >>> images = np.reshape(np.arange(1, 28, dtype="float32"), [3, 3, 3])
+    >>> images[:,:,0] # print the first channel of the images
+    array([[ 1.,  4.,  7.],
+           [10., 13., 16.],
+           [19., 22., 25.]], dtype=float32)
+    >>> cropped_images = keras.image.crop_images(images, 0, 0, 2, 2)
+    >>> cropped_images[:,:,0] # print the first channel of the cropped images
+    array([[ 1.,  4.],
+           [10., 13.]], dtype=float32)"""
+
+    if any_symbolic_tensors((images,)):
+        return CropImages(
+            top_cropping,
+            left_cropping,
+            bottom_cropping,
+            right_cropping,
+            target_height,
+            target_width,
+            data_format,
+        ).symbolic_call(images)
+
+    return _crop_images(
+        images,
+        top_cropping,
+        left_cropping,
+        bottom_cropping,
+        right_cropping,
+        target_height,
+        target_width,
+        data_format,
+    )
+
+
+def _crop_images(
+    images,
+    top_cropping,
+    left_cropping,
+    bottom_cropping,
+    right_cropping,
+    target_height,
+    target_width,
+    data_format=None,
+):
+    data_format = backend.standardize_data_format(data_format)
+    images = backend.convert_to_tensor(images)
+    images_shape = ops.shape(images)
+
+    # Check
+    if len(images_shape) not in (3, 4):
+        raise ValueError(
+            f"Invalid shape for argument `images`: "
+            "it must have rank 3 or 4. "
+            f"Received: images.shape={images_shape}"
+        )
+    if [top_cropping, bottom_cropping, target_height].count(None) != 1:
+        raise ValueError(
+            "Must specify exactly two of "
+            "top_cropping, bottom_cropping, target_height. "
+            f"Received: top_cropping={top_cropping}, "
+            f"bottom_cropping={bottom_cropping}, "
+            f"target_height={target_height}"
+        )
+    if [left_cropping, right_cropping, target_width].count(None) != 1:
+        raise ValueError(
+            "Must specify exactly two of "
+            "left_cropping, right_cropping, target_width. "
+            f"Received: left_cropping={left_cropping}, "
+            f"right_cropping={right_cropping}, "
+            f"target_width={target_width}"
+        )
+
+    is_batch = False if len(images_shape) == 3 else True
+    if data_format == "channels_last":
+        height, width = images_shape[-3], images_shape[-2]
+        channels = images_shape[-1]
+    else:
+        height, width = images_shape[-2], images_shape[-1]
+        channels = images_shape[-3]
+
+    # Infer padding
+    if top_cropping is None:
+        top_cropping = height - target_height - bottom_cropping
+    if target_height is None:
+        target_height = height - bottom_cropping - top_cropping
+    if left_cropping is None:
+        left_cropping = width - target_width - right_cropping
+    if target_width is None:
+        target_width = width - right_cropping - left_cropping
+
+    if top_cropping < 0:
+        raise ValueError(
+            "top_cropping must be >= 0. "
+            f"Received: top_cropping={top_cropping}"
+        )
+    if target_height < 0:
+        raise ValueError(
+            "target_height must be >= 0. "
+            f"Received: target_height={target_height}"
+        )
+    if left_cropping < 0:
+        raise ValueError(
+            "left_cropping must be >= 0. "
+            f"Received: left_cropping={left_cropping}"
+        )
+    if target_width < 0:
+        raise ValueError(
+            "target_width must be >= 0. "
+            f"Received: target_width={target_width}"
+        )
+
+    # Compute start_indices and shape
+    start_indices = [top_cropping, left_cropping]
+    shape = [target_height, target_width]
+    if data_format == "channels_last":
+        start_indices = start_indices + [0]
+        shape = shape + [channels]
+    else:
+        start_indices = [0] + start_indices
+        shape = [channels] + shape
+    if is_batch:
+        batch_size = images_shape[0]
+        start_indices = [0] + start_indices
+        shape = [batch_size] + shape
+
+    cropped_images = ops.slice(images, start_indices, shape)
+    return cropped_images

SwarmUI/dlbackend/ComfyUI/venv/lib/python3.10/site-packages/keras/src/ops/linalg.py

@@ -0,0 +1,707 @@
+from keras.src import backend
+from keras.src.api_export import keras_export
+from keras.src.backend import KerasTensor
+from keras.src.backend import any_symbolic_tensors
+from keras.src.ops.operation import Operation
+from keras.src.ops.operation_utils import reduce_shape
+
+
+class Cholesky(Operation):
+    def __init__(self):
+        super().__init__()
+
+    def call(self, x):
+        return _cholesky(x)
+
+    def compute_output_spec(self, x):
+        _assert_2d(x)
+        _assert_square(x)
+        return KerasTensor(x.shape, x.dtype)
+
+
+@keras_export(["keras.ops.cholesky", "keras.ops.linalg.cholesky"])
+def cholesky(x):
+    """Computes the Cholesky decomposition of a positive semi-definite matrix.
+
+    Args:
+        x: Input tensor of shape `(..., M, M)`.
+
+    Returns:
+        A tensor of shape `(..., M, M)` representing the lower triangular
+        Cholesky factor of `x`.
+
+    """
+    if any_symbolic_tensors((x,)):
+        return Cholesky().symbolic_call(x)
+    return _cholesky(x)
+
+
+def _cholesky(x):
+    x = backend.convert_to_tensor(x)
+    _assert_2d(x)
+    _assert_square(x)
+    try:
+        return backend.linalg.cholesky(x)
+    except Exception as e:
+        raise ValueError(f"Cholesky decomposition failed: {e}")
+
+
+class Det(Operation):
+    def __init__(self):
+        super().__init__()
+
+    def call(self, x):
+        return _det(x)
+
+    def compute_output_spec(self, x):
+        _assert_2d(x)
+        _assert_square(x)
+        return KerasTensor(x.shape[:-2], x.dtype)
+
+
+@keras_export(["keras.ops.det", "keras.ops.linalg.det"])
+def det(x):
+    """Computes the determinant of a square tensor.
+
+    Args:
+        x: Input tensor of shape `(..., M, M)`.
+
+    Returns:
+        A tensor of shape `(...,)` representing the determinant of `x`.
+
+    """
+    if any_symbolic_tensors((x,)):
+        return Det().symbolic_call(x)
+    return _det(x)
+
+
+def _det(x):
+    x = backend.convert_to_tensor(x)
+    _assert_2d(x)
+    _assert_square(x)
+    return backend.linalg.det(x)
+
+
+class Eig(Operation):
+    def __init__(self):
+        super().__init__()
+
+    def call(self, x):
+        return _eig(x)
+
+    def compute_output_spec(self, x):
+        _assert_square(x)
+        _assert_2d(x)
+        return (
+            KerasTensor(x.shape[:-1], x.dtype),
+            KerasTensor(x.shape, x.dtype),
+        )
+
+
+@keras_export(["keras.ops.eig", "keras.ops.linalg.eig"])
+def eig(x):
+    """Computes the eigenvalues and eigenvectors of a square matrix.
+
+    Args:
+        x: Input tensor of shape `(..., M, M)`.
+
+    Returns:
+        A tuple of two tensors: a tensor of shape `(..., M)` containing
+        eigenvalues and a tensor of shape `(..., M, M)` containing eigenvectors.
+    """
+    if any_symbolic_tensors((x,)):
+        return Eig().symbolic_call(x)
+    return _eig(x)
+
+
+def _eig(x):
+    x = backend.convert_to_tensor(x)
+    _assert_square(x)
+    _assert_2d(x)
+    return backend.linalg.eig(x)
+
+
+class Eigh(Operation):
+    def __init__(self):
+        super().__init__()
+
+    def call(self, x):
+        return _eigh(x)
+
+    def compute_output_spec(self, x):
+        _assert_square(x)
+        _assert_2d(x)
+        return (
+            KerasTensor(x.shape[:-1], x.dtype),
+            KerasTensor(x.shape, x.dtype),
+        )
+
+
+@keras_export(["keras.ops.eigh", "keras.ops.linalg.eigh"])
+def eigh(x):
+    """Computes the eigenvalues and eigenvectors of a complex Hermitian.
+
+    Args:
+        x: Input tensor of shape `(..., M, M)`.
+
+    Returns:
+        A tuple of two tensors: a tensor of shape `(..., M)` containing
+        eigenvalues and a tensor of shape `(..., M, M)` containing eigenvectors.
+
+    """
+    if any_symbolic_tensors((x,)):
+        return Eigh().symbolic_call(x)
+    return _eigh(x)
+
+
+def _eigh(x):
+    x = backend.convert_to_tensor(x)
+    _assert_square(x)
+    _assert_2d(x)
+    return backend.linalg.eigh(x)
+
+
+class Inv(Operation):
+    def __init__(self):
+        super().__init__()
+
+    def call(self, x):
+        return _inv(x)
+
+    def compute_output_spec(self, x):
+        _assert_2d(x)
+        _assert_square(x)
+        return KerasTensor(x.shape, x.dtype)
+
+
+@keras_export(["keras.ops.inv", "keras.ops.linalg.inv"])
+def inv(x):
+    """Computes the inverse of a square tensor.
+
+    Args:
+        x: Input tensor of shape `(..., M, M)`.
+
+    Returns:
+        A tensor of shape `(..., M, M)` representing the inverse of `x`.
+
+    """
+    if any_symbolic_tensors((x,)):
+        return Inv().symbolic_call(x)
+    return _inv(x)
+
+
+def _inv(x):
+    x = backend.convert_to_tensor(x)
+    _assert_2d(x)
+    _assert_square(x)
+    return backend.linalg.inv(x)
+
+
+class LuFactor(Operation):
+    def __init__(self):
+        super().__init__()
+
+    def call(self, x):
+        return _lu_factor(x)
+
+    def compute_output_spec(self, x):
+        _assert_2d(x)
+        batch_shape = x.shape[:-2]
+        m, n = x.shape[-2:]
+        k = min(m, n)
+        return (
+            KerasTensor(batch_shape + (m, n), x.dtype),
+            KerasTensor(batch_shape + (k,), x.dtype),
+        )
+
+
+@keras_export(["keras.ops.lu_factor", "keras.ops.linalg.lu_factor"])
+def lu_factor(x):
+    """Computes the lower-upper decomposition of a square matrix.
+
+    Args:
+        x: A tensor of shape `(..., M, M)`.
+
+    Returns:
+        A tuple of two tensors: a tensor of shape `(..., M, M)` containing the
+        lower and upper triangular matrices and a tensor of shape `(..., M)`
+        containing the pivots.
+
+    """
+    if any_symbolic_tensors((x,)):
+        return LuFactor().symbolic_call(x)
+    return _lu_factor(x)
+
+
+def _lu_factor(x):
+    x = backend.convert_to_tensor(x)
+    _assert_2d(x)
+    if backend.backend() == "tensorflow":
+        try:
+            _assert_square(x)
+        except ValueError as e:
+            raise ValueError(
+                f"LU decomposition failed: {e}. LU decomposition is only "
+                "supported for square matrices in Tensorflow."
+            )
+    return backend.linalg.lu_factor(x)
+
+
+class Norm(Operation):
+    def __init__(self, ord=None, axis=None, keepdims=False):
+        super().__init__()
+        if isinstance(ord, str):
+            if ord not in ("fro", "nuc"):
+                raise ValueError(
+                    "Invalid `ord` argument. "
+                    "Expected one of {'fro', 'nuc'} when using string. "
+                    f"Received: ord={ord}"
+                )
+        if isinstance(axis, int):
+            axis = [axis]
+        self.ord = ord
+        self.axis = axis
+        self.keepdims = keepdims
+
+    def compute_output_spec(self, x):
+        output_dtype = backend.standardize_dtype(x.dtype)
+        if "int" in output_dtype or output_dtype == "bool":
+            output_dtype = backend.floatx()
+        if self.axis is None:
+            axis = tuple(range(len(x.shape)))
+        else:
+            axis = self.axis
+        num_axes = len(axis)
+        if num_axes == 1 and isinstance(self.ord, str):
+            raise ValueError(
+                "Invalid `ord` argument for vector norm. "
+                f"Received: ord={self.ord}"
+            )
+        elif num_axes == 2 and self.ord not in (
+            None,
+            "fro",
+            "nuc",
+            float("inf"),
+            float("-inf"),
+            1,
+            -1,
+            2,
+            -2,
+        ):
+            raise ValueError(
+                "Invalid `ord` argument for matrix norm. "
+                f"Received: ord={self.ord}"
+            )
+        return KerasTensor(
+            reduce_shape(x.shape, axis=self.axis, keepdims=self.keepdims),
+            dtype=output_dtype,
+        )
+
+    def call(self, x):
+        x = backend.convert_to_tensor(x)
+        return backend.linalg.norm(
+            x, ord=self.ord, axis=self.axis, keepdims=self.keepdims
+        )
+
+
+@keras_export(["keras.ops.norm", "keras.ops.linalg.norm"])
+def norm(x, ord=None, axis=None, keepdims=False):
+    """Matrix or vector norm.
+
+    This function is able to return one of eight different matrix norms, or one
+    of an infinite number of vector norms (described below), depending on the
+    value of the `ord` parameter.
+
+    Args:
+        x: Input tensor.
+        ord: Order of the norm (see table under Notes). The default is `None`.
+        axis: If `axis` is an integer, it specifies the axis of `x` along which
+            to compute the vector norms. If `axis` is a 2-tuple, it specifies
+            the axes that hold 2-D matrices, and the matrix norms of these
+            matrices are computed.
+        keepdims: If this is set to `True`, the axes which are reduced are left
+            in the result as dimensions with size one.
+
+    Note:
+        For values of `ord < 1`, the result is, strictly speaking, not a
+        mathematical 'norm', but it may still be useful for various numerical
+        purposes. The following norms can be calculated:
+        - For matrices:
+            - `ord=None`: Frobenius norm
+            - `ord="fro"`: Frobenius norm
+            - `ord="nuc"`: nuclear norm
+            - `ord=np.inf`: `max(sum(abs(x), axis=1))`
+            - `ord=-np.inf`: `min(sum(abs(x), axis=1))`
+            - `ord=0`: not supported
+            - `ord=1`: `max(sum(abs(x), axis=0))`
+            - `ord=-1`: `min(sum(abs(x), axis=0))`
+            - `ord=2`: 2-norm (largest sing. value)
+            - `ord=-2`: smallest singular value
+            - other: not supported
+        - For vectors:
+            - `ord=None`: 2-norm
+            - `ord="fro"`: not supported
+            - `ord="nuc"`: not supported
+            - `ord=np.inf`: `max(abs(x))`
+            - `ord=-np.inf`: `min(abs(x))`
+            - `ord=0`: `sum(x != 0)`
+            - `ord=1`: as below
+            - `ord=-1`: as below
+            - `ord=2`: as below
+            - `ord=-2`: as below
+            - other: `sum(abs(x)**ord)**(1./ord)`
+
+    Returns:
+        Norm of the matrix or vector(s).
+
+    Example:
+
+    >>> x = keras.ops.reshape(keras.ops.arange(9, dtype="float32") - 4, (3, 3))
+    >>> keras.ops.linalg.norm(x)
+    7.7459664
+    """
+    if any_symbolic_tensors((x,)):
+        return Norm(ord=ord, axis=axis, keepdims=keepdims).symbolic_call(x)
+    x = backend.convert_to_tensor(x)
+    return backend.linalg.norm(x, ord=ord, axis=axis, keepdims=keepdims)
+
+
+class Qr(Operation):
+    def __init__(self, mode="reduced"):
+        super().__init__()
+        if mode not in {"reduced", "complete"}:
+            raise ValueError(
+                "`mode` argument value not supported. "
+                "Expected one of {'reduced', 'complete'}. "
+                f"Received: mode={mode}"
+            )
+        self.mode = mode
+
+    def compute_output_spec(self, x):
+        if len(x.shape) < 2:
+            raise ValueError(
+                "Input should have rank >= 2. Received: "
+                f"input.shape = {x.shape}"
+            )
+        m = x.shape[-2]
+        n = x.shape[-1]
+        if m is None or n is None:
+            raise ValueError(
+                "Input should have its last 2 dimensions "
+                "fully-defined. Received: "
+                f"input.shape = {x.shape}"
+            )
+        k = min(m, n)
+        base = tuple(x.shape[:-2])
+        if self.mode == "reduced":
+            return (
+                KerasTensor(shape=base + (m, k), dtype=x.dtype),
+                KerasTensor(shape=base + (k, n), dtype=x.dtype),
+            )
+        # 'complete' mode.
+        return (
+            KerasTensor(shape=base + (m, m), dtype=x.dtype),
+            KerasTensor(shape=base + (m, n), dtype=x.dtype),
+        )
+
+    def call(self, x):
+        x = backend.convert_to_tensor(x)
+        return backend.linalg.qr(x, mode=self.mode)
+
+
+@keras_export(["keras.ops.qr", "keras.ops.linalg.qr"])
+def qr(x, mode="reduced"):
+    """Computes the QR decomposition of a tensor.
+
+    Args:
+        x: Input tensor of shape `(..., M, N)`.
+        mode: A string specifying the mode of the QR decomposition.
+            - 'reduced': Returns the reduced QR decomposition. (default)
+            - 'complete': Returns the complete QR decomposition.
+
+    Returns:
+        A tuple containing two tensors. The first tensor of shape `(..., M, K)`
+        is the orthogonal matrix `q` and the second tensor of shape
+        `(..., K, N)` is the upper triangular matrix `r`, where `K = min(M, N)`.
+
+    Example:
+
+    >>> x = keras.ops.convert_to_tensor([[1., 2.], [3., 4.], [5., 6.]])
+    >>> q, r = qr(x)
+    >>> print(q)
+    array([[-0.16903079  0.897085]
+           [-0.5070925   0.2760267 ]
+           [-0.8451542  -0.34503305]], shape=(3, 2), dtype=float32)
+    """
+    if any_symbolic_tensors((x,)):
+        return Qr(mode=mode).symbolic_call(x)
+    x = backend.convert_to_tensor(x)
+    return backend.linalg.qr(x, mode=mode)
+
+
+class Solve(Operation):
+    def __init__(self):
+        super().__init__()
+
+    def call(self, a, b):
+        return _solve(a, b)
+
+    def compute_output_spec(self, a, b):
+        _assert_2d(a)
+        _assert_square(a)
+        _assert_1d(b)
+        _assert_a_b_compat(a, b)
+        return KerasTensor(b.shape, b.dtype)
+
+
+@keras_export(["keras.ops.solve", "keras.ops.linalg.solve"])
+def solve(a, b):
+    """Solves a linear system of equations given by `a x = b`.
+
+    Args:
+        a: A tensor of shape `(..., M, M)` representing the coefficients matrix.
+        b: A tensor of shape `(..., M)` or `(..., M, N)` representing the
+        right-hand side or "dependent variable" matrix.
+
+    Returns:
+        A tensor of shape `(..., M)` or `(..., M, N)` representing the solution
+        of the linear system. Returned shape is identical to `b`.
+
+    """
+    if any_symbolic_tensors((a, b)):
+        return Solve().symbolic_call(a, b)
+    return _solve(a, b)
+
+
+def _solve(a, b):
+    a = backend.convert_to_tensor(a)
+    b = backend.convert_to_tensor(b)
+    _assert_2d(a)
+    _assert_square(a)
+    _assert_1d(b)
+    _assert_a_b_compat(a, b)
+    return backend.linalg.solve(a, b)
+
+
+class SolveTriangular(Operation):
+    def __init__(self, lower=False):
+        super().__init__()
+        self.lower = lower
+
+    def call(self, a, b):
+        return _solve_triangular(a, b, self.lower)
+
+    def compute_output_spec(self, a, b):
+        _assert_2d(a)
+        _assert_square(a)
+        _assert_1d(b)
+        _assert_a_b_compat(a, b)
+        return KerasTensor(b.shape, b.dtype)
+
+
+@keras_export(
+    ["keras.ops.solve_triangular", "keras.ops.linalg.solve_triangular"]
+)
+def solve_triangular(a, b, lower=False):
+    """Solves a linear system of equations given by `a x = b`.
+
+    Args:
+        a: A tensor of shape `(..., M, M)` representing the coefficients matrix.
+        b: A tensor of shape `(..., M)` or `(..., M, N)` representing the
+        right-hand side or "dependent variable" matrix.
+
+    Returns:
+        A tensor of shape `(..., M)` or `(..., M, N)` representing the solution
+        of the linear system. Returned shape is identical to `b`.
+
+    """
+    if any_symbolic_tensors((a, b)):
+        return SolveTriangular(lower).symbolic_call(a, b)
+    return _solve_triangular(a, b, lower)
+
+
+def _solve_triangular(a, b, lower=False):
+    a = backend.convert_to_tensor(a)
+    b = backend.convert_to_tensor(b)
+    _assert_2d(a)
+    _assert_square(a)
+    _assert_1d(b)
+    _assert_a_b_compat(a, b)
+    return backend.linalg.solve_triangular(a, b, lower)
+
+
+class SVD(Operation):
+    def __init__(self, full_matrices=True, compute_uv=True):
+        super().__init__()
+        self.full_matrices = full_matrices
+        self.compute_uv = compute_uv
+
+    def call(self, x):
+        return _svd(x, self.full_matrices, self.compute_uv)
+
+    def compute_output_spec(self, x):
+        _assert_2d(x)
+        rows, columns = x.shape[-2:]
+        batches = x.shape[:-2]
+        s_shape = batches + (min(rows, columns),)
+        if self.full_matrices:
+            u_shape = batches + (rows, rows)
+            v_shape = batches + (columns, columns)
+        else:
+            u_shape = batches + (rows, min(rows, columns))
+            v_shape = batches + (min(rows, columns), columns)
+
+        if self.compute_uv:
+            return (
+                KerasTensor(u_shape, x.dtype),
+                KerasTensor(s_shape, x.dtype),
+                KerasTensor(v_shape, x.dtype),
+            )
+        return KerasTensor(s_shape, x.dtype)
+
+
+@keras_export(["keras.ops.svd", "keras.ops.linalg.svd"])
+def svd(x, full_matrices=True, compute_uv=True):
+    """Computes the singular value decomposition of a matrix.
+
+    Args:
+        x: Input tensor of shape `(..., M, N)`.
+
+    Returns:
+        A tuple of three tensors: a tensor of shape `(..., M, M)` containing the
+        left singular vectors, a tensor of shape `(..., M, N)` containing the
+        singular values and a tensor of shape `(..., N, N)` containing the
+        right singular vectors.
+
+    """
+    if any_symbolic_tensors((x,)):
+        return SVD(full_matrices, compute_uv).symbolic_call(x)
+    return _svd(x, full_matrices, compute_uv)
+
+
+def _svd(x, full_matrices=True, compute_uv=True):
+    x = backend.convert_to_tensor(x)
+    _assert_2d(x)
+    return backend.linalg.svd(x, full_matrices, compute_uv)
+
+
+class Lstsq(Operation):
+    def __init__(self, rcond=None):
+        super().__init__()
+        self.rcond = rcond
+
+    def call(self, a, b):
+        return backend.linalg.lstsq(a, b, rcond=self.rcond)
+
+    def compute_output_spec(self, a, b):
+        if len(a.shape) != 2:
+            raise ValueError(
+                "Expected a to have rank 2. " f"Received: a.shape={a.shape}"
+            )
+        if len(b.shape) not in (1, 2):
+            raise ValueError(
+                "Expected b to have rank 1 or 2. "
+                f"Received: b.shape={b.shape}"
+            )
+        m, n = a.shape
+        if b.shape[0] != m:
+            raise ValueError(
+                "Expected b.shape[0] to be equal to "
+                "a.shape[0]. Received: "
+                f"a.shape={a.shape}, b.shape={b.shape}"
+            )
+        if len(b.shape) == 2:
+            k = b.shape[1]
+            x = KerasTensor((n, k), dtype=a.dtype)
+        else:
+            x = KerasTensor((n,), dtype=a.dtype)
+        return x
+
+
+@keras_export(["keras.ops.lstsq", "keras.ops.linalg.lstsq"])
+def lstsq(a, b, rcond=None):
+    """Return the least-squares solution to a linear matrix equation.
+
+    Computes the vector x that approximately solves the equation
+    `a @ x = b`. The equation may be under-, well-, or over-determined
+    (i.e., the number of linearly independent rows of a can be less than,
+    equal to, or greater than its number of linearly independent columns).
+    If a is square and of full rank, then `x` (but for round-off error)
+    is the exact solution of the equation. Else, `x` minimizes the
+    L2 norm of `b - a * x`.
+
+    If there are multiple minimizing solutions,
+    the one with the smallest L2 norm  is returned.
+
+    Args:
+        a: "Coefficient" matrix of shape `(M, N)`.
+        b: Ordinate or "dependent variable" values,
+            of shape `(M,)` or `(M, K)`.
+            If `b` is two-dimensional, the least-squares solution
+            is calculated for each of the K columns of `b`.
+        rcond: Cut-off ratio for small singular values of `a`.
+            For the purposes of rank determination,
+            singular values are treated as zero if they are
+            smaller than rcond times the largest
+            singular value of `a`.
+
+    Returns:
+        Tensor with shape `(N,)` or `(N, K)` containing
+        the least-squares solutions.
+
+    **NOTE:** The output differs from `numpy.linalg.lstsq`.
+    NumPy returns a tuple with four elements, the first of which
+    being the least-squares solutions and the others
+    being essentially never used.
+    Keras only returns the first value. This is done both
+    to ensure consistency across backends (which cannot be achieved
+    for the other values) and to simplify the API.
+    """
+    if any_symbolic_tensors((a, b)):
+        return Lstsq(rcond=rcond).symbolic_call(a, b)
+    return backend.linalg.lstsq(a, b, rcond=rcond)
+
+
+def _assert_1d(*arrays):
+    for a in arrays:
+        if a.ndim < 1:
+            raise ValueError(
+                "Expected input to have rank >= 1. "
+                "Received scalar input {a}."
+            )
+
+
+def _assert_2d(*arrays):
+    for a in arrays:
+        if a.ndim < 2:
+            raise ValueError(
+                "Expected input to have rank >= 2. "
+                "Received input with shape {a.shape}."
+            )
+
+
+def _assert_square(*arrays):
+    for a in arrays:
+        m, n = a.shape[-2:]
+        if m != n:
+            raise ValueError(
+                "Expected a square matrix. "
+                f"Received non-square input with shape {a.shape}"
+            )
+
+
+def _assert_a_b_compat(a, b):
+    if a.ndim == b.ndim:
+        if a.shape[-2] != b.shape[-2]:
+            raise ValueError(
+                "Incompatible shapes between `a` and `b`. "
+                "Expected `a.shape[-2] == b.shape[-2]`. "
+                f"Received: a.shape={a.shape}, b.shape={b.shape}"
+            )
+    elif a.ndim == b.ndim - 1:
+        if a.shape[-1] != b.shape[-1]:
+            raise ValueError(
+                "Incompatible shapes between `a` and `b`. "
+                "Expected `a.shape[-1] == b.shape[-1]`. "
+                f"Received: a.shape={a.shape}, b.shape={b.shape}"
+            )

SwarmUI/dlbackend/ComfyUI/venv/lib/python3.10/site-packages/keras/src/ops/math.py

@@ -0,0 +1,1046 @@
+"""Commonly used math operations not included in NumPy."""
+
+from keras.src import backend
+from keras.src.api_export import keras_export
+from keras.src.backend import KerasTensor
+from keras.src.backend import any_symbolic_tensors
+from keras.src.ops.operation import Operation
+from keras.src.ops.operation_utils import reduce_shape
+
+
+def _segment_reduce_validation(data, segment_ids):
+    data_shape = data.shape
+    segment_ids_shape = segment_ids.shape
+    if len(segment_ids_shape) > 1:
+        raise ValueError(
+            "Argument `segment_ids` should be an 1-D vector, got shape: "
+            f"{len(segment_ids_shape)}. Consider either flatten input with "
+            "segment_ids.reshape((-1)) and "
+            "data.reshape((-1, ) + data.shape[len(segment_ids.shape):]) or "
+            "vectorize with vmap."
+        )
+    if (
+        segment_ids_shape[0] is not None
+        and data_shape[0] is not None
+        and segment_ids_shape[0] != data_shape[0]
+    ):
+        raise ValueError(
+            "Argument `segment_ids` and `data` should have same leading "
+            f"dimension. Got {segment_ids_shape} v.s. "
+            f"{data_shape}."
+        )
+
+
+class SegmentReduction(Operation):
+    def __init__(self, num_segments=None, sorted=False):
+        super().__init__()
+        self.num_segments = num_segments
+        self.sorted = sorted
+
+    def compute_output_spec(self, data, _):
+        output_shape = (self.num_segments,) + tuple(data.shape[1:])
+        return KerasTensor(shape=output_shape, dtype=data.dtype)
+
+
+class SegmentSum(SegmentReduction):
+    def call(self, data, segment_ids):
+        _segment_reduce_validation(data, segment_ids)
+        return backend.math.segment_sum(
+            data,
+            segment_ids,
+            num_segments=self.num_segments,
+            sorted=self.sorted,
+        )
+
+
+@keras_export("keras.ops.segment_sum")
+def segment_sum(data, segment_ids, num_segments=None, sorted=False):
+    """Computes the sum of segments in a tensor.
+
+    Args:
+        data: Input tensor.
+        segment_ids: A N-D tensor containing segment indices for each
+            element in `data`. Num dims for segment ids should be strictly
+            smaller or equal to number of dims in data.
+        num_segments: An integer representing the total number of
+            segments. If not specified, it is inferred from the maximum
+            value in `segment_ids`.
+        sorted: A boolean indicating whether `segment_ids` is sorted.
+            Defaults to `False`.
+
+    Returns:
+        A tensor containing the sum of segments, where each element
+        represents the sum of the corresponding segment in `data`.
+
+    Example:
+
+    >>> data = keras.ops.convert_to_tensor([1, 2, 10, 20, 100, 200])
+    >>> segment_ids = keras.ops.convert_to_tensor([0, 0, 1, 1, 2, 2])
+    >>> num_segments = 3
+    >>> keras.ops.segment_sum(data, segment_ids,num_segments)
+    array([3, 30, 300], dtype=int32)
+    """
+    _segment_reduce_validation(data, segment_ids)
+    if any_symbolic_tensors((data,)):
+        return SegmentSum(num_segments, sorted).symbolic_call(data, segment_ids)
+    return backend.math.segment_sum(
+        data, segment_ids, num_segments=num_segments, sorted=sorted
+    )
+
+
+class SegmentMax(SegmentReduction):
+    def call(self, data, segment_ids):
+        _segment_reduce_validation(data, segment_ids)
+        return backend.math.segment_max(
+            data,
+            segment_ids,
+            num_segments=self.num_segments,
+            sorted=self.sorted,
+        )
+
+
+@keras_export("keras.ops.segment_max")
+def segment_max(data, segment_ids, num_segments=None, sorted=False):
+    """Computes the max of segments in a tensor.
+
+    Args:
+        data: Input tensor.
+        segment_ids: A N-D tensor containing segment indices for each
+            element in `data`. data.shape[:len(segment_ids.shape)] should match.
+        num_segments: An integer representing the total number of
+            segments. If not specified, it is inferred from the maximum
+            value in `segment_ids`.
+        sorted: A boolean indicating whether `segment_ids` is sorted.
+            Defaults to `False`.
+
+    Returns:
+        A tensor containing the max of segments, where each element
+        represents the max of the corresponding segment in `data`.
+
+    Example:
+
+    >>> data = keras.ops.convert_to_tensor([1, 2, 10, 20, 100, 200])
+    >>> segment_ids = keras.ops.convert_to_tensor([0, 0, 1, 1, 2, 2])
+    >>> num_segments = 3
+    >>> keras.ops.segment_max(data, segment_ids, num_segments)
+    array([2, 20, 200], dtype=int32)
+    """
+    _segment_reduce_validation(data, segment_ids)
+    if any_symbolic_tensors((data,)):
+        return SegmentMax(num_segments, sorted).symbolic_call(data, segment_ids)
+    return backend.math.segment_max(
+        data, segment_ids, num_segments=num_segments, sorted=sorted
+    )
+
+
+class TopK(Operation):
+    def __init__(self, k, sorted=False):
+        super().__init__()
+        self.k = k
+        self.sorted = sorted
+
+    def compute_output_spec(self, x):
+        output_shape = list(x.shape)
+        output_shape[-1] = self.k
+        # Return a tuple (values, indices).
+        return (
+            KerasTensor(shape=output_shape, dtype=x.dtype),
+            KerasTensor(shape=output_shape, dtype="int32"),
+        )
+
+    def call(self, x):
+        return backend.math.top_k(x, self.k, self.sorted)
+
+
+@keras_export("keras.ops.top_k")
+def top_k(x, k, sorted=True):
+    """Finds the top-k values and their indices in a tensor.
+
+    Args:
+        x: Input tensor.
+        k: An integer representing the number of top elements to retrieve.
+        sorted: A boolean indicating whether to sort the output in
+        descending order. Defaults to `True`.
+
+    Returns:
+        A tuple containing two tensors. The first tensor contains the
+        top-k values, and the second tensor contains the indices of the
+        top-k values in the input tensor.
+
+    Example:
+
+    >>> x = keras.ops.convert_to_tensor([5, 2, 7, 1, 9, 3])
+    >>> values, indices = top_k(x, k=3)
+    >>> print(values)
+    array([9 7 5], shape=(3,), dtype=int32)
+    >>> print(indices)
+    array([4 2 0], shape=(3,), dtype=int32)
+
+    """
+    if any_symbolic_tensors((x,)):
+        return TopK(k, sorted).symbolic_call(x)
+    return backend.math.top_k(x, k, sorted)
+
+
+class InTopK(Operation):
+    def __init__(self, k):
+        super().__init__()
+        self.k = k
+
+    def compute_output_spec(self, targets, predictions):
+        return KerasTensor(shape=targets.shape, dtype="bool")
+
+    def call(self, targets, predictions):
+        return backend.math.in_top_k(targets, predictions, self.k)
+
+
+@keras_export("keras.ops.in_top_k")
+def in_top_k(targets, predictions, k):
+    """Checks if the targets are in the top-k predictions.
+
+    Args:
+        targets: A tensor of true labels.
+        predictions: A tensor of predicted labels.
+        k: An integer representing the number of predictions to consider.
+
+    Returns:
+        A boolean tensor of the same shape as `targets`, where each element
+        indicates whether the corresponding target is in the top-k predictions.
+
+    Example:
+
+    >>> targets = keras.ops.convert_to_tensor([2, 5, 3])
+    >>> predictions = keras.ops.convert_to_tensor(
+    ... [[0.1, 0.4, 0.6, 0.9, 0.5],
+    ...  [0.1, 0.7, 0.9, 0.8, 0.3],
+    ...  [0.1, 0.6, 0.9, 0.9, 0.5]])
+    >>> in_top_k(targets, predictions, k=3)
+    array([ True False  True], shape=(3,), dtype=bool)
+    """
+    if any_symbolic_tensors((targets, predictions)):
+        return InTopK(k).symbolic_call(targets, predictions)
+    return backend.math.in_top_k(targets, predictions, k)
+
+
+class Logsumexp(Operation):
+    def __init__(self, axis=None, keepdims=False):
+        super().__init__()
+        self.axis = axis
+        self.keepdims = keepdims
+
+    def compute_output_spec(self, x):
+        output_shape = reduce_shape(x.shape, self.axis, self.keepdims)
+        return KerasTensor(shape=output_shape)
+
+    def call(self, x):
+        return backend.math.logsumexp(x, axis=self.axis, keepdims=self.keepdims)
+
+
+@keras_export("keras.ops.logsumexp")
+def logsumexp(x, axis=None, keepdims=False):
+    """Computes the logarithm of sum of exponentials of elements in a tensor.
+
+    Args:
+        x: Input tensor.
+        axis: An integer or a tuple of integers specifying the axis/axes
+            along which to compute the sum. If `None`, the sum is computed
+            over all elements. Defaults to `None`.
+        keepdims: A boolean indicating whether to keep the dimensions of
+            the input tensor when computing the sum. Defaults to `False`.
+
+    Returns:
+        A tensor containing the logarithm of the sum of exponentials of
+        elements in `x`.
+
+    Example:
+
+    >>> x = keras.ops.convert_to_tensor([1., 2., 3.])
+    >>> logsumexp(x)
+    3.407606
+    """
+    if any_symbolic_tensors((x,)):
+        return Logsumexp(axis, keepdims).symbolic_call(x)
+    return backend.math.logsumexp(x, axis=axis, keepdims=keepdims)
+
+
+class ExtractSequences(Operation):
+    def __init__(self, sequence_length, sequence_stride):
+        super().__init__()
+        self.sequence_length = sequence_length
+        self.sequence_stride = sequence_stride
+
+    def compute_output_spec(self, x):
+        if len(x.shape) < 1:
+            raise ValueError(
+                f"Input should have rank >= 1. "
+                f"Received: input.shape = {x.shape}"
+            )
+        if x.shape[-1] is not None:
+            num_sequences = (
+                1 + (x.shape[-1] - self.sequence_length) // self.sequence_stride
+            )
+        else:
+            num_sequences = None
+        new_shape = x.shape[:-1] + (num_sequences, self.sequence_length)
+        return KerasTensor(shape=new_shape, dtype=x.dtype)
+
+    def call(self, x):
+        return backend.math.extract_sequences(
+            x,
+            sequence_length=self.sequence_length,
+            sequence_stride=self.sequence_stride,
+        )
+
+
+@keras_export("keras.ops.extract_sequences")
+def extract_sequences(x, sequence_length, sequence_stride):
+    """Expands the dimension of last axis into sequences of `sequence_length`.
+
+    Slides a window of size `sequence_length` over the last axis of the input
+    with a stride of `sequence_stride`, replacing the last axis with
+    `[num_sequences, sequence_length]` sequences.
+
+    If the dimension along the last axis is N, the number of sequences can be
+    computed by:
+
+    `num_sequences = 1 + (N - sequence_length) // sequence_stride`
+
+    Args:
+        x: Input tensor.
+        sequence_length: An integer representing the sequences length.
+        sequence_stride: An integer representing the sequences hop size.
+
+    Returns:
+        A tensor of sequences with shape [..., num_sequences, sequence_length].
+
+    Example:
+
+    >>> x = keras.ops.convert_to_tensor([1, 2, 3, 4, 5, 6])
+    >>> extract_sequences(x, 3, 2)
+    array([[1, 2, 3],
+       [3, 4, 5]])
+    """
+    if any_symbolic_tensors((x,)):
+        return ExtractSequences(sequence_length, sequence_stride).symbolic_call(
+            x
+        )
+    return backend.math.extract_sequences(x, sequence_length, sequence_stride)
+
+
+class FFT(Operation):
+    def __init__(self, axis=-1):
+        super().__init__()
+        self.axis = axis
+
+    def compute_output_spec(self, x):
+        if not isinstance(x, (tuple, list)) or len(x) != 2:
+            raise ValueError(
+                "Input `x` should be a tuple of two tensors - real and "
+                f"imaginary. Received: x={x}"
+            )
+
+        real, imag = x
+        # Both real and imaginary parts should have the same shape.
+        if real.shape != imag.shape:
+            raise ValueError(
+                "Input `x` should be a tuple of two tensors - real and "
+                "imaginary. Both the real and imaginary parts should have the "
+                f"same shape. Received: x[0].shape = {real.shape}, "
+                f"x[1].shape = {imag.shape}"
+            )
+
+        # We are calculating 1D FFT. Hence, rank >= 1.
+        if len(real.shape) < 1:
+            raise ValueError(
+                f"Input should have rank >= 1. "
+                f"Received: input.shape = {real.shape}"
+            )
+
+        # The axis along which we are calculating FFT should be fully-defined.
+        m = real.shape[-1]
+        if m is None:
+            raise ValueError(
+                f"Input should have its {self.axis}th axis fully-defined. "
+                f"Received: input.shape = {real.shape}"
+            )
+
+        return (
+            KerasTensor(shape=real.shape, dtype=real.dtype),
+            KerasTensor(shape=imag.shape, dtype=imag.dtype),
+        )
+
+    def call(self, x):
+        return backend.math.fft(x)
+
+
+@keras_export("keras.ops.fft")
+def fft(x):
+    """Computes the Fast Fourier Transform along last axis of input.
+
+    Args:
+        x: Tuple of the real and imaginary parts of the input tensor. Both
+            tensors in the tuple should be of floating type.
+
+    Returns:
+        A tuple containing two tensors - the real and imaginary parts of the
+        output tensor.
+
+    Example:
+
+    >>> x = (
+    ...     keras.ops.convert_to_tensor([1., 2.]),
+    ...     keras.ops.convert_to_tensor([0., 1.]),
+    ... )
+    >>> fft(x)
+    (array([ 3., -1.], dtype=float32), array([ 1., -1.], dtype=float32))
+    """
+    if any_symbolic_tensors(x):
+        return FFT().symbolic_call(x)
+    return backend.math.fft(x)
+
+
+class FFT2(Operation):
+    def __init__(self):
+        super().__init__()
+        self.axes = (-2, -1)
+
+    def compute_output_spec(self, x):
+        if not isinstance(x, (tuple, list)) or len(x) != 2:
+            raise ValueError(
+                "Input `x` should be a tuple of two tensors - real and "
+                f"imaginary. Received: x={x}"
+            )
+
+        real, imag = x
+        # Both real and imaginary parts should have the same shape.
+        if real.shape != imag.shape:
+            raise ValueError(
+                "Input `x` should be a tuple of two tensors - real and "
+                "imaginary. Both the real and imaginary parts should have the "
+                f"same shape. Received: x[0].shape = {real.shape}, "
+                f"x[1].shape = {imag.shape}"
+            )
+        # We are calculating 2D FFT. Hence, rank >= 2.
+        if len(real.shape) < 2:
+            raise ValueError(
+                f"Input should have rank >= 2. "
+                f"Received: input.shape = {real.shape}"
+            )
+
+        # The axes along which we are calculating FFT should be fully-defined.
+        m = real.shape[self.axes[0]]
+        n = real.shape[self.axes[1]]
+        if m is None or n is None:
+            raise ValueError(
+                f"Input should have its {self.axes} axes fully-defined. "
+                f"Received: input.shape = {real.shape}"
+            )
+
+        return (
+            KerasTensor(shape=real.shape, dtype=real.dtype),
+            KerasTensor(shape=imag.shape, dtype=imag.dtype),
+        )
+
+    def call(self, x):
+        return backend.math.fft2(x)
+
+
+@keras_export("keras.ops.fft2")
+def fft2(x):
+    """Computes the 2D Fast Fourier Transform along the last two axes of input.
+
+    Args:
+        x: Tuple of the real and imaginary parts of the input tensor. Both
+            tensors in the tuple should be of floating type.
+
+    Returns:
+        A tuple containing two tensors - the real and imaginary parts of the
+        output.
+
+    Example:
+
+    >>> x = (
+    ...     keras.ops.convert_to_tensor([[1., 2.], [2., 1.]]),
+    ...     keras.ops.convert_to_tensor([[0., 1.], [1., 0.]]),
+    ... )
+    >>> fft2(x)
+    (array([[ 6.,  0.],
+        [ 0., -2.]], dtype=float32), array([[ 2.,  0.],
+        [ 0., -2.]], dtype=float32))
+    """
+    if any_symbolic_tensors(x):
+        return FFT2().symbolic_call(x)
+    return backend.math.fft2(x)
+
+
+class IFFT2(Operation):
+    def __init__(self):
+        super().__init__()
+        self.axes = (-2, -1)
+
+    def compute_output_spec(self, x):
+        if not isinstance(x, (tuple, list)) or len(x) != 2:
+            raise ValueError(
+                "Input `x` should be a tuple of two tensors - real and "
+                f"imaginary. Received: x={x}"
+            )
+
+        real, imag = x
+        # Both real and imaginary parts should have the same shape.
+        if real.shape != imag.shape:
+            raise ValueError(
+                "Input `x` should be a tuple of two tensors - real and "
+                "imaginary. Both the real and imaginary parts should have the "
+                f"same shape. Received: x[0].shape = {real.shape}, "
+                f"x[1].shape = {imag.shape}"
+            )
+        # We are calculating 2D IFFT. Hence, rank >= 2.
+        if len(real.shape) < 2:
+            raise ValueError(
+                f"Input should have rank >= 2. "
+                f"Received: input.shape = {real.shape}"
+            )
+
+        # The axes along which we are calculating IFFT should be fully-defined.
+        m = real.shape[self.axes[0]]
+        n = real.shape[self.axes[1]]
+        if m is None or n is None:
+            raise ValueError(
+                f"Input should have its {self.axes} axes fully-defined. "
+                f"Received: input.shape = {real.shape}"
+            )
+
+        return (
+            KerasTensor(shape=real.shape, dtype=real.dtype),
+            KerasTensor(shape=imag.shape, dtype=imag.dtype),
+        )
+
+    def call(self, x):
+        return backend.math.ifft2(x)
+
+
+@keras_export("keras.ops.ifft2")
+def ifft2(x):
+    """Computes the 2D Inverse Fast Fourier Transform along the last two axes of
+        input.
+
+    Args:
+        x: Tuple of the real and imaginary parts of the input tensor. Both
+            tensors in the tuple should be of floating type.
+
+    Returns:
+        A tuple containing two tensors - the real and imaginary parts of the
+        output.
+
+    Example:
+
+    >>> x = (
+    ...     keras.ops.convert_to_tensor([[1., 2.], [2., 1.]]),
+    ...     keras.ops.convert_to_tensor([[0., 1.], [1., 0.]]),
+    ... )
+    >>> ifft2(x)
+    (array([[ 6.,  0.],
+        [ 0., -2.]], dtype=float32), array([[ 2.,  0.],
+        [ 0., -2.]], dtype=float32))
+    """
+    if any_symbolic_tensors(x):
+        return IFFT2().symbolic_call(x)
+    return backend.math.ifft2(x)
+
+
+class RFFT(Operation):
+    def __init__(self, fft_length=None):
+        super().__init__()
+        self.fft_length = fft_length
+
+    def compute_output_spec(self, x):
+        # We are calculating 1D RFFT. Hence, rank >= 1.
+        if len(x.shape) < 1:
+            raise ValueError(
+                f"Input should have rank >= 1. "
+                f"Received: input.shape = {x.shape}"
+            )
+
+        if self.fft_length is not None:
+            new_last_dimension = self.fft_length // 2 + 1
+        else:
+            if x.shape[-1] is not None:
+                new_last_dimension = x.shape[-1] // 2 + 1
+            else:
+                new_last_dimension = None
+        new_shape = x.shape[:-1] + (new_last_dimension,)
+
+        return (
+            KerasTensor(shape=new_shape, dtype=x.dtype),
+            KerasTensor(shape=new_shape, dtype=x.dtype),
+        )
+
+    def call(self, x):
+        return backend.math.rfft(x, fft_length=self.fft_length)
+
+
+@keras_export("keras.ops.rfft")
+def rfft(x, fft_length=None):
+    """Real-valued Fast Fourier Transform along the last axis of the input.
+
+    Computes the 1D Discrete Fourier Transform of a real-valued signal over the
+    inner-most dimension of input.
+
+    Since the Discrete Fourier Transform of a real-valued signal is
+    Hermitian-symmetric, RFFT only returns the `fft_length / 2 + 1` unique
+    components of the FFT: the zero-frequency term, followed by the
+    `fft_length / 2` positive-frequency terms.
+
+    Along the axis RFFT is computed on, if `fft_length` is smaller than the
+    corresponding dimension of the input, the dimension is cropped. If it is
+    larger, the dimension is padded with zeros.
+
+    Args:
+        x: Input tensor.
+        fft_length: An integer representing the number of the fft length. If not
+            specified, it is inferred from the length of the last axis of `x`.
+            Defaults to `None`.
+
+    Returns:
+        A tuple containing two tensors - the real and imaginary parts of the
+        output.
+
+    Examples:
+
+    >>> x = keras.ops.convert_to_tensor([0.0, 1.0, 2.0, 3.0, 4.0])
+    >>> rfft(x)
+    (array([10.0, -2.5, -2.5]), array([0.0, 3.4409548, 0.81229924]))
+
+    >>> rfft(x, 3)
+    (array([3.0, -1.5]), array([0.0, 0.8660254]))
+    """
+    if any_symbolic_tensors((x,)):
+        return RFFT(fft_length).symbolic_call(x)
+    return backend.math.rfft(x, fft_length)
+
+
+class IRFFT(Operation):
+    def __init__(self, fft_length=None):
+        super().__init__()
+        self.fft_length = fft_length
+
+    def compute_output_spec(self, x):
+        if not isinstance(x, (tuple, list)) or len(x) != 2:
+            raise ValueError(
+                "Input `x` should be a tuple of two tensors - real and "
+                f"imaginary. Received: x={x}"
+            )
+        real, imag = x
+        # Both real and imaginary parts should have the same shape.
+        if real.shape != imag.shape:
+            raise ValueError(
+                "Input `x` should be a tuple of two tensors - real and "
+                "imaginary. Both the real and imaginary parts should have the "
+                f"same shape. Received: x[0].shape = {real.shape}, "
+                f"x[1].shape = {imag.shape}"
+            )
+        # We are calculating 1D IRFFT. Hence, rank >= 1.
+        if len(real.shape) < 1:
+            raise ValueError(
+                f"Input should have rank >= 1. "
+                f"Received: input.shape = {real.shape}"
+            )
+
+        if self.fft_length is not None:
+            new_last_dimension = self.fft_length
+        else:
+            if real.shape[-1] is not None:
+                new_last_dimension = 2 * (real.shape[-1] - 1)
+            else:
+                new_last_dimension = None
+        new_shape = real.shape[:-1] + (new_last_dimension,)
+        return KerasTensor(shape=new_shape, dtype=real.dtype)
+
+    def call(self, x):
+        return backend.math.irfft(x, fft_length=self.fft_length)
+
+
+@keras_export("keras.ops.irfft")
+def irfft(x, fft_length=None):
+    """Inverse real-valued Fast Fourier transform along the last axis.
+
+    Computes the inverse 1D Discrete Fourier Transform of a real-valued signal
+    over the inner-most dimension of input.
+
+    The inner-most dimension of the input is assumed to be the result of RFFT:
+    the `fft_length / 2 + 1` unique components of the DFT of a real-valued
+    signal. If `fft_length` is not provided, it is computed from the size of the
+    inner-most dimension of the input `(fft_length = 2 * (inner - 1))`. If the
+    FFT length used to compute is odd, it should be provided since it cannot
+    be inferred properly.
+
+    Along the axis IRFFT is computed on, if `fft_length / 2 + 1` is smaller than
+    the corresponding dimension of the input, the dimension is cropped. If it is
+    larger, the dimension is padded with zeros.
+
+    Args:
+        x: Tuple of the real and imaginary parts of the input tensor. Both
+            tensors in the tuple should be of floating type.
+        fft_length: An integer representing the number of the fft length. If not
+            specified, it is inferred from the length of the last axis of `x`.
+            Defaults to `None`.
+
+    Returns:
+        A tensor containing the inverse real-valued Fast Fourier Transform
+        along the last axis of `x`.
+
+    Examples:
+
+    >>> real = keras.ops.convert_to_tensor([0.0, 1.0, 2.0, 3.0, 4.0])
+    >>> imag = keras.ops.convert_to_tensor([0.0, 1.0, 2.0, 3.0, 4.0])
+    >>> irfft((real, imag))
+    array([0.66666667, -0.9106836, 0.24401694])
+
+    >>> irfft(rfft(real, 5), 5)
+    array([0.0, 1.0, 2.0, 3.0, 4.0])
+    """
+    if any_symbolic_tensors(x):
+        return IRFFT(fft_length).symbolic_call(x)
+    return backend.math.irfft(x, fft_length)
+
+
+class STFT(Operation):
+    def __init__(
+        self,
+        sequence_length,
+        sequence_stride,
+        fft_length,
+        window="hann",
+        center=True,
+    ):
+        super().__init__()
+        self.sequence_length = sequence_length
+        self.sequence_stride = sequence_stride
+        self.fft_length = fft_length
+        self.window = window
+        self.center = center
+
+    def compute_output_spec(self, x):
+        if x.shape[-1] is not None:
+            padded = 0 if self.center is False else (self.fft_length // 2) * 2
+            num_sequences = (
+                1
+                + (x.shape[-1] + padded - self.fft_length)
+                // self.sequence_stride
+            )
+        else:
+            num_sequences = None
+        new_shape = x.shape[:-1] + (num_sequences, self.fft_length // 2 + 1)
+        return (
+            KerasTensor(shape=new_shape, dtype=x.dtype),
+            KerasTensor(shape=new_shape, dtype=x.dtype),
+        )
+
+    def call(self, x):
+        return backend.math.stft(
+            x,
+            sequence_length=self.sequence_length,
+            sequence_stride=self.sequence_stride,
+            fft_length=self.fft_length,
+            window=self.window,
+            center=self.center,
+        )
+
+
+@keras_export("keras.ops.stft")
+def stft(
+    x, sequence_length, sequence_stride, fft_length, window="hann", center=True
+):
+    """Short-Time Fourier Transform along the last axis of the input.
+
+    The STFT computes the Fourier transform of short overlapping windows of the
+    input. This giving frequency components of the signal as they change over
+    time.
+
+    Args:
+        x: Input tensor.
+        sequence_length: An integer representing the sequence length.
+        sequence_stride: An integer representing the sequence hop size.
+        fft_length: An integer representing the size of the FFT to apply. If not
+            specified, uses the smallest power of 2 enclosing `sequence_length`.
+        window: A string, a tensor of the window or `None`. If `window` is a
+            string, available values are `"hann"` and `"hamming"`. If `window`
+            is a tensor, it will be used directly as the window and its length
+            must be `sequence_length`. If `window` is `None`, no windowing is
+            used. Defaults to `"hann"`.
+        center: Whether to pad `x` on both sides so that the t-th sequence is
+            centered at time `t * sequence_stride`. Otherwise, the t-th sequence
+            begins at time `t * sequence_stride`. Defaults to `True`.
+
+    Returns:
+        A tuple containing two tensors - the real and imaginary parts of the
+        STFT output.
+
+    Example:
+
+    >>> x = keras.ops.convert_to_tensor([0.0, 1.0, 2.0, 3.0, 4.0])
+    >>> stft(x, 3, 2, 3)
+    (array([[0.75, -0.375],
+       [3.75, -1.875],
+       [5.25, -2.625]]), array([[0.0, 0.64951905],
+       [0.0, 0.64951905],
+       [0.0, -0.64951905]]))
+    """
+    if any_symbolic_tensors((x,)):
+        return STFT(
+            sequence_length=sequence_length,
+            sequence_stride=sequence_stride,
+            fft_length=fft_length,
+            window=window,
+            center=center,
+        ).symbolic_call(x)
+    return backend.math.stft(
+        x,
+        sequence_length=sequence_length,
+        sequence_stride=sequence_stride,
+        fft_length=fft_length,
+        window=window,
+        center=center,
+    )
+
+
+class ISTFT(Operation):
+    def __init__(
+        self,
+        sequence_length,
+        sequence_stride,
+        fft_length,
+        length=None,
+        window="hann",
+        center=True,
+    ):
+        super().__init__()
+        self.sequence_length = sequence_length
+        self.sequence_stride = sequence_stride
+        self.fft_length = fft_length
+        self.length = length
+        self.window = window
+        self.center = center
+
+    def compute_output_spec(self, x):
+        if not isinstance(x, (tuple, list)) or len(x) != 2:
+            raise ValueError(
+                "Input `x` should be a tuple of two tensors - real and "
+                f"imaginary. Received: x={x}"
+            )
+        real, imag = x
+        # Both real and imaginary parts should have the same shape.
+        if real.shape != imag.shape:
+            raise ValueError(
+                "Input `x` should be a tuple of two tensors - real and "
+                "imaginary. Both the real and imaginary parts should have the "
+                f"same shape. Received: x[0].shape = {real.shape}, "
+                f"x[1].shape = {imag.shape}"
+            )
+        if len(real.shape) < 2:
+            raise ValueError(
+                f"Input should have rank >= 2. "
+                f"Received: input.shape = {real.shape}"
+            )
+        if real.shape[-2] is not None:
+            output_size = (
+                real.shape[-2] - 1
+            ) * self.sequence_stride + self.fft_length
+            if self.length is not None:
+                output_size = self.length
+            elif self.center:
+                output_size = output_size - (self.fft_length // 2) * 2
+        else:
+            output_size = None
+        new_shape = real.shape[:-2] + (output_size,)
+        return KerasTensor(shape=new_shape, dtype=real.dtype)
+
+    def call(self, x):
+        return backend.math.istft(
+            x,
+            sequence_length=self.sequence_length,
+            sequence_stride=self.sequence_stride,
+            fft_length=self.fft_length,
+            length=self.length,
+            window=self.window,
+            center=self.center,
+        )
+
+
+@keras_export("keras.ops.istft")
+def istft(
+    x,
+    sequence_length,
+    sequence_stride,
+    fft_length,
+    length=None,
+    window="hann",
+    center=True,
+):
+    """Inverse Short-Time Fourier Transform along the last axis of the input.
+
+    To reconstruct an original waveform, the parameters should be the same in
+    `stft`.
+
+    Args:
+        x: Tuple of the real and imaginary parts of the input tensor. Both
+            tensors in the tuple should be of floating type.
+        sequence_length: An integer representing the sequence length.
+        sequence_stride: An integer representing the sequence hop size.
+        fft_length: An integer representing the size of the FFT that produced
+            `stft`. Should be of type `int32`.
+        length: An integer representing the output is clipped to exactly length.
+            If not specified, no padding or clipping take place. Defaults to
+            `None`.
+        window: A string, a tensor of the window or `None`. If `window` is a
+            string, available values are `"hann"` and `"hamming"`. If `window`
+            is a tensor, it will be used directly as the window and its length
+            must be `sequence_length`. If `window` is `None`, no windowing is
+            used. Defaults to `"hann"`.
+        center: Whether `x` was padded on both sides so that the t-th sequence
+            is centered at time `t * sequence_stride`. Defaults to `True`.
+
+    Returns:
+        A tensor containing the inverse Short-Time Fourier Transform along the
+        last axis of `x`.
+
+    Example:
+
+    >>> x = keras.ops.convert_to_tensor([0.0, 1.0, 2.0, 3.0, 4.0])
+    >>> istft(stft(x, 1, 1, 1), 1, 1, 1)
+    array([0.0, 1.0, 2.0, 3.0, 4.0])
+    """
+    if any_symbolic_tensors(x):
+        return ISTFT(
+            sequence_length=sequence_length,
+            sequence_stride=sequence_stride,
+            fft_length=fft_length,
+            window=window,
+            center=center,
+        ).symbolic_call(x)
+    return backend.math.istft(
+        x,
+        sequence_length=sequence_length,
+        sequence_stride=sequence_stride,
+        fft_length=fft_length,
+        length=length,
+        window=window,
+        center=center,
+    )
+
+
+class Rsqrt(Operation):
+    def call(self, x):
+        x = backend.convert_to_tensor(x)
+        return backend.math.rsqrt(x)
+
+    def compute_output_spec(self, x):
+        return KerasTensor(x.shape, dtype=x.dtype)
+
+
+@keras_export("keras.ops.rsqrt")
+def rsqrt(x):
+    """Computes reciprocal of square root of x element-wise.
+
+    Args:
+        x: input tensor
+
+    Returns:
+        A tensor with the same dtype as `x`.
+
+    Example:
+
+    >>> x = keras.ops.convert_to_tensor([1.0, 10.0, 100.0])
+    >>> keras.ops.rsqrt(x)
+    array([1.0, 0.31622776, 0.1], dtype=float32)
+    """
+    if any_symbolic_tensors((x,)):
+        return Rsqrt().symbolic_call(x)
+    x = backend.convert_to_tensor(x)
+    return backend.math.rsqrt(x)
+
+
+class Erf(Operation):
+    def compute_output_spec(self, x):
+        return KerasTensor(shape=x.shape, dtype=x.dtype)
+
+    def call(self, x):
+        return backend.math.erf(x)
+
+
+@keras_export("keras.ops.erf")
+def erf(x):
+    """Computes the error function of `x`, element-wise.
+
+    Args:
+        x: Input tensor.
+
+    Returns:
+        A tensor with the same dtype as `x`.
+
+    Example:
+
+    >>> x = np.array([-3.0, -2.0, -1.0, 0.0, 1.0])
+    >>> keras.ops.erf(x)
+    array([-0.99998 , -0.99532, -0.842701,  0.,  0.842701], dtype=float32)
+    """
+    if any_symbolic_tensors((x,)):
+        return Erf().symbolic_call(x)
+    x = backend.convert_to_tensor(x)
+    return backend.math.erf(x)
+
+
+class Erfinv(Operation):
+    def compute_output_spec(self, x):
+        return KerasTensor(shape=x.shape, dtype=x.dtype)
+
+    def call(self, x):
+        return backend.math.erfinv(x)
+
+
+@keras_export("keras.ops.erfinv")
+def erfinv(x):
+    """Computes the inverse error function of `x`, element-wise.
+
+    Args:
+        x: Input tensor.
+
+    Returns:
+        A tensor with the same dtype as `x`.
+
+    Example:
+
+    >>> x = np.array([-0.5, -0.2, -0.1, 0.0, 0.3])
+    >>> keras.ops.erfinv(x)
+    array([-0.47694, -0.17914, -0.08886,  0. ,  0.27246], dtype=float32)
+    """
+    if any_symbolic_tensors((x,)):
+        return Erfinv().symbolic_call(x)
+    x = backend.convert_to_tensor(x)
+    return backend.math.erfinv(x)
+
+
+class Logdet(Operation):
+    def __init__(self):
+        super().__init__()
+
+    def call(self, x):
+        return backend.math.logdet(x)
+
+    def compute_output_spec(self, x):
+        return KerasTensor(x.shape[:-2], dtype=x.dtype)
+
+
+@keras_export(["keras.ops.logdet"])
+def logdet(x):
+    """Computes log of the determinant of a hermitian positive definite matrix.
+
+    Args:
+        x: Input matrix. It must 2D and square.
+
+    Returns:
+        The natural log of the determinant of matrix.
+    """
+    if any_symbolic_tensors((x,)):
+        return Logdet().symbolic_call(x)
+    return backend.math.logdet(x)

SwarmUI/dlbackend/ComfyUI/venv/lib/python3.10/site-packages/keras/src/ops/nn.py

@@ -0,0 +1,2653 @@
+"""Commonly-used neural network operations not included in NumPy."""
+
+import warnings
+
+from keras.src import backend
+from keras.src.api_export import keras_export
+from keras.src.backend import KerasTensor
+from keras.src.backend import any_symbolic_tensors
+from keras.src.backend import standardize_data_format
+from keras.src.backend.common.backend_utils import (
+    compute_conv_transpose_output_shape,
+)
+from keras.src.ops import operation_utils
+from keras.src.ops.operation import Operation
+from keras.src.ops.operation_utils import reduce_shape
+
+
+class Relu(Operation):
+    def call(self, x):
+        return backend.nn.relu(x)
+
+    def compute_output_spec(self, x):
+        return KerasTensor(x.shape, dtype=x.dtype)
+
+
+@keras_export(["keras.ops.relu", "keras.ops.nn.relu"])
+def relu(x):
+    """Rectified linear unit activation function.
+
+    It is defined as `f(x) = max(0, x)`.
+
+    Args:
+        x: Input tensor.
+
+    Returns:
+        A tensor with the same shape as `x`.
+
+    Example:
+
+    >>> x1 = keras.ops.convert_to_tensor([-1.0, 0.0, 1.0, 0.2])
+    >>> keras.ops.relu(x1)
+    array([0.0, 0.0, 1.0, 0.2], dtype=float32)
+    """
+    if any_symbolic_tensors((x,)):
+        return Relu().symbolic_call(x)
+    return backend.nn.relu(x)
+
+
+class Relu6(Operation):
+    def call(self, x):
+        return backend.nn.relu6(x)
+
+    def compute_output_spec(self, x):
+        return KerasTensor(x.shape, dtype=x.dtype)
+
+
+@keras_export(["keras.ops.relu6", "keras.ops.nn.relu6"])
+def relu6(x):
+    """Rectified linear unit activation function with upper bound of 6.
+
+    It is defined as `f(x) = np.clip(x, 0, 6)`.
+
+    Args:
+        x: Input tensor.
+
+    Returns:
+        A tensor with the same shape as `x`.
+
+    Example:
+
+    >>> x = keras.ops.convert_to_tensor([-3.0, -2.0, 0.1, 0.2, 6.0, 8.0])
+    >>> keras.ops.relu6(x)
+    array([0.0, 0.0, 0.1, 0.2, 6.0, 6.0], dtype=float32)
+    """
+    if any_symbolic_tensors((x,)):
+        return Relu6().symbolic_call(x)
+    return backend.nn.relu6(x)
+
+
+class Sigmoid(Operation):
+    def call(self, x):
+        return backend.nn.sigmoid(x)
+
+    def compute_output_spec(self, x):
+        return KerasTensor(x.shape, dtype=x.dtype)
+
+
+@keras_export(["keras.ops.sigmoid", "keras.ops.nn.sigmoid"])
+def sigmoid(x):
+    """Sigmoid activation function.
+
+    It is defined as `f(x) = 1 / (1 + exp(-x))`.
+
+    Args:
+        x: Input tensor.
+
+    Returns:
+        A tensor with the same shape as `x`.
+
+    Example:
+
+    >>> x = keras.ops.convert_to_tensor([-6.0, 1.0, 0.0, 1.0, 6.0])
+    >>> keras.ops.sigmoid(x)
+    array([0.00247262, 0.7310586, 0.5, 0.7310586, 0.9975274], dtype=float32)
+
+    """
+    if any_symbolic_tensors((x,)):
+        return Sigmoid().symbolic_call(x)
+    return backend.nn.sigmoid(x)
+
+
+class Softplus(Operation):
+    def call(self, x):
+        return backend.nn.softplus(x)
+
+    def compute_output_spec(self, x):
+        return KerasTensor(x.shape, dtype=x.dtype)
+
+
+@keras_export(["keras.ops.softplus", "keras.ops.nn.softplus"])
+def softplus(x):
+    """Softplus activation function.
+
+    It is defined as `f(x) = log(exp(x) + 1)`, where `log` is the natural
+    logarithm and `exp` is the exponential function.
+
+    Args:
+        x: Input tensor.
+
+    Returns:
+        A tensor with the same shape as `x`.
+
+    Example:
+
+    >>> x = keras.ops.convert_to_tensor([-0.555, 0.0, 0.555])
+    >>> keras.ops.softplus(x)
+    array([0.45366603, 0.6931472, 1.008666], dtype=float32)
+
+    """
+    if any_symbolic_tensors((x,)):
+        return Softplus().symbolic_call(x)
+    return backend.nn.softplus(x)
+
+
+class Softsign(Operation):
+    def call(self, x):
+        return backend.nn.softsign(x)
+
+    def compute_output_spec(self, x):
+        return KerasTensor(x.shape, dtype=x.dtype)
+
+
+@keras_export(["keras.ops.softsign", "keras.ops.nn.softsign"])
+def softsign(x):
+    """Softsign activation function.
+
+    It is defined as `f(x) = x / (abs(x) + 1)`.
+
+    Args:
+        x: Input tensor.
+
+    Returns:
+        A tensor with the same shape as `x`.
+
+    Example:
+
+    >>> x = keras.ops.convert_to_tensor([-0.100, -10.0, 1.0, 0.0, 100.0])
+    >>> keras.ops.softsign(x)
+    Array([-0.09090909, -0.90909094, 0.5, 0.0, 0.990099], dtype=float32)
+
+    """
+    if any_symbolic_tensors((x,)):
+        return Softsign().symbolic_call(x)
+    return backend.nn.softsign(x)
+
+
+class SoftShrink(Operation):
+    def __init__(self, threshold=0.5):
+        super().__init__()
+        self.threshold = threshold
+
+    def call(self, x):
+        return backend.nn.soft_shrink(x, self.threshold)
+
+    def compute_output_spec(self, x):
+        return KerasTensor(x.shape, dtype=x.dtype)
+
+
+@keras_export(["keras.ops.soft_shrink", "keras.ops.nn.soft_shrink"])
+def soft_shrink(x, threshold=0.5):
+    """Soft Shrink activation function.
+
+    It is defined as
+
+    `f(x) = x - threshold` if `x > threshold`,
+    `f(x) = x + threshold` if `x < -threshold`,
+    `f(x) = 0` otherwise.
+
+    Args:
+        x: Input tensor.
+        threshold: Threshold value. Defaults to 0.5.
+
+    Returns:
+        A tensor with the same shape as `x`.
+
+    Example:
+
+    >>> x = np.array([-1.0, 0.0, 1.0])
+    >>> x_soft_shrink = keras.ops.soft_shrink(x)
+    >>> print(x_soft_shrink)
+    array([-0.5  0.   0.5], shape=(3,), dtype=float64)
+
+    """
+    if any_symbolic_tensors((x,)):
+        return SoftShrink(threshold).symbolic_call(x)
+    return backend.nn.soft_shrink(x, threshold)
+
+
+class SparsePlus(Operation):
+    def call(self, x):
+        return backend.nn.sparse_plus(x)
+
+    def compute_output_spec(self, x):
+        return KerasTensor(x.shape, dtype=x.dtype)
+
+
+@keras_export(["keras.ops.sparse_plus", "keras.ops.nn.sparse_plus"])
+def sparse_plus(x):
+    """SparsePlus activation function.
+
+    It is defined as
+
+    `f(x) = 0` for `x <= -1`.
+    `f(x) = (1/4) * (x + 1)^2` for `-1 < x < 1`.
+    `f(x) = x` for `x >= 1`.
+
+
+    Args:
+        x: Input tensor.
+
+    Returns:
+        A tensor with the same shape as `x`.
+
+    Example:
+
+    >>> x = np.array([-1.0, 0.0, 1.0])
+    >>> x_sparse_plus = keras.ops.sparse_plus(x)
+    >>> print(x_sparse_plus)
+    Array([0.   0.25 1.  ], shape=(3,), dtype=float32)
+
+    """
+    if any_symbolic_tensors((x,)):
+        return SparsePlus().symbolic_call(x)
+    return backend.nn.sparse_plus(x)
+
+
+class Silu(Operation):
+    def call(self, x):
+        return backend.nn.silu(x)
+
+    def compute_output_spec(self, x):
+        return KerasTensor(x.shape, dtype=x.dtype)
+
+
+@keras_export(
+    [
+        "keras.ops.silu",
+        "keras.ops.nn.silu",
+        "keras.ops.swish",
+        "keras.ops.nn.swish",
+    ]
+)
+def silu(x):
+    """Sigmoid Linear Unit (SiLU) activation function, also known as Swish.
+
+    The SiLU activation function is computed by the sigmoid function multiplied
+    by its input. It is defined as `f(x) = x * sigmoid(x)`.
+
+    Args:
+        x: Input tensor.
+
+    Returns:
+        A tensor with the same shape as `x`.
+
+    Example:
+
+    >>> x = keras.ops.convert_to_tensor([-6.0, 1.0, 0.0, 1.0, 6.0])
+    >>> keras.ops.sigmoid(x)
+    array([0.00247262, 0.7310586, 0.5, 0.7310586, 0.9975274], dtype=float32)
+    >>> keras.ops.silu(x)
+    array([-0.0148357, 0.7310586, 0.0, 0.7310586, 5.9851646], dtype=float32)
+
+    """
+    if any_symbolic_tensors((x,)):
+        return Silu().symbolic_call(x)
+    return backend.nn.silu(x)
+
+
+class Squareplus(Operation):
+    def __init__(self, b=4):
+        super().__init__()
+        self.b = b
+
+    def call(self, x):
+        return backend.nn.squareplus(x, self.b)
+
+    def compute_output_spec(self, x):
+        return KerasTensor(x.shape, dtype=x.dtype)
+
+
+@keras_export(["keras.ops.squareplus", "keras.ops.nn.squareplus"])
+def squareplus(x, b=4):
+    """Squareplus activation function.
+
+    The Squareplus activation function is defined as:
+
+    `f(x) = (x + sqrt(x^2 + b)) / 2`
+
+    Args:
+        x: Input tensor.
+        b: Smoothness parameter. Defaults to 4.
+
+    Returns:
+        A tensor with the same shape as `x`.
+
+    Example:
+
+    >>> x = np.array([-1.0, 0.0, 1.0])
+    >>> x_squareplus = keras.ops.squareplus(x)
+    >>> print(x_squareplus)
+    array([0.6180, 1.0000, 1.6180], dtype=float32)
+
+    """
+    if any_symbolic_tensors((x,)):
+        return Squareplus(b).symbolic_call(x)
+    return backend.nn.squareplus(x, b)
+
+
+class LogSigmoid(Operation):
+    def call(self, x):
+        return backend.nn.log_sigmoid(x)
+
+    def compute_output_spec(self, x):
+        return KerasTensor(x.shape, dtype=x.dtype)
+
+
+@keras_export(
+    [
+        "keras.ops.log_sigmoid",
+        "keras.ops.nn.log_sigmoid",
+    ]
+)
+def log_sigmoid(x):
+    """Logarithm of the sigmoid activation function.
+
+    It is defined as `f(x) = log(1 / (1 + exp(-x)))`.
+
+    Args:
+        x: Input tensor.
+
+    Returns:
+        A tensor with the same shape as `x`.
+
+    Example:
+
+    >>> x = keras.ops.convert_to_tensor([-0.541391, 0.0, 0.50, 5.0])
+    >>> keras.ops.log_sigmoid(x)
+    array([-1.0000418, -0.6931472, -0.474077, -0.00671535], dtype=float32)
+
+    """
+    if any_symbolic_tensors((x,)):
+        return LogSigmoid().symbolic_call(x)
+    return backend.nn.log_sigmoid(x)
+
+
+class LeakyRelu(Operation):
+    def __init__(self, negative_slope=0.2):
+        super().__init__()
+        self.negative_slope = negative_slope
+
+    def call(self, x):
+        return backend.nn.leaky_relu(x, self.negative_slope)
+
+    def compute_output_spec(self, x):
+        return KerasTensor(x.shape, dtype=x.dtype)
+
+
+@keras_export(["keras.ops.leaky_relu", "keras.ops.nn.leaky_relu"])
+def leaky_relu(x, negative_slope=0.2):
+    """Leaky version of a Rectified Linear Unit activation function.
+
+    It allows a small gradient when the unit is not active, it is defined as:
+
+    `f(x) = alpha * x for x < 0` or `f(x) = x for x >= 0`.
+
+    Args:
+        x: Input tensor.
+        negative_slope: Slope of the activation function at x < 0.
+            Defaults to `0.2`.
+
+    Returns:
+        A tensor with the same shape as `x`.
+
+    Example:
+
+    >>> x = np.array([-1., 0., 1.])
+    >>> x_leaky_relu = keras.ops.leaky_relu(x)
+    >>> print(x_leaky_relu)
+    array([-0.2,  0. ,  1. ], shape=(3,), dtype=float64)
+
+    """
+    if any_symbolic_tensors((x,)):
+        return LeakyRelu(negative_slope).symbolic_call(x)
+    return backend.nn.leaky_relu(x, negative_slope=negative_slope)
+
+
+class HardSigmoid(Operation):
+    def call(self, x):
+        return backend.nn.hard_sigmoid(x)
+
+    def compute_output_spec(self, x):
+        return KerasTensor(x.shape, dtype=x.dtype)
+
+
+@keras_export(
+    [
+        "keras.ops.hard_sigmoid",
+        "keras.ops.nn.hard_sigmoid",
+    ]
+)
+def hard_sigmoid(x):
+    """Hard sigmoid activation function.
+
+    It is defined as:
+
+    `0 if x < -2.5`, `1 if x > 2.5`, `(0.2 * x) + 0.5 if -2.5 <= x <= 2.5`.
+
+    Args:
+        x: Input tensor.
+
+    Returns:
+        A tensor with the same shape as `x`.
+
+    Example:
+
+    >>> x = np.array([-1., 0., 1.])
+    >>> x_hard_sigmoid = keras.ops.hard_sigmoid(x)
+    >>> print(x_hard_sigmoid)
+    array([0.3, 0.5, 0.7], shape=(3,), dtype=float64)
+
+    """
+    if any_symbolic_tensors((x,)):
+        return HardSigmoid().symbolic_call(x)
+    return backend.nn.hard_sigmoid(x)
+
+
+class HardSilu(Operation):
+    def call(self, x):
+        return backend.nn.hard_silu(x)
+
+    def compute_output_spec(self, x):
+        return KerasTensor(x.shape, dtype=x.dtype)
+
+
+@keras_export(
+    [
+        "keras.ops.hard_silu",
+        "keras.ops.nn.hard_silu",
+        "keras.ops.hard_swish",
+        "keras.ops.nn.hard_swish",
+    ]
+)
+def hard_silu(x):
+    """Hard SiLU activation function, also known as Hard Swish.
+
+    It is defined as:
+
+    - `0` if `if x < -3`
+    - `x` if `x > 3`
+    - `x * (x + 3) / 6` if `-3 <= x <= 3`
+
+    It's a faster, piecewise linear approximation of the silu activation.
+
+    Args:
+        x: Input tensor.
+
+    Returns:
+        A tensor with the same shape as `x`.
+
+    Example:
+
+    >>> x = keras.ops.convert_to_tensor([-3.0, -1.0, 0.0, 1.0, 3.0])
+    >>> keras.ops.hard_silu(x)
+    array([-0.0, -0.3333333, 0.0, 0.6666667, 3.0], shape=(5,), dtype=float32)
+
+    """
+    if any_symbolic_tensors((x,)):
+        return HardSilu().symbolic_call(x)
+    return backend.nn.hard_silu(x)
+
+
+class Elu(Operation):
+    def __init__(self, alpha=1.0):
+        super().__init__()
+        self.alpha = alpha
+
+    def call(self, x):
+        return backend.nn.elu(x, alpha=self.alpha)
+
+    def compute_output_spec(self, x):
+        return KerasTensor(x.shape, dtype=x.dtype)
+
+
+@keras_export(["keras.ops.elu", "keras.ops.nn.elu"])
+def elu(x, alpha=1.0):
+    """Exponential Linear Unit activation function.
+
+    It is defined as:
+
+    `f(x) =  alpha * (exp(x) - 1.) for x < 0`, `f(x) = x for x >= 0`.
+
+    Args:
+        x: Input tensor.
+        alpha: A scalar, slope of positive section. Defaults to `1.0`.
+
+    Returns:
+        A tensor with the same shape as `x`.
+
+    Example:
+
+    >>> x = np.array([-1., 0., 1.])
+    >>> x_elu = keras.ops.elu(x)
+    >>> print(x_elu)
+    array([-0.63212055, 0., 1.], shape=(3,), dtype=float64)
+
+    """
+    if any_symbolic_tensors((x,)):
+        return Elu(alpha).symbolic_call(x)
+    return backend.nn.elu(x, alpha=alpha)
+
+
+class Selu(Operation):
+    def call(self, x):
+        return backend.nn.selu(x)
+
+    def compute_output_spec(self, x):
+        return KerasTensor(x.shape, dtype=x.dtype)
+
+
+@keras_export(["keras.ops.selu", "keras.ops.nn.selu"])
+def selu(x):
+    """Scaled Exponential Linear Unit (SELU) activation function.
+
+    It is defined as:
+
+    `f(x) =  scale * alpha * (exp(x) - 1.) for x < 0`,
+    `f(x) = scale * x for x >= 0`.
+
+    Args:
+        x: Input tensor.
+
+    Returns:
+        A tensor with the same shape as `x`.
+
+    Example:
+
+    >>> x = np.array([-1., 0., 1.])
+    >>> x_selu = keras.ops.selu(x)
+    >>> print(x_selu)
+    array([-1.11133055, 0., 1.05070098], shape=(3,), dtype=float64)
+
+    """
+    if any_symbolic_tensors((x,)):
+        return Selu().symbolic_call(x)
+    return backend.nn.selu(x)
+
+
+class Gelu(Operation):
+    def __init__(self, approximate=True):
+        super().__init__()
+        self.approximate = approximate
+
+    def call(self, x):
+        return backend.nn.gelu(x, self.approximate)
+
+    def compute_output_spec(self, x):
+        return KerasTensor(x.shape, dtype=x.dtype)
+
+
+@keras_export(["keras.ops.gelu", "keras.ops.nn.gelu"])
+def gelu(x, approximate=True):
+    """Gaussian Error Linear Unit (GELU) activation function.
+
+    If `approximate` is `True`, it is defined as:
+    `f(x) = 0.5 * x * (1 + tanh(sqrt(2 / pi) * (x + 0.044715 * x^3)))`
+
+    Or if `approximate` is `False`, it is defined as:
+    `f(x) = x * P(X <= x) = 0.5 * x * (1 + erf(x / sqrt(2)))`,
+    where `P(X) ~ N(0, 1)`.
+
+    Args:
+        x: Input tensor.
+        approximate: Approximate version of GELU activation. Defaults to `True`.
+
+    Returns:
+        A tensor with the same shape as `x`.
+
+    Example:
+
+    >>> x = np.array([-1., 0., 1.])
+    >>> x_gelu = keras.ops.gelu(x)
+    >>> print(x_gelu)
+    array([-0.15865525, 0., 0.84134475], shape=(3,), dtype=float64)
+
+    """
+    if any_symbolic_tensors((x,)):
+        return Gelu(approximate).symbolic_call(x)
+    return backend.nn.gelu(x, approximate)
+
+
+class Celu(Operation):
+    def __init__(self, alpha=1.0):
+        super().__init__()
+        self.alpha = alpha
+
+    def call(self, x):
+        return backend.nn.celu(x, self.alpha)
+
+    def compute_output_spec(self, x):
+        return KerasTensor(x.shape, dtype=x.dtype)
+
+
+@keras_export(["keras.ops.celu", "keras.ops.nn.celu"])
+def celu(x, alpha=1.0):
+    """Continuously-differentiable exponential linear unit.
+
+    It is defined as:
+
+    `f(x) =  alpha * (exp(x / alpha) - 1) for x < 0`, `f(x) = x for x >= 0`.
+
+    Args:
+        x: Input tensor.
+        alpha: the α value for the CELU formulation. Defaults to `1.0`.
+
+    Returns:
+        A tensor with the same shape as `x`.
+
+    Example:
+
+    >>> x = np.array([-1., 0., 1.])
+    >>> x_celu = keras.ops.celu(x)
+    >>> print(x_celu)
+    array([-0.63212056, 0. , 1. ], shape=(3,), dtype=float64)
+
+    """
+    if any_symbolic_tensors((x,)):
+        return Celu(alpha).symbolic_call(x)
+    return backend.nn.celu(x, alpha)
+
+
+class Glu(Operation):
+    def __init__(self, axis=-1):
+        super().__init__()
+        self.axis = axis
+
+    def call(self, x):
+        return backend.nn.glu(x, axis=self.axis)
+
+    def compute_output_spec(self, x):
+        return KerasTensor(x.shape, dtype=x.dtype)
+
+
+@keras_export(["keras.ops.glu", "keras.ops.nn.glu"])
+def glu(x, axis=-1):
+    """Gated Linear Unit (GLU) activation function.
+
+    It is defined as:
+
+    `f(x) = a * sigmoid(b)`
+    where `x` is split into `a` and `b` along the given axis.
+
+    Args:
+        x: Input tensor.
+        axis: The axis along which to split the input tensor. Defaults to `-1`.
+
+    Returns:
+        A tensor with the same shape as half of the input.
+
+    Example:
+
+    >>> x = np.array([-1., 0., 1. , 1.])
+    >>> x_glu = keras.ops.glu(x)
+    >>> print(x_glu)
+    array([-0.73105858, 0. ], shape=(2,), dtype=float64)
+
+    """
+    if any_symbolic_tensors((x,)):
+        return Glu(axis).symbolic_call(x)
+    return backend.nn.glu(x, axis=axis)
+
+
+class TanhShrink(Operation):
+    def __init__(self):
+        super().__init__()
+
+    def call(self, x):
+        return backend.nn.tanh_shrink(x)
+
+    def compute_output_spec(self, x):
+        return KerasTensor(x.shape, dtype=x.dtype)
+
+
+@keras_export(["keras.ops.tanh_shrink", "keras.ops.nn.tanh_shrink"])
+def tanh_shrink(x):
+    """Applies the tanh shrink function element-wise.
+
+    It is defined as:
+
+    `f(x) = x - tanh(x)`.
+
+    Args:
+        x: Input tensor.
+
+    Returns:
+        Output tensor of the same shape as `x`, where each element is
+        transformed according to the tanh shrink operation.
+
+    Example:
+
+    >>> x = np.array([ -1., 0., 1.])
+    >>> x_tanh_shrink = keras.ops.tanh_shrink(x)
+    >>> print(x_tanh_shrink)
+    array([-0.23840584  0.  0.23840584], shape=(3,), dtype=float64)
+
+    """
+    if any_symbolic_tensors((x,)):
+        return TanhShrink().symbolic_call(x)
+    return backend.nn.tanh_shrink(x)
+
+
+class HardTanh(Operation):
+    def __init__(self):
+        super().__init__()
+
+    def call(self, x):
+        return backend.nn.hard_tanh(x)
+
+    def compute_output_spec(self, x):
+        return KerasTensor(x.shape, dtype=x.dtype)
+
+
+@keras_export(["keras.ops.hard_tanh", "keras.ops.nn.hard_tanh"])
+def hard_tanh(x):
+    """Applies the HardTanh function element-wise.
+
+    It is defined as:
+
+    `f(x) = -1 for x < -1`, `f(x) = x for -1 <= x <= 1`, `f(x) = 1 for x > 1`.
+
+    Args:
+        x: Input tensor.
+
+    Returns:
+        Output tensor of same shape as `x`
+        where values are clamped between -1 and 1.
+
+    Example:
+
+    >>> x = np.array([-2., -1., 0., 1., 2.])
+    >>> x_hard_tanh = keras.ops.hard_tanh(x)
+    >>> print(x_hard_tanh)
+    array([-1. -1.  0.  1.  1.], shape=(5,), dtype=float64)
+
+    """
+    if any_symbolic_tensors((x,)):
+        return HardTanh().symbolic_call(x)
+    return backend.nn.hard_tanh(x)
+
+
+class HardShrink(Operation):
+    def __init__(self, threshold=0.5):
+        super().__init__()
+        self.threshold = threshold
+
+    def call(self, x):
+        return backend.nn.hard_shrink(x, self.threshold)
+
+    def compute_output_spec(self, x):
+        return KerasTensor(x.shape, dtype=x.dtype)
+
+
+@keras_export(["keras.ops.hard_shrink", "keras.ops.nn.hard_shrink"])
+def hard_shrink(x, threshold=0.5):
+    """Hard Shrink activation function.
+
+    The Hard Shrink function is a thresholding operation defined as:
+
+    `f(x) = x` if `|x| > threshold`,
+    `f(x) = 0` otherwise.
+
+    Args:
+        x: Input tensor.
+        threshold: Threshold value. Defaults to 0.5.
+
+    Returns:
+        A tensor with the same shape as `x`.
+
+    Example:
+
+    >>> x = np.array([-0.5, 0., 1.])
+    >>> x_hard_shrink = keras.ops.hard_shrink(x)
+    >>> print(x_hard_shrink)
+    array([0. 0. 1.], shape=(3,), dtype=float64)
+
+    """
+    if any_symbolic_tensors((x,)):
+        return HardShrink(threshold).symbolic_call(x)
+    return backend.nn.hard_shrink(x, threshold)
+
+
+class Threshold(Operation):
+    def __init__(self, threshold_value, value):
+        super().__init__()
+        self.threshold_value = threshold_value
+        self.value = value
+
+    def call(self, x):
+        return backend.nn.threshold(x, self.threshold_value, self.value)
+
+    def compute_output_spec(self, x):
+        return KerasTensor(x.shape, dtype=x.dtype)
+
+
+@keras_export(["keras.ops.threshold", "keras.ops.nn.threshold"])
+def threshold(x, threshold, default_value):
+    """Threshold activation function.
+
+    The function thresholds the input `x` as follows:
+    `f(x) = x` if `x > threshold`,
+    `f(x) = default_value` otherwise.
+
+    Args:
+        x: Input tensor.
+        threshold: The value that decides when to retain or replace x.
+        default_value: Value to assign when `x <= threshold`.
+
+    Returns:
+        A tensor with the same shape as `x`.
+
+    Example:
+
+    >>> x = np.array([-1.0, 0.0, 1.0, 2.0])
+    >>> x_threshold = keras.ops.threshold(x, 1, 0)
+    >>> print(x_threshold)
+    array([0., 0., 0., 2.], shape=(4,), dtype=float64)
+
+    """
+    if any_symbolic_tensors((x,)):
+        return Threshold(threshold, default_value).symbolic_call(x)
+    return backend.nn.threshold(x, threshold, default_value)
+
+
+class Softmax(Operation):
+    def __init__(self, axis=-1):
+        super().__init__()
+        self.axis = axis
+
+    def call(self, x):
+        return backend.nn.softmax(x, axis=self.axis)
+
+    def compute_output_spec(self, x):
+        return KerasTensor(x.shape, dtype=x.dtype)
+
+
+@keras_export(["keras.ops.softmax", "keras.ops.nn.softmax"])
+def softmax(x, axis=-1):
+    """Softmax activation function.
+
+    The elements of the output vector lie within the range `(0, 1)`, and their
+    total sum is exactly 1 (excluding the floating point rounding error).
+
+    Each vector is processed independently. The `axis` argument specifies the
+    axis along which the function is applied within the input.
+
+    It is defined as:
+    `f(x) = exp(x) / sum(exp(x))`
+
+    Args:
+        x: Input tensor.
+        axis: Integer, axis along which the softmax is applied.
+
+    Returns:
+        A tensor with the same shape as `x`.
+
+    Example:
+
+    >>> x = np.array([-1., 0., 1.])
+    >>> x_softmax = keras.ops.softmax(x)
+    >>> print(x_softmax)
+    array([0.09003057, 0.24472847, 0.66524096], shape=(3,), dtype=float64)
+
+    """
+    # Don't use `backend.shape` since TensorFlow returns
+    # symbolic tensors for unknown shape which can trigger
+    # an error in TensorFlow graph execution.
+    if isinstance(axis, int) and x.shape[axis] == 1:
+        warnings.warn(
+            f"You are using a softmax over axis {axis} "
+            f"of a tensor of shape {x.shape}. This axis "
+            "has size 1. The softmax operation will always return "
+            "the value 1, which is likely not what you intended. "
+            "Did you mean to use a sigmoid instead?"
+        )
+    if any_symbolic_tensors((x,)):
+        return Softmax(axis).symbolic_call(x)
+    if isinstance(axis, tuple):
+        axis_to_keep = [v for v in range(len(x.shape)) if v not in axis]
+
+        x_transposed = backend.numpy.transpose(x, axes=(*axis_to_keep, *axis))
+        x_reshaped = backend.numpy.reshape(
+            x_transposed, (*[x.shape[v] for v in axis_to_keep], -1)
+        )
+
+        x = backend.nn.softmax(x_reshaped, axis=-1)
+
+        x = backend.numpy.reshape(x, x_transposed.shape)
+        x = backend.numpy.transpose(
+            x, axes=list(backend.numpy.argsort([*axis_to_keep, *axis]))
+        )
+        return x
+    else:
+        return backend.nn.softmax(x, axis=axis)
+
+
+class LogSoftmax(Operation):
+    def __init__(self, axis=-1):
+        super().__init__()
+        self.axis = axis
+
+    def call(self, x):
+        return backend.nn.log_softmax(x, axis=self.axis)
+
+    def compute_output_spec(self, x):
+        return KerasTensor(x.shape, dtype=x.dtype)
+
+
+@keras_export(
+    [
+        "keras.ops.log_softmax",
+        "keras.ops.nn.log_softmax",
+    ]
+)
+def log_softmax(x, axis=-1):
+    """Log-softmax activation function.
+
+    It is defined as:
+    `f(x) = x - max(x) - log(sum(exp(x - max(x))))`
+
+    Args:
+        x: Input tensor.
+        axis: Integer, axis along which the log-softmax is applied.
+            Defaults to `-1`.
+
+    Returns:
+        A tensor with the same shape as `x`.
+
+    Example:
+
+    >>> x = np.array([-1., 0., 1.])
+    >>> x_log_softmax = keras.ops.log_softmax(x)
+    >>> print(x_log_softmax)
+    array([-2.40760596, -1.40760596, -0.40760596], shape=(3,), dtype=float64)
+
+    """
+    if any_symbolic_tensors((x,)):
+        return LogSoftmax(axis).symbolic_call(x)
+    if isinstance(axis, tuple):
+        axis_to_keep = [v for v in range(len(x.shape)) if v not in axis]
+
+        x_transposed = backend.numpy.transpose(x, axes=(*axis_to_keep, *axis))
+        x_reshaped = backend.numpy.reshape(
+            x_transposed, (*[x.shape[v] for v in axis_to_keep], -1)
+        )
+
+        x = backend.nn.log_softmax(x_reshaped, axis=-1)
+
+        x = backend.numpy.reshape(x, x_transposed.shape)
+        x = backend.numpy.transpose(
+            x, axes=list(backend.numpy.argsort([*axis_to_keep, *axis]))
+        )
+        return x
+    else:
+        return backend.nn.log_softmax(x, axis=axis)
+
+
+class Sparsemax(Operation):
+    def __init__(self, axis=-1):
+        super().__init__()
+        self.axis = axis
+
+    def call(self, x):
+        return backend.nn.sparsemax(x, axis=self.axis)
+
+    def compute_output_spec(self, x):
+        return KerasTensor(x.shape, dtype=x.dtype)
+
+
+@keras_export(["keras.ops.sparsemax", "keras.ops.nn.sparsemax"])
+def sparsemax(x, axis=-1):
+    """Sparsemax activation function.
+
+    For each batch `i`, and class `j`,
+    sparsemax activation function is defined as:
+
+    `sparsemax(x)[i, j] = max(x[i, j] - τ(x[i, :]), 0).`
+
+    Args:
+        x: Input tensor.
+        axis: `int`, axis along which the sparsemax operation is applied.
+
+    Returns:
+        A tensor, output of sparsemax transformation. Has the same type and
+        shape as `x`.
+
+    Example:
+
+    >>> x = np.array([-1., 0., 1.])
+    >>> x_sparsemax = keras.ops.sparsemax(x)
+    >>> print(x_sparsemax)
+    array([0., 0., 1.], shape=(3,), dtype=float64)
+
+    """
+    if any_symbolic_tensors((x,)):
+        return Sparsemax(axis).symbolic_call(x)
+    return backend.nn.sparsemax(x, axis=axis)
+
+
+class MaxPool(Operation):
+    def __init__(
+        self,
+        pool_size,
+        strides=None,
+        padding="valid",
+        data_format=None,
+    ):
+        super().__init__()
+        self.pool_size = pool_size
+        self.strides = strides
+        self.padding = padding.lower()
+        self.data_format = data_format
+
+    def call(self, inputs):
+        return backend.nn.max_pool(
+            inputs,
+            self.pool_size,
+            self.strides,
+            self.padding,
+            self.data_format,
+        )
+
+    def compute_output_spec(self, inputs):
+        output_shape = operation_utils.compute_pooling_output_shape(
+            inputs.shape,
+            self.pool_size,
+            self.strides,
+            self.padding,
+            self.data_format,
+        )
+        return KerasTensor(output_shape, dtype=inputs.dtype)
+
+
+@keras_export(["keras.ops.max_pool", "keras.ops.nn.max_pool"])
+def max_pool(
+    inputs,
+    pool_size,
+    strides=None,
+    padding="valid",
+    data_format=None,
+):
+    """Max pooling operation.
+
+    Args:
+        inputs: Tensor of rank N+2. `inputs` has shape
+            `(batch_size,) + inputs_spatial_shape + (num_channels,)` if
+            `data_format="channels_last"`, or
+            `(batch_size, num_channels) + inputs_spatial_shape` if
+            `data_format="channels_first"`. Pooling happens over the spatial
+            dimensions only.
+        pool_size: int or tuple/list of integers of size
+            `len(inputs_spatial_shape)`, specifying the size of the pooling
+            window for each spatial dimension of the input tensor. If
+            `pool_size` is int, then every spatial dimension shares the same
+            `pool_size`.
+        strides: int or tuple/list of integers of size
+            `len(inputs_spatial_shape)`. The stride of the sliding window for
+            each spatial dimension of the input tensor. If `strides` is int,
+            then every spatial dimension shares the same `strides`.
+        padding: string, either `"valid"` or `"same"`. `"valid"` means no
+            padding is applied, and `"same"` results in padding evenly to the
+            left/right or up/down of the input such that output has the
+            same height/width dimension as the input when `strides=1`.
+        data_format: A string, either `"channels_last"` or `"channels_first"`.
+            `data_format` determines the ordering of the dimensions in the
+            inputs. If `data_format="channels_last"`, `inputs` is of shape
+            `(batch_size, ..., channels)` while if
+            `data_format="channels_first"`, `inputs` is of shape
+            `(batch_size, channels, ...)`.
+
+    Returns:
+        A tensor of rank N+2, the result of the max pooling operation.
+    """
+    data_format = standardize_data_format(data_format)
+    padding = padding.lower()
+    if any_symbolic_tensors((inputs,)):
+        return MaxPool(
+            pool_size,
+            strides,
+            padding,
+            data_format,
+        ).symbolic_call(inputs)
+    return backend.nn.max_pool(inputs, pool_size, strides, padding, data_format)
+
+
+class AveragePool(Operation):
+    def __init__(
+        self,
+        pool_size,
+        strides=None,
+        padding="valid",
+        data_format=None,
+    ):
+        super().__init__()
+        self.pool_size = pool_size
+        self.strides = strides
+        self.padding = padding.lower()
+        self.data_format = data_format
+
+    def call(self, inputs):
+        return backend.nn.average_pool(
+            inputs,
+            self.pool_size,
+            self.strides,
+            self.padding,
+            self.data_format,
+        )
+
+    def compute_output_spec(self, inputs):
+        output_shape = operation_utils.compute_pooling_output_shape(
+            inputs.shape,
+            self.pool_size,
+            self.strides,
+            self.padding,
+            self.data_format,
+        )
+        return KerasTensor(output_shape, dtype=inputs.dtype)
+
+
+@keras_export(
+    [
+        "keras.ops.average_pool",
+        "keras.ops.nn.average_pool",
+    ]
+)
+def average_pool(
+    inputs,
+    pool_size,
+    strides=None,
+    padding="valid",
+    data_format=None,
+):
+    """Average pooling operation.
+
+    Args:
+        inputs: Tensor of rank N+2. `inputs` has shape
+            `(batch_size,) + inputs_spatial_shape + (num_channels,)` if
+            `data_format="channels_last"`, or
+            `(batch_size, num_channels) + inputs_spatial_shape` if
+            `data_format="channels_first"`. Pooling happens over the spatial
+            dimensions only.
+        pool_size: int or tuple/list of integers of size
+            `len(inputs_spatial_shape)`, specifying the size of the pooling
+            window for each spatial dimension of the input tensor. If
+            `pool_size` is int, then every spatial dimension shares the same
+            `pool_size`.
+        strides: int or tuple/list of integers of size
+            `len(inputs_spatial_shape)`. The stride of the sliding window for
+            each spatial dimension of the input tensor. If `strides` is int,
+            then every spatial dimension shares the same `strides`.
+        padding: string, either `"valid"` or `"same"`. `"valid"` means no
+            padding is applied, and `"same"` results in padding evenly to the
+            left/right or up/down of the input such that output has the
+            same height/width dimension as the input when `strides=1`.
+        data_format: A string, either `"channels_last"` or `"channels_first"`.
+            `data_format` determines the ordering of the dimensions in the
+            inputs. If `data_format="channels_last"`, `inputs` is of shape
+            `(batch_size, ..., channels)` while if
+            `data_format="channels_first"`, `inputs` is of shape
+            `(batch_size, channels, ...)`.
+
+    Returns:
+        A tensor of rank N+2, the result of the average pooling operation.
+    """
+    data_format = standardize_data_format(data_format)
+    padding = padding.lower()
+    if any_symbolic_tensors((inputs,)):
+        return AveragePool(
+            pool_size,
+            strides,
+            padding,
+            data_format,
+        ).symbolic_call(inputs)
+    return backend.nn.average_pool(
+        inputs, pool_size, strides, padding, data_format
+    )
+
+
+class Conv(Operation):
+    def __init__(
+        self,
+        strides=1,
+        padding="valid",
+        data_format=None,
+        dilation_rate=1,
+    ):
+        super().__init__()
+        self.strides = strides
+        self.padding = padding.lower()
+        self.data_format = data_format
+        self.dilation_rate = dilation_rate
+
+    def call(self, inputs, kernel):
+        return backend.nn.conv(
+            inputs,
+            kernel,
+            strides=self.strides,
+            padding=self.padding,
+            data_format=self.data_format,
+            dilation_rate=self.dilation_rate,
+        )
+
+    def compute_output_spec(self, inputs, kernel):
+        output_shape = operation_utils.compute_conv_output_shape(
+            inputs.shape,
+            kernel.shape[-1],
+            kernel.shape[:-2],
+            self.strides,
+            self.padding,
+            self.data_format,
+            self.dilation_rate,
+        )
+        return KerasTensor(output_shape, dtype=inputs.dtype)
+
+
+@keras_export(["keras.ops.conv", "keras.ops.nn.conv"])
+def conv(
+    inputs,
+    kernel,
+    strides=1,
+    padding="valid",
+    data_format=None,
+    dilation_rate=1,
+):
+    """General N-D convolution.
+
+    This ops supports 1D, 2D and 3D convolution.
+
+    Args:
+        inputs: Tensor of rank N+2. `inputs` has shape
+            `(batch_size,) + inputs_spatial_shape + (num_channels,)` if
+            `data_format="channels_last"`, or
+            `(batch_size, num_channels) + inputs_spatial_shape` if
+            `data_format="channels_first"`.
+        kernel: Tensor of rank N+2. `kernel` has shape
+            `(kernel_spatial_shape, num_input_channels, num_output_channels)`.
+            `num_input_channels` should match the number of channels in
+            `inputs`.
+        strides: int or int tuple/list of `len(inputs_spatial_shape)`,
+            specifying the strides of the convolution along each spatial
+            dimension. If `strides` is int, then every spatial dimension shares
+            the same `strides`.
+        padding: string, either `"valid"` or `"same"`. `"valid"` means no
+            padding is applied, and `"same"` results in padding evenly to the
+            left/right or up/down of the input such that output has the
+            same height/width dimension as the input when `strides=1`.
+        data_format: A string, either `"channels_last"` or `"channels_first"`.
+            `data_format` determines the ordering of the dimensions in the
+            inputs. If `data_format="channels_last"`, `inputs` is of shape
+            `(batch_size, ..., channels)` while if
+            `data_format="channels_first"`, `inputs` is of shape
+            `(batch_size, channels, ...)`.
+        dilation_rate: int or int tuple/list of `len(inputs_spatial_shape)`,
+            specifying the dilation rate to use for dilated convolution. If
+            `dilation_rate` is int, then every spatial dimension shares
+            the same `dilation_rate`.
+
+    Returns:
+        A tensor of rank N+2, the result of the conv operation.
+    """
+    data_format = standardize_data_format(data_format)
+    padding = padding.lower()
+    if any_symbolic_tensors((inputs,)):
+        return Conv(strides, padding, data_format, dilation_rate).symbolic_call(
+            inputs, kernel
+        )
+    return backend.nn.conv(
+        inputs, kernel, strides, padding, data_format, dilation_rate
+    )
+
+
+class DepthwiseConv(Operation):
+    def __init__(
+        self,
+        strides=1,
+        padding="valid",
+        data_format=None,
+        dilation_rate=1,
+    ):
+        super().__init__()
+        self.strides = strides
+        self.padding = padding.lower()
+        self.data_format = data_format
+        self.dilation_rate = dilation_rate
+
+    def call(self, inputs, kernel):
+        return backend.nn.depthwise_conv(
+            inputs,
+            kernel,
+            self.strides,
+            self.padding,
+            self.data_format,
+            self.dilation_rate,
+        )
+
+    def compute_output_spec(self, inputs, kernel):
+        output_shape = operation_utils.compute_conv_output_shape(
+            inputs.shape,
+            kernel.shape[-1] * kernel.shape[-2],
+            kernel.shape[:-2],
+            self.strides,
+            self.padding,
+            self.data_format,
+            self.dilation_rate,
+        )
+        return KerasTensor(output_shape, dtype=inputs.dtype)
+
+
+@keras_export(
+    [
+        "keras.ops.depthwise_conv",
+        "keras.ops.nn.depthwise_conv",
+    ]
+)
+def depthwise_conv(
+    inputs,
+    kernel,
+    strides=1,
+    padding="valid",
+    data_format=None,
+    dilation_rate=1,
+):
+    """General N-D depthwise convolution.
+
+    This ops supports 1D and 2D depthwise convolution.
+
+    Args:
+        inputs: Tensor of rank N+2. `inputs` has shape
+            `(batch_size,) + inputs_spatial_shape + (num_channels,)` if
+            `data_format="channels_last"`, or
+            `(batch_size, num_channels) + inputs_spatial_shape` if
+            `data_format="channels_first"`.
+        kernel: Tensor of rank N+2. `kernel` has shape
+            [kernel_spatial_shape, num_input_channels, num_channels_multiplier],
+            `num_input_channels` should match the number of channels in
+            `inputs`.
+        strides: int or int tuple/list of `len(inputs_spatial_shape)`,
+            specifying the strides of the convolution along each spatial
+            dimension. If `strides` is int, then every spatial dimension shares
+            the same `strides`.
+        padding: string, either `"valid"` or `"same"`. `"valid"` means no
+            padding is applied, and `"same"` results in padding evenly to the
+            left/right or up/down of the input such that output has the
+            same height/width dimension as the input when `strides=1`.
+        data_format: A string, either `"channels_last"` or `"channels_first"`.
+            `data_format` determines the ordering of the dimensions in the
+            inputs. If `data_format="channels_last"`, `inputs` is of shape
+            `(batch_size, ..., channels)` while if
+            `data_format="channels_first"`, `inputs` is of shape
+            `(batch_size, channels, ...)`.
+        dilation_rate: int or int tuple/list of `len(inputs_spatial_shape)`,
+            specifying the dilation rate to use for dilated convolution. If
+            `dilation_rate` is int, then every spatial dimension shares
+            the same `dilation_rate`.
+
+    Returns:
+        A tensor of rank N+2, the result of the depthwise conv operation.
+    """
+    data_format = standardize_data_format(data_format)
+    padding = padding.lower()
+    if any_symbolic_tensors((inputs,)):
+        return DepthwiseConv(
+            strides, padding, data_format, dilation_rate
+        ).symbolic_call(inputs, kernel)
+    return backend.nn.depthwise_conv(
+        inputs,
+        kernel,
+        strides,
+        padding,
+        data_format,
+        dilation_rate,
+    )
+
+
+class SeparableConv(Operation):
+    def __init__(
+        self,
+        strides=1,
+        padding="valid",
+        data_format=None,
+        dilation_rate=1,
+    ):
+        super().__init__()
+        self.strides = strides
+        self.padding = padding.lower()
+        self.data_format = data_format
+        self.dilation_rate = dilation_rate
+
+    def call(self, inputs, depthwise_kernel, pointwise_kernel):
+        return backend.nn.separable_conv(
+            inputs,
+            depthwise_kernel,
+            pointwise_kernel,
+            self.strides,
+            self.padding,
+            self.data_format,
+            self.dilation_rate,
+        )
+
+    def compute_output_spec(self, inputs, depthwise_kernel, pointwise_kernel):
+        output_shape = list(
+            depthwise_conv(
+                inputs,
+                depthwise_kernel,
+                self.strides,
+                self.padding,
+                self.data_format,
+                self.dilation_rate,
+            ).shape
+        )
+        if self.data_format == "channels_last":
+            output_shape[-1] = pointwise_kernel.shape[-1]
+        else:
+            output_shape[1] = pointwise_kernel.shape[-1]
+        return KerasTensor(output_shape, dtype=inputs.dtype)
+
+
+@keras_export(
+    [
+        "keras.ops.separable_conv",
+        "keras.ops.nn.separable_conv",
+    ]
+)
+def separable_conv(
+    inputs,
+    depthwise_kernel,
+    pointwise_kernel,
+    strides=1,
+    padding="valid",
+    data_format=None,
+    dilation_rate=1,
+):
+    """General N-D separable convolution.
+
+    This ops supports 1D and 2D separable convolution. `separable_conv` is
+    a depthwise conv followed by a pointwise conv.
+
+    Args:
+        inputs: Tensor of rank N+2. `inputs` has shape
+            `(batch_size,) + inputs_spatial_shape + (num_channels,)` if
+            `data_format="channels_last"`, or
+            `(batch_size, num_channels) + inputs_spatial_shape` if
+            `data_format="channels_first"`.
+        depthwise_kernel: Tensor of rank N+2. `depthwise_kernel` has shape
+            [kernel_spatial_shape, num_input_channels, num_channels_multiplier],
+            `num_input_channels` should match the number of channels in
+            `inputs`.
+        pointwise_kernel: Tensor of rank N+2. `pointwise_kernel` has shape
+            `(*ones_like(kernel_spatial_shape),
+            num_input_channels * num_channels_multiplier, num_output_channels)`.
+        strides: int or int tuple/list of `len(inputs_spatial_shape)`,
+            specifying the strides of the convolution along each spatial
+            dimension. If `strides` is int, then every spatial dimension shares
+            the same `strides`.
+        padding: string, either `"valid"` or `"same"`. `"valid"` means no
+            padding is applied, and `"same"` results in padding evenly to the
+            left/right or up/down of the input such that output has the
+            same height/width dimension as the input when `strides=1`.
+        data_format: A string, either `"channels_last"` or `"channels_first"`.
+            `data_format` determines the ordering of the dimensions in the
+            inputs. If `data_format="channels_last"`, `inputs` is of shape
+            `(batch_size, ..., channels)` while if
+            `data_format="channels_first"`, `inputs` is of shape
+            `(batch_size, channels, ...)`.
+        dilation_rate: int or int tuple/list of `len(inputs_spatial_shape)`,
+            specifying the dilation rate to use for dilated convolution. If
+            `dilation_rate` is int, then every spatial dimension shares
+            the same `dilation_rate`.
+
+    Returns:
+        A tensor of rank N+2, the result of the depthwise conv operation.
+    """
+    data_format = standardize_data_format(data_format)
+    padding = padding.lower()
+    if any_symbolic_tensors((inputs,)):
+        return SeparableConv(
+            strides,
+            padding,
+            data_format,
+            dilation_rate,
+        ).symbolic_call(inputs, depthwise_kernel, pointwise_kernel)
+    return backend.nn.separable_conv(
+        inputs,
+        depthwise_kernel,
+        pointwise_kernel,
+        strides,
+        padding,
+        data_format,
+        dilation_rate,
+    )
+
+
+class ConvTranspose(Operation):
+    def __init__(
+        self,
+        strides,
+        padding="valid",
+        output_padding=None,
+        data_format=None,
+        dilation_rate=1,
+    ):
+        super().__init__()
+        self.strides = strides
+        self.output_padding = output_padding
+        self.padding = padding.lower()
+        self.data_format = data_format
+        self.dilation_rate = dilation_rate
+
+    def call(
+        self,
+        inputs,
+        kernel,
+    ):
+        return backend.nn.conv_transpose(
+            inputs,
+            kernel,
+            self.strides,
+            self.output_padding,
+            self.padding,
+            self.data_format,
+            self.dilation_rate,
+        )
+
+    def compute_output_spec(self, inputs, kernel):
+        kernel_size = kernel.shape[:-2]
+        filters = kernel.shape[-2]
+        output_shape = compute_conv_transpose_output_shape(
+            inputs.shape,
+            kernel_size,
+            filters,
+            self.strides,
+            self.padding,
+            self.output_padding,
+            self.data_format,
+            self.dilation_rate,
+        )
+        return KerasTensor(output_shape, dtype=inputs.dtype)
+
+
+@keras_export(
+    [
+        "keras.ops.conv_transpose",
+        "keras.ops.nn.conv_transpose",
+    ]
+)
+def conv_transpose(
+    inputs,
+    kernel,
+    strides,
+    padding="valid",
+    output_padding=None,
+    data_format=None,
+    dilation_rate=1,
+):
+    """General N-D convolution transpose.
+
+    Also known as de-convolution. This ops supports 1D, 2D and 3D convolution.
+
+    Args:
+        inputs: Tensor of rank N+2. `inputs` has shape
+            `(batch_size,) + inputs_spatial_shape + (num_channels,)` if
+            `data_format="channels_last"`, or
+            `(batch_size, num_channels) + inputs_spatial_shape` if
+            `data_format="channels_first"`.
+        kernel: Tensor of rank N+2. `kernel` has shape
+            [kernel_spatial_shape, num_output_channels, num_input_channels],
+            `num_input_channels` should match the number of channels in
+            `inputs`.
+        strides: int or int tuple/list of `len(inputs_spatial_shape)`,
+            specifying the strides of the convolution along each spatial
+            dimension. If `strides` is int, then every spatial dimension shares
+            the same `strides`.
+        padding: string, either `"valid"` or `"same"`. `"valid"` means no
+            padding is applied, and `"same"` results in padding evenly to the
+            left/right or up/down of the input such that output has the
+            same height/width dimension as the input when `strides=1`.
+        output_padding: int or int tuple/list of `len(inputs_spatial_shape)`,
+            specifying the amount of padding along the height and width of
+            the output tensor. Can be a single integer to specify the same
+            value for all spatial dimensions. The amount of output padding
+            along a given dimension must be lower than the stride along that
+            same dimension. If set to `None` (default), the output shape is
+            inferred.
+        data_format: A string, either `"channels_last"` or `"channels_first"`.
+            `data_format` determines the ordering of the dimensions in the
+            inputs. If `data_format="channels_last"`, `inputs` is of shape
+            `(batch_size, ..., channels)` while if
+            `data_format="channels_first"`, `inputs` is of shape
+            `(batch_size, channels, ...)`.
+        dilation_rate: int or int tuple/list of `len(inputs_spatial_shape)`,
+            specifying the dilation rate to use for dilated convolution. If
+            `dilation_rate` is int, then every spatial dimension shares
+            the same `dilation_rate`.
+
+    Returns:
+        A tensor of rank N+2, the result of the conv operation.
+    """
+    data_format = standardize_data_format(data_format)
+    padding = padding.lower()
+    if any_symbolic_tensors((inputs,)):
+        return ConvTranspose(
+            strides, padding, output_padding, data_format, dilation_rate
+        ).symbolic_call(inputs, kernel)
+    return backend.nn.conv_transpose(
+        inputs,
+        kernel,
+        strides,
+        padding,
+        output_padding,
+        data_format,
+        dilation_rate,
+    )
+
+
+class OneHot(Operation):
+    def __init__(self, num_classes, axis=-1, dtype=None, sparse=False):
+        super().__init__()
+        self.num_classes = num_classes
+        self.axis = axis
+        self.dtype = dtype or backend.floatx()
+        self.sparse = sparse
+
+    def call(self, x):
+        return backend.nn.one_hot(
+            x,
+            self.num_classes,
+            axis=self.axis,
+            dtype=self.dtype,
+            sparse=self.sparse,
+        )
+
+    def compute_output_spec(self, x):
+        x_shape = list(getattr(x, "shape", []))
+        if self.axis == -1:
+            x_shape.append(self.num_classes)
+        elif self.axis >= 0 and self.axis < len(x_shape):
+            x_shape.insert(self.axis, self.num_classes)
+        else:
+            raise ValueError(
+                f"axis must be -1 or between [0, {len(x.shape)}), but "
+                f"received {self.axis}."
+            )
+        return KerasTensor(x_shape, dtype=self.dtype, sparse=self.sparse)
+
+
+@keras_export(["keras.ops.one_hot", "keras.ops.nn.one_hot"])
+def one_hot(x, num_classes, axis=-1, dtype=None, sparse=False):
+    """Converts integer tensor `x` into a one-hot tensor.
+
+    The one-hot encoding is a representation where each integer value is
+    converted into a binary vector with a length equal to `num_classes`,
+    and the index corresponding to the integer value is marked as 1, while
+    all other indices are marked as 0.
+
+    Args:
+        x: Integer tensor to be encoded. The shape can be
+            arbitrary, but the dtype should be integer.
+        num_classes: Number of classes for the one-hot encoding.
+        axis: Axis along which the encoding is performed.
+            `-1` represents the last axis. Defaults to `-1`.
+        dtype: (Optional) Data type of the output tensor. If not
+            provided, it defaults to the default data type of the backend.
+        sparse: Whether to return a sparse tensor; for backends that support
+            sparse tensors.
+
+    Returns:
+        Integer tensor: One-hot encoded tensor with the same shape as `x`
+        except for the specified `axis` dimension, which will have
+        a length of `num_classes`. The dtype of the output tensor
+        is determined by `dtype` or the default data type of the backend.
+
+    Example:
+
+    >>> x = keras.ops.convert_to_tensor([1, 3, 2, 0])
+    >>> one_hot(x, num_classes=4)
+    array([[0. 1. 0. 0.]
+           [0. 0. 0. 1.]
+           [0. 0. 1. 0.]
+           [1. 0. 0. 0.]], shape=(4, 4), dtype=float32)
+    """
+    if any_symbolic_tensors((x,)):
+        return OneHot(
+            num_classes, axis=axis, dtype=dtype, sparse=sparse
+        ).symbolic_call(x)
+    return backend.nn.one_hot(
+        x,
+        num_classes,
+        axis=axis,
+        dtype=dtype or backend.floatx(),
+        sparse=sparse,
+    )
+
+
+class BinaryCrossentropy(Operation):
+    def __init__(self, from_logits=False):
+        super().__init__()
+        self.from_logits = from_logits
+
+    def call(self, target, output):
+        return backend.nn.binary_crossentropy(
+            target, output, from_logits=self.from_logits
+        )
+
+    def compute_output_spec(self, target, output):
+        if target.shape != output.shape:
+            raise ValueError(
+                "Arguments `target` and `output` must have the same shape. "
+                "Received: "
+                f"target.shape={target.shape}, output.shape={output.shape}"
+            )
+        return KerasTensor(output.shape, dtype=output.dtype)
+
+
+@keras_export(
+    [
+        "keras.ops.binary_crossentropy",
+        "keras.ops.nn.binary_crossentropy",
+    ]
+)
+def binary_crossentropy(target, output, from_logits=False):
+    """Computes binary cross-entropy loss between target and output tensor.
+
+    The binary cross-entropy loss is commonly used in binary
+    classification tasks where each input sample belongs to one
+    of the two classes. It measures the dissimilarity between the
+    target and output probabilities or logits.
+
+    Args:
+        target: The target tensor representing the true binary labels.
+            Its shape should match the shape of the `output` tensor.
+        output: The output tensor representing the predicted probabilities
+            or logits. Its shape should match the shape of the
+            `target` tensor.
+        from_logits: (optional) Whether `output` is a tensor of logits or
+            probabilities.
+            Set it to `True` if `output` represents logits; otherwise,
+            set it to `False` if `output` represents probabilities.
+            Defaults to `False`.
+
+    Returns:
+        Integer tensor: The computed binary cross-entropy loss between
+        `target` and `output`.
+
+    Example:
+
+    >>> target = keras.ops.convert_to_tensor([0, 1, 1, 0])
+    >>> output = keras.ops.convert_to_tensor([0.1, 0.9, 0.8, 0.2])
+    >>> binary_crossentropy(target, output)
+    array([0.10536054 0.10536054 0.22314355 0.22314355],
+          shape=(4,), dtype=float32)
+    """
+    if any_symbolic_tensors((target, output)):
+        return BinaryCrossentropy(from_logits=from_logits).symbolic_call(
+            target, output
+        )
+    return backend.nn.binary_crossentropy(
+        target, output, from_logits=from_logits
+    )
+
+
+class CategoricalCrossentropy(Operation):
+    def __init__(self, from_logits=False, axis=-1):
+        super().__init__()
+        self.from_logits = from_logits
+        self.axis = axis
+
+    def call(self, target, output):
+        return backend.nn.categorical_crossentropy(
+            target, output, from_logits=self.from_logits, axis=self.axis
+        )
+
+    def compute_output_spec(self, target, output):
+        if target.shape != output.shape:
+            raise ValueError(
+                "Arguments `target` and `output` must have the same shape. "
+                "Received: "
+                f"target.shape={target.shape}, output.shape={output.shape}"
+            )
+        if len(target.shape) < 1:
+            raise ValueError(
+                "Arguments `target` and `output` must be at least rank 1. "
+                "Received: "
+                f"target.shape={target.shape}, output.shape={output.shape}"
+            )
+        return KerasTensor(output.shape[:-1], dtype=output.dtype)
+
+
+@keras_export(
+    [
+        "keras.ops.categorical_crossentropy",
+        "keras.ops.nn.categorical_crossentropy",
+    ]
+)
+def categorical_crossentropy(target, output, from_logits=False, axis=-1):
+    """Computes categorical cross-entropy loss between target and output tensor.
+
+    The categorical cross-entropy loss is commonly used in multi-class
+    classification tasks where each input sample can belong to one of
+    multiple classes. It measures the dissimilarity
+    between the target and output probabilities or logits.
+
+    Args:
+        target: The target tensor representing the true categorical labels.
+            Its shape should match the shape of the `output` tensor
+            except for the last dimension.
+        output: The output tensor representing the predicted probabilities
+            or logits. Its shape should match the shape of the `target`
+            tensor except for the last dimension.
+        from_logits: (optional) Whether `output` is a tensor of logits or
+            probabilities.
+            Set it to `True` if `output` represents logits; otherwise,
+            set it to `False` if `output` represents probabilities.
+            Defaults to `False`.
+        axis: (optional) The axis along which the categorical cross-entropy
+            is computed.
+            Defaults to `-1`, which corresponds to the last dimension of
+            the tensors.
+
+    Returns:
+        Integer tensor: The computed categorical cross-entropy loss between
+        `target` and `output`.
+
+    Example:
+
+    >>> target = keras.ops.convert_to_tensor(
+    ... [[1, 0, 0],
+    ...  [0, 1, 0],
+    ...  [0, 0, 1]])
+    >>> output = keras.ops.convert_to_tensor(
+    ... [[0.9, 0.05, 0.05],
+    ...  [0.1, 0.8, 0.1],
+    ...  [0.2, 0.3, 0.5]])
+    >>> categorical_crossentropy(target, output)
+    array([0.10536054 0.22314355 0.6931472 ], shape=(3,), dtype=float32)
+    """
+    if any_symbolic_tensors((target, output)):
+        return CategoricalCrossentropy(
+            from_logits=from_logits, axis=axis
+        ).symbolic_call(target, output)
+    return backend.nn.categorical_crossentropy(
+        target, output, from_logits=from_logits, axis=axis
+    )
+
+
+class SparseCategoricalCrossentropy(Operation):
+    def __init__(self, from_logits=False, axis=-1):
+        super().__init__()
+        self.from_logits = from_logits
+        self.axis = axis
+
+    def call(self, target, output):
+        return backend.nn.sparse_categorical_crossentropy(
+            target, output, from_logits=self.from_logits, axis=self.axis
+        )
+
+    def compute_output_spec(self, target, output):
+        if len(output.shape) < 1:
+            raise ValueError(
+                "Argument `output` must be at least rank 1. "
+                "Received: "
+                f"output.shape={output.shape}"
+            )
+        target_shape = target.shape
+        if len(target_shape) == len(output.shape) and target_shape[-1] == 1:
+            target_shape = target_shape[:-1]
+        if target_shape != output.shape[:-1]:
+            raise ValueError(
+                "Arguments `target` and `output` must have the same shape "
+                "up until the last dimension: "
+                f"target.shape={target.shape}, output.shape={output.shape}"
+            )
+        return KerasTensor(output.shape[:-1], dtype=output.dtype)
+
+
+@keras_export(
+    [
+        "keras.ops.sparse_categorical_crossentropy",
+        "keras.ops.nn.sparse_categorical_crossentropy",
+    ]
+)
+def sparse_categorical_crossentropy(target, output, from_logits=False, axis=-1):
+    """Computes sparse categorical cross-entropy loss.
+
+    The sparse categorical cross-entropy loss is similar to categorical
+    cross-entropy, but it is used when the target tensor contains integer
+    class labels instead of one-hot encoded vectors. It measures the
+    dissimilarity between the target and output probabilities or logits.
+
+    Args:
+        target: The target tensor representing the true class labels as
+            integers. Its shape should match the shape of the `output`
+            tensor except for the last dimension.
+        output: The output tensor representing the predicted probabilities
+            or logits.
+            Its shape should match the shape of the `target` tensor except
+            for the last dimension.
+        from_logits: (optional) Whether `output` is a tensor of logits
+            or probabilities.
+            Set it to `True` if `output` represents logits; otherwise,
+            set it to `False` if `output` represents probabilities.
+            Defaults to `False`.
+        axis: (optional) The axis along which the sparse categorical
+            cross-entropy is computed.
+            Defaults to `-1`, which corresponds to the last dimension
+            of the tensors.
+
+    Returns:
+        Integer tensor: The computed sparse categorical cross-entropy
+        loss between `target` and `output`.
+
+    Example:
+
+    >>> target = keras.ops.convert_to_tensor([0, 1, 2], dtype=int32)
+    >>> output = keras.ops.convert_to_tensor(
+    ... [[0.9, 0.05, 0.05],
+    ...  [0.1, 0.8, 0.1],
+    ...  [0.2, 0.3, 0.5]])
+    >>> sparse_categorical_crossentropy(target, output)
+    array([0.10536056 0.22314355 0.6931472 ], shape=(3,), dtype=float32)
+    """
+    if any_symbolic_tensors((target, output)):
+        return SparseCategoricalCrossentropy(
+            from_logits=from_logits, axis=axis
+        ).symbolic_call(target, output)
+    return backend.nn.sparse_categorical_crossentropy(
+        target, output, from_logits=from_logits, axis=axis
+    )
+
+
+class MultiHot(Operation):
+    def __init__(
+        self, num_classes=None, axis=-1, dtype=None, sparse=False, **kwargs
+    ):
+        if num_classes is None and "num_tokens" in kwargs:
+            num_classes = kwargs.pop("num_tokens")
+        if num_classes is None:
+            raise ValueError("Argument `num_classes` must be specified.")
+        super().__init__(**kwargs)
+        self.num_classes = num_classes
+        self.axis = axis
+        self.dtype = dtype or backend.floatx()
+        self.sparse = sparse
+
+    def call(self, inputs):
+        return backend.nn.multi_hot(
+            inputs,
+            num_classes=self.num_classes,
+            axis=self.axis,
+            dtype=self.dtype,
+        )
+
+    def compute_output_spec(self, inputs):
+        x_shape = list(getattr(inputs, "shape", []))
+        if self.axis == -1:
+            x_shape.append(self.num_classes)
+        elif self.axis >= 0 and self.axis < len(x_shape):
+            x_shape.insert(self.axis, self.num_classes)
+        else:
+            raise ValueError(
+                f"axis must be -1 or between [0, {len(inputs.shape)}), but "
+                f"received {self.axis}."
+            )
+
+        if len(x_shape) == 2:
+            x_shape = [x_shape[-1]]
+        else:
+            x_shape = [x_shape[0]] + x_shape[2:]
+
+        return KerasTensor(x_shape, dtype=inputs.dtype, sparse=self.sparse)
+
+
+@keras_export(
+    [
+        "keras.ops.multi_hot",
+        "keras.ops.nn.multi_hot",
+    ]
+)
+def multi_hot(
+    inputs, num_classes=None, axis=-1, dtype=None, sparse=False, **kwargs
+):
+    """Encodes integer labels as multi-hot vectors.
+
+    This function encodes integer labels as multi-hot vectors, where each label
+    is mapped to a binary value in the resulting vector.
+
+    Args:
+        inputs: Tensor of integer labels to be converted to multi-hot vectors.
+        num_classes: Integer, the total number of unique classes.
+        axis: (optional) Axis along which the multi-hot encoding should be
+            added. Defaults to `-1`, which corresponds to the last dimension.
+        dtype: (optional) The data type of the resulting tensor. Default
+            is backend's float type.
+        sparse: Whether to return a sparse tensor; for backends that support
+            sparse tensors.
+
+    Returns:
+        Tensor: The multi-hot encoded tensor.
+
+    Example:
+
+    >>> data = keras.ops.convert_to_tensor([0, 4])
+    >>> keras.ops.multi_hot(data, num_classes=5)
+    array([1.0, 0.0, 0.0, 0.0, 1.0], dtype=float32)
+
+    """
+    if num_classes is None and "num_tokens" in kwargs:
+        num_classes = kwargs.pop("num_tokens")
+    if num_classes is None:
+        raise ValueError("Argument `num_classes` must be specified.")
+
+    if any_symbolic_tensors((inputs,)):
+        return MultiHot(num_classes, axis, dtype, sparse).symbolic_call(inputs)
+
+    return backend.nn.multi_hot(inputs, num_classes, axis, dtype, sparse)
+
+
+class Moments(Operation):
+    def __init__(self, axes, keepdims=False, synchronized=False):
+        super().__init__()
+        self.axes = axes
+        self.keepdims = keepdims
+        self.synchronized = synchronized
+
+    def call(self, x):
+        return backend.nn.moments(
+            x,
+            axes=self.axes,
+            keepdims=self.keepdims,
+            synchronized=self.synchronized,
+        )
+
+    def compute_output_spec(self, x):
+        return (
+            KerasTensor(
+                reduce_shape(x.shape, axis=self.axes, keepdims=self.keepdims),
+                dtype=x.dtype,
+            ),
+            KerasTensor(
+                reduce_shape(x.shape, axis=self.axes, keepdims=self.keepdims),
+                dtype=x.dtype,
+            ),
+        )
+
+
+@keras_export(
+    [
+        "keras.ops.moments",
+        "keras.ops.nn.moments",
+    ]
+)
+def moments(x, axes, keepdims=False, synchronized=False):
+    """Calculates the mean and variance of `x`.
+
+    The mean and variance are calculated by aggregating the contents of `x`
+    across `axes`. If `x` is 1-D and `axes = [0]` this is just the mean and
+    variance of a vector.
+
+    Args:
+        x: Input tensor.
+        axes: A list of axes which to compute mean and variance.
+        keepdims: If this is set to `True`, the axes which are reduced are left
+            in the result as dimensions with size one.
+        synchronized: Only applicable with the TensorFlow backend.
+            If `True`, synchronizes the global batch statistics (mean and
+            variance) across all devices at each training step in a
+            distributed training strategy. If `False`, each replica uses its own
+            local batch statistics.
+
+    Returns:
+        A tuple containing two tensors - mean and variance.
+
+    Example:
+
+    >>> x = keras.ops.convert_to_tensor([0, 1, 2, 3, 100], dtype="float32")
+    >>> keras.ops.moments(x, axes=[0])
+    (array(21.2, dtype=float32), array(1553.3601, dtype=float32))
+
+    """
+    if any_symbolic_tensors((x,)):
+        return Moments(axes, keepdims, synchronized=synchronized).symbolic_call(
+            x
+        )
+
+    return backend.nn.moments(x, axes, keepdims, synchronized=synchronized)
+
+
+class BatchNorm(Operation):
+    def __init__(self, axis, epsilon):
+        super().__init__()
+        self.axis = axis
+        self.epsilon = epsilon
+
+    def _check_shape(self, name, shape, expected_shape):
+        if shape != expected_shape:
+            raise ValueError(
+                f"Arguments `{name}` must be a vector of length "
+                f"`x.shape[axis]`. Expected: `{expected_shape}`. "
+                f"Received: `{shape}."
+            )
+
+    def compute_output_spec(self, x, mean, variance, offset, scale):
+        shape = (x.shape[self.axis],)
+        self._check_shape("mean", tuple(mean.shape), shape)
+        self._check_shape("variance", tuple(variance.shape), shape)
+        if offset is not None:
+            self._check_shape("offset", tuple(offset.shape), shape)
+        if offset is not scale:
+            self._check_shape("scale", tuple(scale.shape), shape)
+        return KerasTensor(x.shape, dtype=x.dtype)
+
+
+@keras_export(
+    [
+        "keras.ops.batch_normalization",
+        "keras.ops.nn.batch_normalization",
+    ]
+)
+def batch_normalization(
+    x, mean, variance, axis, offset=None, scale=None, epsilon=1e-3
+):
+    """Normalizes `x` by `mean` and `variance`.
+
+    This op is typically used by the batch normalization step in a neural
+    network. It normalizes the input tensor along the given axis.
+
+    Args:
+        x: Input tensor.
+        mean: A mean vector of the same length as the `axis` dimension of the
+            input thensor.
+        variance: A variance vector of the same length as the `axis` dimension
+            of the input tensor.
+        axis: Integer, the axis that should be normalized.
+        offset: An offset vector of the same length as the `axis` dimension of
+            the input tensor. If not `None`, `offset` is added to the normalized
+            tensor. Defaults to `None`.
+        scale: A scale vector of the same length as the `axis` dimension of the
+            input tensor. If not `None`, the normalized tensor is multiplied by
+            `scale`. Defaults to `None`.
+        epsilon: Small float added to variance to avoid dividing by zero.
+            Defaults to 1e-3.
+
+    Returns:
+        The normalized tensor.
+
+    Example:
+
+    >>> x = keras.ops.convert_to_tensor(
+    ...     [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6], [0.7, 0.8, 0.9]]
+    ... )
+    >>> keras.ops.batch_normalization(
+    ...     x,
+    ...     mean=[0.4, 0.5, 0.6],
+    ...     variance=[0.67, 0.67, 0.67],
+    ...     axis=-1
+    ... )
+    array([[-3.6624e-01, -3.6624e-01, -3.6624e-01],
+           [-4.6445e-09,  0.0000e+00, -1.8578e-08],
+           [ 3.6624e-01,  3.6624e-01,  3.6624e-01]])
+
+    """
+    if any_symbolic_tensors((x, mean, variance, offset, scale)):
+        return BatchNorm(axis, epsilon).symbolic_call(
+            x, mean, variance, offset, scale
+        )
+
+    return backend.nn.batch_normalization(
+        x, mean, variance, axis, offset, scale, epsilon
+    )
+
+
+class CTCLoss(Operation):
+    def __init__(self, mask_index=0):
+        super().__init__()
+        self.mask_index = mask_index
+
+    def call(self, target, output, target_length, output_length):
+        return backend.nn.ctc_loss(
+            target, output, target_length, output_length, self.mask_index
+        )
+
+    def _check_shape_first_dim(self, name1, shape1, name2, shape2):
+        if shape1[0] != shape2[0]:
+            raise ValueError(
+                f"Arguments `{name1}` and `{name2}` must have the same "
+                "first dimension. "
+                f"Received shapes: `{shape1}` and `{shape2}`."
+            )
+
+    def compute_output_spec(self, target, output, target_length, output_length):
+        self._check_shape_first_dim(
+            "target", target.shape, "output", output.shape
+        )
+        self._check_shape_first_dim(
+            "target_length", target_length.shape, "target", target.shape
+        )
+        self._check_shape_first_dim(
+            "output_length", output_length.shape, "output", output.shape
+        )
+        dtype = backend.result_type(output.dtype, "float32")
+        return KerasTensor((target.shape[0],), dtype=dtype)
+
+
+@keras_export(
+    [
+        "keras.ops.ctc_loss",
+        "keras.ops.nn.ctc_loss",
+    ]
+)
+def ctc_loss(target, output, target_length, output_length, mask_index=0):
+    """CTC (Connectionist Temporal Classification) loss.
+
+    Args:
+        target: A tensor of shape `(batch_size, max_length)` containing
+            the true labels in integer format.
+        output: A tensor of shape `(batch_size, max_length, num_classes)`
+            containing logits (the output of your model).
+        target_length: A tensor of shape `(batch_size,)` containing the
+            true label lengths.
+        output_length: A tensor of shape `(batch_size,)` containing the
+            output lengths.
+        mask_index: The index of the mask character in the vocabulary.
+            Defaults to `0`.
+    """
+
+    if any_symbolic_tensors((target, output, target_length, output_length)):
+        return CTCLoss(mask_index).symbolic_call(
+            target, output, target_length, output_length
+        )
+    return backend.nn.ctc_loss(
+        target, output, target_length, output_length, mask_index
+    )
+
+
+class CTCDecode(Operation):
+    def __init__(
+        self,
+        strategy="greedy",
+        beam_width=100,
+        top_paths=1,
+        merge_repeated=True,
+        mask_index=0,
+    ):
+        super().__init__()
+        self.strategy = strategy
+        self.beam_width = beam_width
+        self.top_paths = top_paths
+        self.merge_repeated = merge_repeated
+        self.mask_index = mask_index
+
+    def call(self, inputs, sequence_lengths):
+        return backend.nn.ctc_decode(
+            inputs,
+            sequence_lengths,
+            strategy=self.strategy,
+            beam_width=self.beam_width,
+            top_paths=self.top_paths,
+            merge_repeated=self.merge_repeated,
+            mask_index=self.mask_index,
+        )
+
+    def compute_output_spec(self, inputs, sequence_lengths):
+        inputs_shape = inputs.shape
+        if self.strategy == "greedy":
+            top_paths = 1
+        else:
+            top_paths = self.top_paths
+        dtype = backend.result_type(inputs.dtype, "float32")
+        return (
+            KerasTensor(
+                (top_paths, inputs_shape[0], inputs_shape[1]), dtype="int32"
+            ),
+            KerasTensor((inputs_shape[0], top_paths), dtype=dtype),
+        )
+
+
+@keras_export(
+    [
+        "keras.ops.ctc_decode",
+        "keras.ops.nn.ctc_decode",
+    ]
+)
+def ctc_decode(
+    inputs,
+    sequence_lengths,
+    strategy="greedy",
+    beam_width=100,
+    top_paths=1,
+    merge_repeated=True,
+    mask_index=0,
+):
+    """Decodes the output of a CTC model.
+
+    Args:
+        inputs: A tensor of shape `(batch_size, max_length, num_classes)`
+            containing the logits (the output of the model).
+            They should *not* be normalized via softmax.
+        sequence_lengths: A tensor of shape `(batch_size,)` containing the
+            sequence lengths for the batch.
+        strategy: A string for the decoding strategy. Supported values are
+            `"greedy"` and `"beam_search"`.
+        beam_width: An integer scalar beam width used in beam search.
+            Defaults to 100.
+        top_paths: An integer scalar, the number of top paths to return.
+            Defaults to 1.
+        merge_repeated: A boolean scalar, whether to merge repeated
+            labels in the output. Defaults to `True`.
+        mask_index: An integer scalar, the index of the mask character in
+            the vocabulary. Defaults to `0`.
+
+    Returns:
+        A tuple containing:
+        - The tensor representing the list of decoded sequences. If
+            `strategy="greedy"`, the shape is `(1, batch_size, max_length)`. If
+            `strategy="beam_search"`, the shape is
+            `(top_paths, batch_size, max_length)`. Note that: `-1` indicates the
+            blank label.
+        - If `strategy="greedy"`, a tensor of shape `(batch_size, 1)`
+            representing the negative of the sum of the probability logits for
+            each sequence. If `strategy="beam_seatch"`, a tensor of shape
+            `(batch_size, top_paths)` representing the log probability for each
+            sequence.
+    """
+
+    if any_symbolic_tensors((inputs, sequence_lengths)):
+        return CTCDecode(
+            strategy=strategy,
+            beam_width=beam_width,
+            top_paths=top_paths,
+            merge_repeated=merge_repeated,
+            mask_index=mask_index,
+        ).symbolic_call(inputs, sequence_lengths)
+    return backend.nn.ctc_decode(
+        inputs=inputs,
+        sequence_lengths=sequence_lengths,
+        strategy=strategy,
+        beam_width=beam_width,
+        top_paths=top_paths,
+        merge_repeated=merge_repeated,
+        mask_index=mask_index,
+    )
+
+
+class Normalize(Operation):
+    def __init__(self, axis=-1, order=2, epsilon=None):
+        super().__init__()
+        self.axis = axis
+        self.order = order
+        self.epsilon = epsilon
+
+    def compute_output_spec(self, x):
+        return KerasTensor(shape=x.shape)
+
+    def call(self, x):
+        return _normalize(
+            x, axis=self.axis, order=self.order, epsilon=self.epsilon
+        )
+
+
+@keras_export(
+    [
+        "keras.ops.normalize",
+        "keras.ops.nn.normalize",
+    ]
+)
+def normalize(x, axis=-1, order=2, epsilon=None):
+    """Normalizes `x` over the specified axis.
+
+    It is defined as: `normalize(x) = x / max(norm(x), epsilon)`.
+
+    Args:
+        x: Input tensor.
+        axis: The axis or axes along which to perform normalization.
+            Default to -1.
+        order: The exponent value in the norm formulation.
+            Defaults to 2.
+        epsilon: A lower bound value for the norm.
+            Defaults to `backend.epsilon()`.
+
+    Returns:
+        The normalized array.
+
+    Example:
+
+    >>> x = keras.ops.convert_to_tensor([[1, 2, 3], [4, 5, 6]])
+    >>> x_norm = keras.ops.math.normalize(x)
+    >>> print(x_norm)
+    array([[0.26726124 0.5345225  0.8017837 ]
+           [0.45584232 0.5698029  0.68376344]], shape=(2, 3), dtype=float32)
+
+    """
+    if any_symbolic_tensors((x,)):
+        return Normalize(axis=axis, order=order, epsilon=epsilon).symbolic_call(
+            x
+        )
+    return _normalize(x, axis=axis, order=order, epsilon=epsilon)
+
+
+def _normalize(x, axis=-1, order=2, epsilon=None):
+    if not isinstance(order, int) or not order >= 1:
+        raise ValueError(
+            f"Argument `order` must be an int >= 1. Received: order={order}"
+        )
+    x = backend.convert_to_tensor(x)
+    if len(x.shape) == 0:
+        x = backend.numpy.expand_dims(x, axis=0)
+    if epsilon is None:
+        epsilon = backend.epsilon()
+    if 2 == order:
+        # A special case: L2 normalization with `x * rsqrt(...)`
+        # instead of `x / sqrt(...)`
+        square_sum = backend.numpy.sum(
+            backend.numpy.square(x), axis=axis, keepdims=True
+        )
+        inv_norm = backend.math.rsqrt(square_sum)
+        inv_norm = backend.numpy.minimum(inv_norm, 1.0 / epsilon)
+        return x * inv_norm
+    norm = backend.linalg.norm(x, ord=order, axis=axis, keepdims=True)
+    denom = backend.numpy.maximum(norm, epsilon)
+    return backend.numpy.divide(x, denom)
+
+
+class PSNR(Operation):
+    def __init__(
+        self,
+        max_val,
+    ):
+        super().__init__()
+        self.max_val = max_val
+
+    def call(self, x1, x2):
+        return backend.nn.psnr(
+            x1=x1,
+            x2=x2,
+            max_val=self.max_val,
+        )
+
+    def compute_output_spec(self, x1, x2):
+        if len(x1.shape) != len(x2.shape):
+            raise ValueError("Inputs must have the same rank")
+
+        return KerasTensor(shape=())
+
+
+@keras_export(
+    [
+        "keras.ops.psnr",
+        "keras.ops.nn.psnr",
+    ]
+)
+def psnr(
+    x1,
+    x2,
+    max_val,
+):
+    """Peak Signal-to-Noise Ratio (PSNR) function.
+
+    This function computes the Peak Signal-to-Noise Ratio between two signals,
+    `x1` and `x2`. PSNR is a measure of the quality of a reconstructed signal.
+    The higher the PSNR, the closer the reconstructed signal is to the original
+    signal. Note that it can become negative when the signal power is
+    smaller that the noise power.
+
+    Args:
+        x1: The first input signal.
+        x2: The second input signal. Must have the same shape as `x1`.
+        max_val: The maximum possible value in the signals.
+
+    Returns:
+        float: The PSNR value between `x1` and `x2`.
+
+    Examples:
+
+    >>> x1 = keras.random.normal((2, 4, 4, 3))
+    >>> x2 = keras.random.normal((2, 4, 4, 3))
+    >>> max_val = 1.0
+    >>> keras.ops.nn.psnr(x1, x2, max_val)
+    -3.1697404
+    """
+    if any_symbolic_tensors(
+        (
+            x1,
+            x2,
+        )
+    ):
+        return PSNR(
+            max_val,
+        ).symbolic_call(x1, x2)
+    return backend.nn.psnr(
+        x1,
+        x2,
+        max_val,
+    )
+
+
+class DotProductAttention(Operation):
+    def __init__(self, is_causal=False):
+        super().__init__()
+        self.is_causal = is_causal
+
+    def call(
+        self,
+        query,
+        key,
+        value,
+        bias=None,
+        mask=None,
+        scale=None,
+        flash_attention=None,
+    ):
+        return backend.nn.dot_product_attention(
+            query,
+            key,
+            value,
+            bias=bias,
+            mask=mask,
+            scale=scale,
+            is_causal=self.is_causal,
+            flash_attention=flash_attention,
+        )
+
+    def compute_output_spec(
+        self,
+        query,
+        key,
+        value,
+        bias=None,
+        mask=None,
+        scale=None,
+        flash_attention=None,
+    ):
+        return KerasTensor(query.shape, dtype=query.dtype)
+
+
+@keras_export(
+    ["keras.ops.dot_product_attention", "keras.ops.nn.dot_product_attention"]
+)
+def dot_product_attention(
+    query,
+    key,
+    value,
+    bias=None,
+    mask=None,
+    scale=None,
+    is_causal=False,
+    flash_attention=None,
+):
+    """Scaled dot product attention function.
+
+    Computes the attention function on Q (`query`), K (`key`), and V(`value`):
+    `attention(Q, K, V) = softmax(Q * K / sqrt(d)) * V`. If we define `logits`
+    as the output of `Q * K` and the `probs` as the output of `softmax`.
+
+    Throughout this function, we utilize the following notation to represent the
+    shape of array:
+    - B: batch size
+    - S: length of the key/value
+    - T: length of the query
+    - N: number of attention heads
+    - H: dimensions of each attention head
+    - K: number of key/value heads
+    - G: number of groups, which equals to `N // K`
+
+    Args:
+        query: The query array with the shape of `(B, T, N, H)`.
+        key: The key array with the shape of `(B, S, K, H)`. When `K` equals
+            `N`, multi-headed attention (MHA) is performed. Otherwise, grouped
+            query attention (GQA) is performed if `N` is a multiple of `K`. and
+            multi-query attention (MQA) is performed if `K==1` (a special case
+            of GQA).
+        value: The value array with the same shape of `key`.
+        bias: Optional bias array to be added to logits. The shape must be
+            broadcastable to `(B, N, T, S)`.
+        mask: Optional mask array used to filter out logits. It is a boolean
+            mask where `True` indicates the element should take part in
+            attention. For an additive mask, users should pass it to bias. The
+            shape must be broadcastable to `(B, N, T, S)`.
+        scale: Optional scale for the logits. If `None`, the scale will be set
+            to `1.0 / sqrt(H)`.
+        is_causal: Whether to apply causal mask.
+        flash_attention: Whether to use flash attention. If `None`, it will
+            attempt to use flash attention if the required conditions are met.
+            Typically, the inputs must be in float16 and bfloat16 dtype and the
+            input layout requirements may vary depending on the backend.
+
+    Returns:
+        An array of the attention output with the same shape of `query`.
+
+    Example:
+
+    >>> query = keras.random.normal((2, 4, 8, 16))
+    >>> key = keras.random.normal((2, 6, 8, 16))
+    >>> value = keras.random.normal((2, 6, 8, 16))
+    >>> keras.ops.nn.dot_product_attention(query, key, value).shape
+    (2, 4, 8, 16)
+    """
+    if any_symbolic_tensors((query, key, value)):
+        return DotProductAttention(is_causal=is_causal).symbolic_call(
+            query,
+            key,
+            value,
+            bias=bias,
+            mask=mask,
+            scale=scale,
+            flash_attention=flash_attention,
+        )
+    return backend.nn.dot_product_attention(
+        query,
+        key,
+        value,
+        bias=bias,
+        mask=mask,
+        scale=scale,
+        is_causal=is_causal,
+        flash_attention=flash_attention,
+    )

SwarmUI/dlbackend/ComfyUI/venv/lib/python3.10/site-packages/keras/src/ops/node.py

@@ -0,0 +1,143 @@
+import collections
+
+from keras.src import tree
+from keras.src.backend import KerasTensor
+from keras.src.ops.symbolic_arguments import SymbolicArguments
+
+
+class Node:
+    """A `Node` describes an operation `__call__()` event.
+
+    A Keras Function is a DAG with `Node` instances as nodes, and
+    `KerasTensor` instances as edges. Nodes aren't `Operation` instances,
+    because a single operation could be called multiple times, which would
+    result in graph cycles.
+
+    A `__call__()` event involves input tensors (and other input arguments),
+    the operation that was called, and the resulting output tensors.
+    A `Node` will include all this information.
+
+    Since a single `Operation` could be called multiple times,
+    the `Node` instances are stored on operations as a list.
+    Each time an operation is called, a node is added to `op._inbound_nodes`.
+    Each time the output of an operation is used by another operation,
+    a node is added to `op._outbound_nodes`.
+
+    Every `KerasTensor` instance has a `KerasHistory` object attached,
+    which tracks the `Node` that records the `__call__()` event that created
+    the tensor. By recursively walking through `Node` instances
+    via the `KerasHistory` metadata of `KerasTensor` instances, once can
+    retrieve the entire DAG of a Keras Function.
+
+    Args:
+        operation: The Operation that was called in the `op.__call__()`
+            event that this node represents.
+        call_args: The positional arguments the operation was called with.
+        call_kwargs: The keyword arguments the operation was called with.
+        outputs: The output tensors of the `op.__call__()` call.
+    """
+
+    def __init__(
+        self, operation, call_args=None, call_kwargs=None, outputs=None
+    ):
+        self.operation = operation
+        self.arguments = SymbolicArguments(*call_args, **call_kwargs)
+        self.outputs = [] if outputs is None else tree.flatten(outputs)
+        for x in self.outputs:
+            if not isinstance(x, KerasTensor):
+                raise ValueError(
+                    "All operation outputs must be tensors. "
+                    f"Operation {operation} returned a non-tensor. "
+                    f"Non-tensor received: {x}"
+                )
+
+        zero_history = any(
+            not x.record_history for x in self.arguments.keras_tensors
+        )
+
+        # If inputs don't have metadata yet, add it.
+        if not zero_history:
+            for tensor in self.arguments.keras_tensors:
+                if not hasattr(tensor, "_keras_history"):
+                    tensor._keras_history = KerasHistory(
+                        operation=None, node_index=0, tensor_index=0
+                    )
+
+        # Wire up Node to Operations.
+        self.operation._inbound_nodes.append(self)
+        for kt in self.arguments.keras_tensors:
+            inbound_op = kt._keras_history.operation
+            if inbound_op is not None:  # It's a graph entry point.
+                inbound_op._outbound_nodes.append(self)
+
+        # Set metadata on outputs.
+        if not zero_history:
+            node_index = len(self.operation._inbound_nodes) - 1
+            for i, tensor in enumerate(self.outputs):
+                tensor._keras_history = KerasHistory(
+                    operation=operation, node_index=node_index, tensor_index=i
+                )
+
+        # Whether this is a root node.
+        self.is_input = not self.arguments.keras_tensors
+
+    def __repr__(self):
+        return f"<Node operation={self.operation.name}, id={id(self)}>"
+
+    @property
+    def input_tensors(self):
+        return self.arguments.keras_tensors
+
+    @property
+    def output_tensors(self):
+        return self.outputs
+
+    @property
+    def parent_nodes(self):
+        """The parent `Node`s.
+
+        Returns:
+            all the `Node`s whose output this node immediately depends on.
+        """
+        node_deps = []
+        for kt in self.arguments.keras_tensors:
+            op = kt._keras_history.operation
+            node_index = kt._keras_history.node_index
+            if op is not None:  # `None` for `Input` tensors.
+                node_deps.append(op._inbound_nodes[node_index])
+        return node_deps
+
+
+class KerasHistory(
+    collections.namedtuple(
+        "KerasHistory", ["operation", "node_index", "tensor_index"]
+    )
+):
+    """Tracks the Operation call that created a Tensor.
+
+    During construction of Keras Functions, this metadata is added to
+    each Tensor produced as the output of an Operation.
+    This allows Keras to track how each Tensor was produced, and
+    this information is later retraced by the `Function` class to
+    reconstruct the Operations graph.
+
+    Attributes:
+      operation: The Operation instance that produced the Tensor.
+      node_index: The specific call to the Operation that produced this Tensor.
+        Operations can be called multiple times in order to share weights. A new
+        node is created every time an Operation is called. The corresponding
+        node that represents the call event that produced the Tensor can be
+        found at `op._inbound_nodes[node_index]`.
+      tensor_index: The output index for this Tensor.
+        Always zero if the Operation that produced this Tensor
+        only has one output. Nested structures of
+        Tensors are deterministically assigned an index via `nest.flatten`.
+    """
+
+    # Added to maintain memory and performance characteristics of `namedtuple`
+    # while subclassing.
+    __slots__ = ()
+
+
+def is_keras_tensor(obj):
+    return hasattr(obj, "_keras_history")

SwarmUI/dlbackend/ComfyUI/venv/lib/python3.10/site-packages/keras/src/ops/operation.py

@@ -0,0 +1,316 @@
+import inspect
+import textwrap
+
+from keras.src import backend
+from keras.src import dtype_policies
+from keras.src import tree
+from keras.src.api_export import keras_export
+from keras.src.backend.common.keras_tensor import any_symbolic_tensors
+from keras.src.ops.node import Node
+from keras.src.utils import python_utils
+from keras.src.utils import traceback_utils
+from keras.src.utils.naming import auto_name
+
+
+@keras_export("keras.Operation")
+class Operation:
+    def __init__(self, dtype=None, name=None):
+        if name is None:
+            name = auto_name(self.__class__.__name__)
+        if not isinstance(name, str) or "/" in name:
+            raise ValueError(
+                "Argument `name` must be a string and "
+                "cannot contain character `/`. "
+                f"Received: name={name} (of type {type(name)})"
+            )
+        self._dtype_policy = dtype_policies.get(dtype)
+        self.name = name
+        self._inbound_nodes = []
+        self._outbound_nodes = []
+
+    @traceback_utils.filter_traceback
+    def __call__(self, *args, **kwargs):
+        if traceback_utils.is_traceback_filtering_enabled():
+            # Wrap self.call to provide helpful info in case of exception
+            if any_symbolic_tensors(args, kwargs):
+                call_fn = self.symbolic_call
+            else:
+                if getattr(self, "quantization_mode", None) is not None:
+                    call_fn = self.quantized_call
+                else:
+                    call_fn = self.call
+            call_fn = traceback_utils.inject_argument_info_in_traceback(
+                call_fn,
+                object_name=(f"{self.__class__.__name__}.call()"),
+            )
+            return call_fn(*args, **kwargs)
+
+        # Plain flow.
+        if any_symbolic_tensors(args, kwargs):
+            return self.symbolic_call(*args, **kwargs)
+        if getattr(self, "quantization_mode", None) is not None:
+            return self.quantized_call(*args, **kwargs)
+        else:
+            return self.call(*args, **kwargs)
+
+    def symbolic_call(self, *args, **kwargs):
+        # Perform shape/dtype inference.
+        outputs = self.compute_output_spec(*args, **kwargs)
+        # Record a new node in the operations graph.
+        # The Node wires itself to inbound and outbound ops.  The
+        # Node constructor updates this op's self._inbound_nodes,
+        # sets _keras_history on the outputs, and adds itself to the
+        # `_outbound_nodes` of the ops that produced the inputs to this
+        # call.
+        Node(
+            operation=self, call_args=args, call_kwargs=kwargs, outputs=outputs
+        )
+        return outputs
+
+    def call(self, *args, **kwargs):
+        raise NotImplementedError
+
+    def quantized_call(self, *args, **kwargs):
+        raise NotImplementedError
+
+    def compute_output_spec(self, *args, **kwargs):
+        try:
+            return backend.compute_output_spec(self.call, *args, **kwargs)
+        except Exception as e:
+            new_e = e.__class__(
+                "Could not automatically infer the output shape / dtype of "
+                f"'{self.name}' (of type {self.__class__.__name__}). "
+                f"Either the `{self.__class__.__name__}.call()` method "
+                f"is incorrect, or you need to implement the "
+                f"`{self.__class__.__name__}.compute_output_spec() / "
+                "compute_output_shape()` method. "
+                f"Error encountered:\n\n{e}"
+            )
+            raise new_e.with_traceback(e.__traceback__) from None
+
+    def __new__(cls, *args, **kwargs):
+        """We override __new__ to saving serializable constructor arguments.
+
+        These arguments are used to auto-generate an object serialization
+        config, which enables user-created subclasses to be serializable
+        out of the box in most cases without forcing the user
+        to manually implement `get_config()`.
+        """
+        instance = super(Operation, cls).__new__(cls)
+
+        # Generate a config to be returned by default by `get_config()`.
+        arg_names = inspect.getfullargspec(cls.__init__).args
+        kwargs.update(dict(zip(arg_names[1 : len(args) + 1], args)))
+
+        # Explicitly serialize `dtype` to support auto_config
+        dtype = kwargs.get("dtype", None)
+        if dtype is not None and isinstance(dtype, dtype_policies.DTypePolicy):
+            # For backward compatibility, we use a str (`name`) for
+            # `DTypePolicy`
+            if dtype.quantization_mode is None:
+                kwargs["dtype"] = dtype.name
+            # Otherwise, use `dtype_policies.serialize`
+            else:
+                kwargs["dtype"] = dtype_policies.serialize(dtype)
+
+        # For safety, we only rely on auto-configs for a small set of
+        # serializable types.
+        supported_types = (str, int, float, bool, type(None))
+        try:
+            flat_arg_values = tree.flatten(kwargs)
+            auto_config = True
+            for value in flat_arg_values:
+                if not isinstance(value, supported_types):
+                    auto_config = False
+                    break
+        except TypeError:
+            auto_config = False
+        try:
+            instance._lock = False
+            if auto_config:
+                from keras.src.saving import serialization_lib
+
+                instance._auto_config = serialization_lib.SerializableDict(
+                    **kwargs
+                )
+            else:
+                instance._auto_config = None
+            instance._lock = True
+        except RecursionError:
+            # Setting an instance attribute in __new__ has the potential
+            # to trigger an infinite recursion if a subclass overrides
+            # setattr in an unsafe way.
+            pass
+        return instance
+
+    @python_utils.default
+    def get_config(self):
+        """Returns the config of the object.
+
+        An object config is a Python dictionary (serializable)
+        containing the information needed to re-instantiate it.
+        """
+        config = {
+            "name": self.name,
+        }
+
+        if not python_utils.is_default(self.get_config):
+            # In this case the subclass implements get_config()
+            return config
+
+        # In this case the subclass doesn't implement get_config():
+        # Let's see if we can autogenerate it.
+        if getattr(self, "_auto_config", None) is not None:
+            xtra_args = set(config.keys())
+            config.update(self._auto_config.config)
+            # Remove args non explicitly supported
+            argspec = inspect.getfullargspec(self.__init__)
+            if argspec.varkw != "kwargs":
+                for key in xtra_args - xtra_args.intersection(argspec.args[1:]):
+                    config.pop(key, None)
+            return config
+        else:
+            raise NotImplementedError(
+                textwrap.dedent(
+                    f"""
+        Object {self.__class__.__name__} was created by passing
+        non-serializable argument values in `__init__()`,
+        and therefore the object must override `get_config()` in
+        order to be serializable. Please implement `get_config()`.
+
+        Example:
+
+        class CustomLayer(keras.layers.Layer):
+            def __init__(self, arg1, arg2, **kwargs):
+                super().__init__(**kwargs)
+                self.arg1 = arg1
+                self.arg2 = arg2
+
+            def get_config(self):
+                config = super().get_config()
+                config.update({{
+                    "arg1": self.arg1,
+                    "arg2": self.arg2,
+                }})
+                return config"""
+                )
+            )
+
+    @classmethod
+    def from_config(cls, config):
+        """Creates an operation from its config.
+
+        This method is the reverse of `get_config`, capable of instantiating the
+        same operation from the config dictionary.
+
+        Note: If you override this method, you might receive a serialized dtype
+        config, which is a `dict`. You can deserialize it as follows:
+
+        ```python
+        if "dtype" in config and isinstance(config["dtype"], dict):
+            policy = dtype_policies.deserialize(config["dtype"])
+        ```
+
+        Args:
+            config: A Python dictionary, typically the output of `get_config`.
+
+        Returns:
+            An operation instance.
+        """
+        # Explicitly deserialize dtype config if needed. This enables users to
+        # directly interact with the instance of `DTypePolicy`.
+        if "dtype" in config and isinstance(config["dtype"], dict):
+            config = config.copy()
+            policy = dtype_policies.deserialize(config["dtype"])
+            if (
+                not isinstance(policy, dtype_policies.DTypePolicyMap)
+                and policy.quantization_mode is None
+            ):
+                # For backward compatibility, we use a str (`name`) for
+                # `DTypePolicy`
+                policy = policy.name
+            config["dtype"] = policy
+        try:
+            return cls(**config)
+        except Exception as e:
+            raise TypeError(
+                f"Error when deserializing class '{cls.__name__}' using "
+                f"config={config}.\n\nException encountered: {e}"
+            )
+
+    def __repr__(self):
+        return f"<Operation name={self.name}>"
+
+    @property
+    def input(self):
+        """Retrieves the input tensor(s) of a symbolic operation.
+
+        Only returns the tensor(s) corresponding to the *first time*
+        the operation was called.
+
+        Returns:
+            Input tensor or list of input tensors.
+        """
+        return self._get_node_attribute_at_index(0, "input_tensors", "input")
+
+    @property
+    def output(self):
+        """Retrieves the output tensor(s) of a layer.
+
+        Only returns the tensor(s) corresponding to the *first time*
+        the operation was called.
+
+        Returns:
+            Output tensor or list of output tensors.
+        """
+        return self._get_node_attribute_at_index(0, "output_tensors", "output")
+
+    def _get_node_attribute_at_index(self, node_index, attr, attr_name):
+        """Private utility to retrieves an attribute (e.g. inputs) from a node.
+
+        This is used to implement the properties:
+        - output
+        - input
+
+        Args:
+            node_index: Integer index of the node from which
+                to retrieve the attribute.
+            attr: Exact node attribute name.
+            attr_name: Human-readable attribute name, for error messages.
+
+        Returns:
+            The operation's attribute `attr` at the node of index `node_index`.
+        """
+        if not self._inbound_nodes:
+            raise AttributeError(
+                f"The layer {self.name} has never been called "
+                f"and thus has no defined {attr_name}."
+            )
+        if not len(self._inbound_nodes) > node_index:
+            raise ValueError(
+                f"Asked to get {attr_name} at node "
+                f"{node_index}, but the operation has only "
+                f"{len(self._inbound_nodes)} inbound nodes."
+            )
+        values = getattr(self._inbound_nodes[node_index], attr)
+        if isinstance(values, list) and len(values) == 1:
+            return values[0]
+        else:
+            return values
+
+    # Hooks for backend layer classes
+    def _post_build(self):
+        """Can be overridden for per backend post build actions."""
+        pass
+
+    def _setattr_hook(self, name, value):
+        """Can be overridden for per backend post build actions."""
+        return name, value
+
+    def _post_track_variable(self, variable):
+        """Can be overridden for per backend post track actions."""
+        pass
+
+    def _post_untrack_variable(self, variable):
+        """Can be overridden for per backend post untrack actions."""
+        pass

SwarmUI/dlbackend/ComfyUI/venv/lib/python3.10/site-packages/keras/src/ops/operation_utils.py

@@ -0,0 +1,421 @@
+import math
+
+import numpy as np
+
+from keras.src import tree
+from keras.src.api_export import keras_export
+from keras.src.backend.common.backend_utils import canonicalize_axis
+from keras.src.backend.common.backend_utils import to_tuple_or_list
+
+
+def broadcast_shapes(shape1, shape2):
+    """Broadcast input shapes to a unified shape.
+
+    Convert to list for mutability.
+
+    Args:
+        shape1: A tuple or list of integers.
+        shape2: A tuple or list of integers.
+
+    Returns:
+        output_shape (list of integers or `None`): The broadcasted shape.
+
+    Example:
+    >>> broadcast_shapes((5, 3), (1, 3))
+    [5, 3]
+    """
+    shape1 = list(shape1)
+    shape2 = list(shape2)
+    origin_shape1 = shape1
+    origin_shape2 = shape2
+
+    if len(shape1) > len(shape2):
+        shape2 = [1] * (len(shape1) - len(shape2)) + shape2
+    if len(shape1) < len(shape2):
+        shape1 = [1] * (len(shape2) - len(shape1)) + shape1
+    output_shape = list(shape1)
+    for i in range(len(shape1)):
+        if shape1[i] == 1:
+            output_shape[i] = shape2[i]
+        elif shape1[i] is None:
+            output_shape[i] = None if shape2[i] == 1 else shape2[i]
+        else:
+            if shape2[i] == 1 or shape2[i] is None or shape2[i] == shape1[i]:
+                output_shape[i] = shape1[i]
+            else:
+                raise ValueError(
+                    "Cannot broadcast shape, the failure dim has value "
+                    f"{shape1[i]}, which cannot be broadcasted to {shape2[i]}. "
+                    f"Input shapes are: {origin_shape1} and {origin_shape2}."
+                )
+
+    return output_shape
+
+
+def compute_expand_dims_output_shape(input_shape, axis):
+    """Compute the output shape for the `expand_dims` operation.
+
+    Args:
+        input_shape: Input shape.
+        axis: int or sequence of ints for the axis to expand.
+
+    Returns:
+        Tuple of ints: The output shape after the `expand_dims` operation.
+    """
+    input_shape = list(input_shape)
+    if axis is None:
+        axis = len(input_shape)
+    axis = to_tuple_or_list(axis)
+    out_ndim = len(axis) + len(input_shape)
+    axis = [canonicalize_axis(a, out_ndim) for a in axis]
+    shape_iter = iter(input_shape)
+    new_shape = [
+        1 if ax in axis else next(shape_iter) for ax in range(out_ndim)
+    ]
+    return tuple(new_shape)
+
+
+def compute_pooling_output_shape(
+    input_shape,
+    pool_size,
+    strides,
+    padding="valid",
+    data_format="channels_last",
+):
+    """Computes the output shape of pooling operations.
+
+    Args:
+        input_shape: Input shape. Must be a tuple of integers.
+        pool_size: Size of the pooling operation. Must be a tuple of integers.
+        strides: Stride of the pooling operation. Must be a tuple of integers.
+            Defaults to `pool_size`.
+        padding: Padding method. Available methods are `"valid"` or `"same"`.
+            Defaults to `"valid"`.
+        data_format: String, either `"channels_last"` or `"channels_first"`.
+            The ordering of the dimensions in the inputs. `"channels_last"`
+            corresponds to inputs with shape `(batch, height, width, channels)`
+            while `"channels_first"` corresponds to inputs with shape
+            `(batch, channels, height, weight)`. Defaults to `"channels_last"`.
+
+    Returns:
+        Tuple of ints: The output shape of the pooling operation.
+
+    Examples:
+
+    # Basic usage with square pooling on a single image
+    >>> compute_pooling_output_shape((1, 4, 4, 1), (2, 2))
+    (1, 2, 2, 1)
+
+    # Strided pooling on a single image with strides different from pool_size
+    >>> compute_pooling_output_shape((1, 4, 4, 1), (2, 2), strides=(1, 1))
+    (1, 3, 3, 1)
+
+    # Pooling on a batch of images
+    >>> compute_pooling_output_shape((32, 4, 4, 3), (2, 2))
+    (32, 2, 2, 3)
+    """
+    strides = pool_size if strides is None else strides
+    input_shape_origin = list(input_shape)
+    input_shape = np.array(input_shape)
+    if data_format == "channels_last":
+        spatial_shape = input_shape[1:-1]
+    else:
+        spatial_shape = input_shape[2:]
+    none_dims = []
+    for i in range(len(spatial_shape)):
+        if spatial_shape[i] is None:
+            # Set `None` shape to a manual value so that we can run numpy
+            # computation on `spatial_shape`.
+            spatial_shape[i] = -1
+            none_dims.append(i)
+    pool_size = np.array(pool_size)
+    if padding == "valid":
+        output_spatial_shape = (
+            np.floor((spatial_shape - pool_size) / strides) + 1
+        )
+        for i in range(len(output_spatial_shape)):
+            if i not in none_dims and output_spatial_shape[i] < 0:
+                raise ValueError(
+                    "Computed output size would be negative. Received: "
+                    f"`inputs.shape={input_shape}` and `pool_size={pool_size}`."
+                )
+    elif padding == "same":
+        output_spatial_shape = np.floor((spatial_shape - 1) / strides) + 1
+    else:
+        raise ValueError(
+            "Argument `padding` must be either 'valid' or 'same'. Received: "
+            f"padding={padding}"
+        )
+    output_spatial_shape = [int(i) for i in output_spatial_shape]
+    for i in none_dims:
+        output_spatial_shape[i] = None
+    output_spatial_shape = tuple(output_spatial_shape)
+    if data_format == "channels_last":
+        output_shape = (
+            (input_shape_origin[0],)
+            + output_spatial_shape
+            + (input_shape_origin[-1],)
+        )
+    else:
+        output_shape = (
+            input_shape_origin[0],
+            input_shape_origin[1],
+        ) + output_spatial_shape
+    return output_shape
+
+
+def compute_conv_output_shape(
+    input_shape,
+    filters,
+    kernel_size,
+    strides=1,
+    padding="valid",
+    data_format="channels_last",
+    dilation_rate=1,
+):
+    """Compute the output shape of conv ops."""
+    if data_format == "channels_last":
+        spatial_shape = input_shape[1:-1]
+        kernel_shape = kernel_size + (input_shape[-1], filters)
+    else:
+        spatial_shape = input_shape[2:]
+        kernel_shape = kernel_size + (input_shape[1], filters)
+    if len(kernel_shape) != len(input_shape):
+        raise ValueError(
+            "Kernel shape must have the same length as input, but received "
+            f"kernel of shape {kernel_shape} and "
+            f"input of shape {input_shape}."
+        )
+    if isinstance(dilation_rate, int):
+        dilation_rate = (dilation_rate,) * len(spatial_shape)
+    if isinstance(strides, int):
+        strides = (strides,) * len(spatial_shape)
+    if len(dilation_rate) != len(spatial_shape):
+        raise ValueError(
+            "Dilation must be None, scalar or tuple/list of length of "
+            "inputs' spatial shape, but received "
+            f"`dilation_rate={dilation_rate}` and "
+            f"input of shape {input_shape}."
+        )
+    none_dims = []
+    spatial_shape = np.array(spatial_shape)
+    for i in range(len(spatial_shape)):
+        if spatial_shape[i] is None:
+            # Set `None` shape to a manual value so that we can run numpy
+            # computation on `spatial_shape`.
+            spatial_shape[i] = -1
+            none_dims.append(i)
+
+    kernel_spatial_shape = np.array(kernel_shape[:-2])
+    dilation_rate = np.array(dilation_rate)
+    if padding == "valid":
+        output_spatial_shape = (
+            np.floor(
+                (spatial_shape - dilation_rate * (kernel_spatial_shape - 1) - 1)
+                / strides
+            )
+            + 1
+        )
+        for i in range(len(output_spatial_shape)):
+            if i not in none_dims and output_spatial_shape[i] < 0:
+                raise ValueError(
+                    "Computed output size would be negative. Received "
+                    f"`inputs shape={input_shape}`, "
+                    f"`kernel shape={kernel_shape}`, "
+                    f"`dilation_rate={dilation_rate}`."
+                )
+    elif padding == "same" or padding == "causal":
+        output_spatial_shape = np.floor((spatial_shape - 1) / strides) + 1
+    else:
+        raise ValueError(
+            "`padding` must be either `'valid'` or `'same'`. Received "
+            f"{padding}."
+        )
+    output_spatial_shape = [int(i) for i in output_spatial_shape]
+    for i in none_dims:
+        output_spatial_shape[i] = None
+    output_spatial_shape = tuple(output_spatial_shape)
+    if data_format == "channels_last":
+        output_shape = (
+            (input_shape[0],) + output_spatial_shape + (kernel_shape[-1],)
+        )
+    else:
+        output_shape = (input_shape[0], kernel_shape[-1]) + output_spatial_shape
+    return output_shape
+
+
+def compute_matmul_output_shape(shape1, shape2):
+    """Compute the output shape of a `matmul` operation.
+
+    Args:
+        shape1: Shape of the left operand.
+        shape2: Shape of the right operand.
+
+    Returns:
+        Tuple of ints: The output shape for the `matmul` operation.
+    """
+    if len(shape1) == 1:
+        shape1 = (1, shape1[0])
+    if len(shape2) == 1:
+        shape2 = (shape2[0], 1)
+    if (
+        shape1[-1] is not None
+        and shape2[-2] is not None
+        and shape1[-1] != shape2[-2]
+    ):
+        raise ValueError(
+            "Inner dimensions (`x1.shape[-1]` and `x2.shape[-2]`) must be "
+            f"equal, but received `x1.shape={shape1}` and "
+            f"`x2.shape={shape2}`."
+        )
+
+    leading_shape = broadcast_shapes(shape1[:-2], shape2[:-2])
+    last_2_dims_shape = [shape1[-2], shape2[-1]]
+    output_shape = leading_shape + last_2_dims_shape
+    if len(shape1) == 1:
+        del output_shape[-2]
+    if len(shape2) == 1:
+        del output_shape[-1]
+    return tuple(output_shape)
+
+
+def compute_reshape_output_shape(input_shape, newshape, newshape_arg_name):
+    """Converts `-1` in `newshape` to either an actual dimension or `None`.
+
+    This utility does not special case the 0th dimension (batch size).
+    """
+    unknown_dim_count = newshape.count(-1)
+    if unknown_dim_count > 1:
+        raise ValueError(
+            "There must be at most one unknown dimension (-1) in "
+            f"{newshape_arg_name}. Received: {newshape_arg_name}={newshape}."
+        )
+
+    # If there is a None in input_shape, we can't infer what the -1 is
+    if None in input_shape:
+        return tuple(dim if dim != -1 else None for dim in newshape)
+
+    input_size = math.prod(input_shape)
+    # If the `newshape` is fully defined, return it
+    if unknown_dim_count == 0:
+        if input_size != math.prod(newshape):
+            raise ValueError(
+                "The total size of the tensor must be unchanged. Received: "
+                f"input_shape={input_shape}, {newshape_arg_name}={newshape}"
+            )
+        return newshape
+
+    # We have one -1 in `newshape`, compute the actual value
+    known_output_size = 1
+    unknown_dim_index = None
+    for index, dim in enumerate(newshape):
+        if dim == -1:
+            unknown_dim_index = index
+        else:
+            known_output_size *= dim
+
+    if known_output_size == 0 or input_size % known_output_size != 0:
+        raise ValueError(
+            "The total size of the tensor must be unchanged, however, the "
+            "input size cannot by divided by the specified dimensions in "
+            f"{newshape_arg_name}. Received: input_shape={input_shape}, "
+            f"{newshape_arg_name}={newshape}"
+        )
+
+    output_shape = list(newshape)
+    output_shape[unknown_dim_index] = input_size // known_output_size
+    return tuple(output_shape)
+
+
+def compute_transpose_output_shape(input_shape, axes):
+    """Compute the output shape for the `transpose` operation.
+
+    Args:
+        input_shape: Input shape.
+        axes: Permutation of the dimensions for the `transpose` operation.
+
+    Returns:
+        Tuple of ints: The output shape after the `transpose` operation.
+    """
+    input_shape = list(input_shape)
+    if axes is None:
+        return tuple(input_shape[::-1])
+
+    if len(axes) != len(input_shape):
+        raise ValueError(
+            "axis must be a list of the same length as the input shape, "
+            f"expected {len(input_shape)}, but received {len(axes)}."
+        )
+    return tuple(input_shape[ax] for ax in axes)
+
+
+def compute_take_along_axis_output_shape(input_shape, indices_shape, axis):
+    input_shape = list(input_shape)
+    indices_shape = list(indices_shape)
+    if axis is None:
+        input_shape = (
+            [None] if None in input_shape else [int(np.prod(input_shape))]
+        )
+
+    if len(input_shape) != len(indices_shape):
+        raise ValueError(
+            "`x` and `indices` must have the same number of dimensions, "
+            f"but receive shape {input_shape} and {indices_shape}."
+        )
+
+    input_shape[axis] = indices_shape[axis]
+    output_shape = broadcast_shapes(input_shape, indices_shape)
+    return output_shape
+
+
+def reduce_shape(shape, axis=None, keepdims=False):
+    shape = list(shape)
+    if axis is None:
+        if keepdims:
+            return tuple([1 for _ in shape])
+        else:
+            return tuple([])
+
+    if keepdims:
+        for ax in axis:
+            shape[ax] = 1
+        return tuple(shape)
+    else:
+        for ax in sorted(axis, reverse=True):
+            del shape[ax]
+        return tuple(shape)
+
+
+@keras_export("keras.utils.get_source_inputs")
+def get_source_inputs(tensor):
+    """Returns the list of input tensors necessary to compute `tensor`.
+
+    Output will always be a list of tensors
+    (potentially with 1 element).
+
+    Args:
+        tensor: The tensor to start from.
+
+    Returns:
+        List of input tensors.
+    """
+    if not hasattr(tensor, "_keras_history"):
+        return tensor
+
+    operation, node_index, _ = tensor._keras_history
+    if not operation or not operation._inbound_nodes:
+        return [tensor]
+    else:
+        node = operation._inbound_nodes[node_index]
+        if node.is_input:
+            # Reached input node, stop recursion.
+            return tree.flatten(node.output_tensors)
+        else:
+            source_tensors = []
+            for tensor in node.input_tensors:
+                previous_sources = get_source_inputs(tensor)
+                # Avoid input redundancy.
+                for x in previous_sources:
+                    if all(x is not t for t in source_tensors):
+                        source_tensors.append(x)
+            return source_tensors

SwarmUI/dlbackend/ComfyUI/venv/lib/python3.10/site-packages/keras/src/ops/symbolic_arguments.py

@@ -0,0 +1,46 @@
+from keras.src import tree
+from keras.src.backend import KerasTensor
+
+
+class SymbolicArguments:
+    def __init__(self, *args, **kwargs):
+        self.args = tree.map_structure(lambda x: x, args)
+        self.kwargs = tree.map_structure(lambda x: x, kwargs)
+        self._flat_arguments = tree.flatten((self.args, self.kwargs))
+
+        # Used to avoid expensive `tree` operations in the most common case.
+        if (
+            not self.kwargs
+            and len(self.args) == 1
+            and isinstance(self.args[0], KerasTensor)
+        ):
+            self._single_positional_tensor = self.args[0]
+        else:
+            self._single_positional_tensor = None
+
+        self.keras_tensors = []
+        for arg in self._flat_arguments:
+            if isinstance(arg, KerasTensor):
+                self.keras_tensors.append(arg)
+
+    def convert(self, conversion_fn):
+        args = tree.map_structure(conversion_fn, self.args)
+        kwargs = tree.map_structure(conversion_fn, self.kwargs)
+        return args, kwargs
+
+    def fill_in(self, tensor_dict):
+        """Maps KerasTensors to computed values using `tensor_dict`.
+
+        `tensor_dict` maps `KerasTensor` instances to their current values.
+        """
+        if self._single_positional_tensor is not None:
+            # Performance optimization for most common case.
+            # Approx. 70x faster.
+            return (tensor_dict[id(self._single_positional_tensor)],), {}
+
+        def switch_fn(x):
+            if isinstance(x, KerasTensor):
+                return tensor_dict.get(id(x), None)
+            return x
+
+        return self.convert(switch_fn)

SwarmUI/dlbackend/ComfyUI/venv/lib/python3.10/site-packages/keras/src/optimizers/__init__.py

@@ -0,0 +1,121 @@
+from keras.src.api_export import keras_export
+from keras.src.optimizers.adadelta import Adadelta
+from keras.src.optimizers.adafactor import Adafactor
+from keras.src.optimizers.adagrad import Adagrad
+from keras.src.optimizers.adam import Adam
+from keras.src.optimizers.adamax import Adamax
+from keras.src.optimizers.adamw import AdamW
+from keras.src.optimizers.ftrl import Ftrl
+from keras.src.optimizers.lion import Lion
+from keras.src.optimizers.loss_scale_optimizer import LossScaleOptimizer
+from keras.src.optimizers.nadam import Nadam
+from keras.src.optimizers.optimizer import Optimizer
+from keras.src.optimizers.rmsprop import RMSprop
+from keras.src.optimizers.sgd import SGD
+from keras.src.saving import serialization_lib
+
+ALL_OBJECTS = {
+    Optimizer,
+    Adam,
+    SGD,
+    RMSprop,
+    Adadelta,
+    AdamW,
+    Adagrad,
+    Adamax,
+    Adafactor,
+    Nadam,
+    Ftrl,
+    Lion,
+    LossScaleOptimizer,
+}
+ALL_OBJECTS_DICT = {cls.__name__.lower(): cls for cls in ALL_OBJECTS}
+
+
+@keras_export("keras.optimizers.serialize")
+def serialize(optimizer):
+    """Returns the optimizer configuration as a Python dict.
+
+    Args:
+        optimizer: An `Optimizer` instance to serialize.
+
+    Returns:
+        Python dict which contains the configuration of the optimizer.
+    """
+    return serialization_lib.serialize_keras_object(optimizer)
+
+
+@keras_export("keras.optimizers.deserialize")
+def deserialize(config, custom_objects=None):
+    """Returns a Keras optimizer object via its configuration.
+
+    Args:
+        config: Optimizer configuration dictionary.
+        custom_objects: Optional dictionary mapping names (strings) to custom
+            objects (classes and functions) to be considered during
+            deserialization.
+
+    Returns:
+        A Keras Optimizer instance.
+    """
+    # Make deserialization case-insensitive for built-in optimizers.
+    if config["class_name"].lower() in ALL_OBJECTS_DICT:
+        config["class_name"] = config["class_name"].lower()
+
+    return serialization_lib.deserialize_keras_object(
+        config,
+        module_objects=ALL_OBJECTS_DICT,
+        custom_objects=custom_objects,
+    )
+
+
+@keras_export("keras.optimizers.get")
+def get(identifier):
+    """Retrieves a Keras Optimizer instance.
+
+    Args:
+        identifier: Optimizer identifier, one of:
+            - String: name of an optimizer
+            - Dictionary: configuration dictionary.
+            - Keras Optimizer instance (it will be returned unchanged).
+
+    Returns:
+        A Keras Optimizer instance.
+    """
+    if identifier is None:
+        return None
+    elif isinstance(identifier, dict):
+        obj = deserialize(identifier)
+    elif isinstance(identifier, str):
+        config = {"class_name": identifier, "config": {}}
+        obj = deserialize(config)
+    else:
+        obj = identifier
+
+    if isinstance(obj, Optimizer):
+        return obj
+    raise ValueError(f"Could not interpret optimizer identifier: {identifier}")
+
+
+# We will add this temporarily so that tensorflow packages that depend on
+# estimators will continue to import (there are a large number). Note that
+# Keras 3 will not work with the estimators API.
+@keras_export(
+    [
+        "keras.optimizers.legacy.Adagrad",
+        "keras.optimizers.legacy.Adam",
+        "keras.optimizers.legacy.Ftrl",
+        "keras.optimizers.legacy.RMSprop",
+        "keras.optimizers.legacy.SGD",
+        "keras.optimizers.legacy.Optimizer",
+    ]
+)
+class LegacyOptimizerWarning:
+    def __init__(self, *args, **kwargs):
+        raise ImportError(
+            "`keras.optimizers.legacy` is not supported in Keras 3. When using "
+            "`tf.keras`, to continue using a `tf.keras.optimizers.legacy` "
+            "optimizer, you can install the `tf_keras` package (Keras 2) and "
+            "set the environment variable `TF_USE_LEGACY_KERAS=True` to "
+            "configure TensorFlow to use `tf_keras` when accessing `tf.keras`."
+        )

SwarmUI/dlbackend/ComfyUI/venv/lib/python3.10/site-packages/keras/src/optimizers/adadelta.py

@@ -0,0 +1,139 @@
+from keras.src import ops
+from keras.src.api_export import keras_export
+from keras.src.optimizers import optimizer
+
+
+@keras_export(["keras.optimizers.Adadelta"])
+class Adadelta(optimizer.Optimizer):
+    """Optimizer that implements the Adadelta algorithm.
+
+    Adadelta optimization is a stochastic gradient descent method that is based
+    on adaptive learning rate per dimension to address two drawbacks:
+
+    - The continual decay of learning rates throughout training.
+    - The need for a manually selected global learning rate.
+
+    Adadelta is a more robust extension of Adagrad that adapts learning rates
+    based on a moving window of gradient updates, instead of accumulating all
+    past gradients. This way, Adadelta continues learning even when many updates
+    have been done. Compared to Adagrad, in the original version of Adadelta you
+    don't have to set an initial learning rate. In this version, the initial
+    learning rate can be set, as in most other Keras optimizers.
+
+    Args:
+        learning_rate: A float, a
+            `keras.optimizers.schedules.LearningRateSchedule` instance, or
+            a callable that takes no arguments and returns the actual value to
+            use. The learning rate. Defaults to `0.001`. Note that `Adadelta`
+            tends to benefit from higher initial learning rate values compared
+            to other optimizers. To match the exact form in the original paper,
+            use 1.0.
+        rho: A floating point value. The decay rate. Defaults to `0.95`.
+        epsilon: Small floating point value for maintaining numerical stability.
+        {{base_optimizer_keyword_args}}
+
+    Reference:
+
+    - [Zeiler, 2012](http://arxiv.org/abs/1212.5701)
+    """
+
+    def __init__(
+        self,
+        learning_rate=0.001,
+        rho=0.95,
+        epsilon=1e-7,
+        weight_decay=None,
+        clipnorm=None,
+        clipvalue=None,
+        global_clipnorm=None,
+        use_ema=False,
+        ema_momentum=0.99,
+        ema_overwrite_frequency=None,
+        loss_scale_factor=None,
+        gradient_accumulation_steps=None,
+        name="adadelta",
+        **kwargs,
+    ):
+        super().__init__(
+            learning_rate=learning_rate,
+            weight_decay=weight_decay,
+            clipnorm=clipnorm,
+            clipvalue=clipvalue,
+            global_clipnorm=global_clipnorm,
+            use_ema=use_ema,
+            ema_momentum=ema_momentum,
+            ema_overwrite_frequency=ema_overwrite_frequency,
+            loss_scale_factor=loss_scale_factor,
+            gradient_accumulation_steps=gradient_accumulation_steps,
+            name=name,
+            **kwargs,
+        )
+        self.rho = rho
+        self.epsilon = epsilon
+
+    def build(self, var_list):
+        if self.built:
+            return
+        super().build(var_list)
+        self._accumulated_grads = []
+        self._accumulated_delta_vars = []
+        for var in var_list:
+            self._accumulated_grads.append(
+                self.add_variable_from_reference(var, "accumulated_grad")
+            )
+            self._accumulated_delta_vars.append(
+                self.add_variable_from_reference(var, "accumulated_delta_var")
+            )
+
+    def update_step(self, grad, variable, learning_rate):
+        """Update step given gradient and the associated model variable."""
+        lr = ops.cast(learning_rate, variable.dtype)
+        grad = ops.cast(grad, variable.dtype)
+
+        rho = self.rho
+        accumulated_grad = self._accumulated_grads[
+            self._get_variable_index(variable)
+        ]
+        accumulated_delta_var = self._accumulated_delta_vars[
+            self._get_variable_index(variable)
+        ]
+
+        def rms(x):
+            return ops.sqrt(ops.add(x, self.epsilon))
+
+        self.assign(
+            accumulated_grad,
+            ops.add(
+                rho * accumulated_grad, ops.multiply(1 - rho, ops.square(grad))
+            ),
+        )
+        delta_var = ops.negative(
+            ops.divide(
+                ops.multiply(rms(accumulated_delta_var), grad),
+                rms(accumulated_grad),
+            )
+        )
+        self.assign(
+            accumulated_delta_var,
+            ops.add(
+                ops.multiply(rho, accumulated_delta_var),
+                ops.multiply(1 - rho, ops.square(delta_var)),
+            ),
+        )
+        self.assign_add(variable, ops.multiply(lr, delta_var))
+
+    def get_config(self):
+        config = super().get_config()
+
+        config.update(
+            {
+                "rho": self.rho,
+                "epsilon": self.epsilon,
+            }
+        )
+        return config
+
+
+Adadelta.__doc__ = Adadelta.__doc__.replace(
+    "{{base_optimizer_keyword_args}}", optimizer.base_optimizer_keyword_args
+)

SwarmUI/dlbackend/ComfyUI/venv/lib/python3.10/site-packages/keras/src/optimizers/adafactor.py

@@ -0,0 +1,208 @@
+from keras.src import backend
+from keras.src import ops
+from keras.src.api_export import keras_export
+from keras.src.optimizers import optimizer
+
+
+@keras_export(["keras.optimizers.Adafactor"])
+class Adafactor(optimizer.Optimizer):
+    """Optimizer that implements the Adafactor algorithm.
+
+    Adafactor is commonly used in NLP tasks, and has the advantage
+    of taking less memory because it only saves partial information of previous
+    gradients.
+
+    The default argument setup is based on the original paper (see reference).
+    When gradients are of dimension > 2, Adafactor optimizer will delete the
+    last 2 dimensions separately in its accumulator variables.
+
+    Args:
+        learning_rate: A float, a
+            `keras.optimizers.schedules.LearningRateSchedule` instance, or
+            a callable that takes no arguments and returns the actual value to
+            use. The learning rate. Defaults to `0.001`.
+        beta_2_decay: float, defaults to -0.8. The decay rate of `beta_2`.
+        epsilon_1: float, defaults to 1e-30. A small offset to keep denominator
+            away from 0.
+        epsilon_2: float, defaults to 1e-3. A small offset to avoid learning
+            rate becoming too small by time.
+        clip_threshold: float, defaults to 1.0. Clipping threshold. This is a
+            part of Adafactor algorithm, independent from `clipnorm`,
+            `clipvalue`, and `global_clipnorm`.
+        relative_step: bool, defaults to `True`. If `learning_rate` is a
+            constant and `relative_step=True`, learning rate will be adjusted
+            based on current iterations. This is a default learning rate decay
+            in Adafactor.
+        {{base_optimizer_keyword_args}}
+
+    Reference:
+
+    - [Shazeer, Noam et al., 2018](https://arxiv.org/abs/1804.04235).
+
+    """
+
+    def __init__(
+        self,
+        learning_rate=0.001,
+        beta_2_decay=-0.8,
+        epsilon_1=1e-30,
+        epsilon_2=1e-3,
+        clip_threshold=1.0,
+        relative_step=True,
+        weight_decay=None,
+        clipnorm=None,
+        clipvalue=None,
+        global_clipnorm=None,
+        use_ema=False,
+        ema_momentum=0.99,
+        ema_overwrite_frequency=None,
+        loss_scale_factor=None,
+        gradient_accumulation_steps=None,
+        name="adafactor",
+        **kwargs,
+    ):
+        super().__init__(
+            learning_rate=learning_rate,
+            name=name,
+            weight_decay=weight_decay,
+            clipnorm=clipnorm,
+            clipvalue=clipvalue,
+            global_clipnorm=global_clipnorm,
+            use_ema=use_ema,
+            ema_momentum=ema_momentum,
+            ema_overwrite_frequency=ema_overwrite_frequency,
+            loss_scale_factor=loss_scale_factor,
+            gradient_accumulation_steps=gradient_accumulation_steps,
+            **kwargs,
+        )
+        self.beta_2_decay = beta_2_decay
+        self.epsilon_1 = epsilon_1
+        self.epsilon_2 = epsilon_2
+        self.clip_threshold = clip_threshold
+        self.relative_step = relative_step
+
+    def build(self, var_list):
+        """Initialize optimizer variables.
+
+        Adam optimizer has 3 types of variables: momentums, velocities and
+        velocity_hat (only set when amsgrad is applied),
+
+        Args:
+            var_list: list of model variables to build Adam variables on.
+        """
+        if self.built:
+            return
+        super().build(var_list)
+        self._r = []
+        self._c = []
+        self._v = []
+        for var in var_list:
+            if len(var.shape) < 2:
+                # Don't factor if variable is of dimension < 2, but we still
+                # need to create dummy variables as placeholder.
+                with backend.name_scope(self.name, caller=self):
+                    self._r.append(
+                        backend.Variable(0, name=var.name, trainable=False)
+                    )
+                    self._c.append(
+                        backend.Variable(0, name=var.name, trainable=False)
+                    )
+            else:
+                # Always factor the last 2 dimensions.
+                r_shape = var.shape[:-1]
+                c_shape = var.shape[:-2] + (var.shape[-1],)
+                self._r.append(
+                    self.add_variable(
+                        shape=r_shape,
+                        dtype=var.dtype,
+                        name=var.name,
+                    )
+                )
+                self._c.append(
+                    self.add_variable(
+                        shape=c_shape,
+                        dtype=var.dtype,
+                        name=var.name,
+                    )
+                )
+            self._v.append(
+                self.add_variable_from_reference(
+                    reference_variable=var, name="velocity"
+                )
+            )
+
+    def _rms(self, x):
+        return ops.sqrt(ops.mean(ops.square(x)))
+
+    def update_step(self, gradient, variable, learning_rate):
+        """Update step given gradient and the associated model variable."""
+
+        lr = ops.cast(learning_rate, variable.dtype)
+        gradient = ops.cast(gradient, variable.dtype)
+        epsilon_2 = ops.cast(self.epsilon_2, variable.dtype)
+        one = ops.cast(1.0, variable.dtype)
+        local_step = ops.cast(self.iterations + 1, variable.dtype)
+        if not callable(self._learning_rate) and self.relative_step:
+            lr = ops.minimum(lr, 1 / ops.sqrt(local_step))
+
+        r = self._r[self._get_variable_index(variable)]
+        c = self._c[self._get_variable_index(variable)]
+        v = self._v[self._get_variable_index(variable)]
+
+        rho_t = ops.minimum(lr, 1 / ops.sqrt(local_step))
+        alpha_t = ops.maximum(epsilon_2, self._rms(variable)) * rho_t
+        regulated_grad_square = ops.add(ops.square(gradient), self.epsilon_1)
+        beta_2_t = 1 - ops.power(local_step, self.beta_2_decay)
+
+        if len(variable.shape) >= 2:
+            # `r` deletes the last dimension of gradient, so it is of shape
+            # `gradient.shape[:-1]`.
+            self.assign(
+                r,
+                beta_2_t * r
+                + (1 - beta_2_t) * ops.mean(regulated_grad_square, axis=-1),
+            )
+            # `c` deletes the second last dimension of gradient, so it is of
+            # shape `gradient.shape[:-2] + gradient.shape[-1]`.
+            self.assign(
+                c,
+                beta_2_t * c
+                + (1 - beta_2_t) * ops.mean(regulated_grad_square, axis=-2),
+            )
+            self.assign(
+                v,
+                ops.expand_dims(
+                    r / ops.mean(r, axis=-1, keepdims=True), axis=-1
+                )
+                * ops.expand_dims(c, -2),
+            )
+        else:
+            self.assign(
+                v, beta_2_t * v + (1 - beta_2_t) * regulated_grad_square
+            )
+
+        u_t = ops.divide(gradient, ops.sqrt(v))
+        u_t_hat = ops.divide(
+            u_t,
+            ops.maximum(one, ops.divide(self._rms(u_t), self.clip_threshold)),
+        )
+        self.assign_sub(variable, ops.multiply(alpha_t, u_t_hat))
+
+    def get_config(self):
+        config = super().get_config()
+
+        config.update(
+            {
+                "beta_2_decay": self.beta_2_decay,
+                "epsilon_1": self.epsilon_1,
+                "epsilon_2": self.epsilon_2,
+                "clip_threshold": self.clip_threshold,
+                "relative_step": self.relative_step,
+            }
+        )
+        return config
+
+
+Adafactor.__doc__ = Adafactor.__doc__.replace(
+    "{{base_optimizer_keyword_args}}", optimizer.base_optimizer_keyword_args
+)

SwarmUI/dlbackend/ComfyUI/venv/lib/python3.10/site-packages/keras/src/optimizers/adagrad.py

@@ -0,0 +1,115 @@
+from keras.src import initializers
+from keras.src import ops
+from keras.src.api_export import keras_export
+from keras.src.optimizers import optimizer
+
+
+@keras_export(["keras.optimizers.Adagrad"])
+class Adagrad(optimizer.Optimizer):
+    """Optimizer that implements the Adagrad algorithm.
+
+    Adagrad is an optimizer with parameter-specific learning rates,
+    which are adapted relative to how frequently a parameter gets
+    updated during training. The more updates a parameter receives,
+    the smaller the updates.
+
+    Args:
+        learning_rate: A float, a
+            `keras.optimizers.schedules.LearningRateSchedule` instance, or
+            a callable that takes no arguments and returns the actual value to
+            use. The learning rate. Defaults to `0.001`. Note that `Adagrad`
+            tends to benefit from higher initial learning rate values compared
+            to other optimizers. To match the exact form in the original paper,
+            use `1.0`.
+        initial_accumulator_value: Floating point value. Starting value for the
+            accumulators (per-parameter momentum values). Must be non-negative.
+        epsilon: Small floating point value for maintaining numerical stability.
+        {{base_optimizer_keyword_args}}
+
+    Reference:
+
+    - [Duchi et al., 2011](
+        http://www.jmlr.org/papers/volume12/duchi11a/duchi11a.pdf).
+    """
+
+    def __init__(
+        self,
+        learning_rate=0.001,
+        initial_accumulator_value=0.1,
+        epsilon=1e-7,
+        weight_decay=None,
+        clipnorm=None,
+        clipvalue=None,
+        global_clipnorm=None,
+        use_ema=False,
+        ema_momentum=0.99,
+        ema_overwrite_frequency=None,
+        loss_scale_factor=None,
+        gradient_accumulation_steps=None,
+        name="adagrad",
+        **kwargs,
+    ):
+        super().__init__(
+            learning_rate=learning_rate,
+            weight_decay=weight_decay,
+            clipnorm=clipnorm,
+            clipvalue=clipvalue,
+            global_clipnorm=global_clipnorm,
+            use_ema=use_ema,
+            ema_momentum=ema_momentum,
+            ema_overwrite_frequency=ema_overwrite_frequency,
+            loss_scale_factor=loss_scale_factor,
+            gradient_accumulation_steps=gradient_accumulation_steps,
+            name=name,
+            **kwargs,
+        )
+        self.initial_accumulator_value = initial_accumulator_value
+        self.epsilon = epsilon
+
+    def build(self, var_list):
+        if self.built:
+            return
+        super().build(var_list)
+        self._accumulators = []
+        initializer = initializers.Constant(self.initial_accumulator_value)
+        for var in var_list:
+            self._accumulators.append(
+                self.add_variable(
+                    shape=var.shape,
+                    initializer=initializer,
+                    dtype=var.dtype,
+                    name="accumulator",
+                )
+            )
+
+    def update_step(self, gradient, variable, learning_rate):
+        """Update step given gradient and the associated model variable."""
+        lr = ops.cast(learning_rate, variable.dtype)
+        gradient = ops.cast(gradient, variable.dtype)
+
+        accumulator = self._accumulators[self._get_variable_index(variable)]
+
+        self.assign_add(accumulator, ops.square(gradient))
+        self.assign_sub(
+            variable,
+            ops.divide(
+                ops.multiply(lr, gradient),
+                ops.sqrt(ops.add(accumulator, self.epsilon)),
+            ),
+        )
+
+    def get_config(self):
+        config = super().get_config()
+
+        config.update(
+            {
+                "initial_accumulator_value": self.initial_accumulator_value,
+                "epsilon": self.epsilon,
+            }
+        )
+        return config
+
+
+Adagrad.__doc__ = Adagrad.__doc__.replace(
+    "{{base_optimizer_keyword_args}}", optimizer.base_optimizer_keyword_args
+)

SwarmUI/dlbackend/ComfyUI/venv/lib/python3.10/site-packages/keras/src/optimizers/adam.py

@@ -0,0 +1,167 @@
+from keras.src import ops
+from keras.src.api_export import keras_export
+from keras.src.optimizers import optimizer
+
+
+@keras_export(["keras.optimizers.Adam"])
+class Adam(optimizer.Optimizer):
+    """Optimizer that implements the Adam algorithm.
+
+    Adam optimization is a stochastic gradient descent method that is based on
+    adaptive estimation of first-order and second-order moments.
+
+    According to
+    [Kingma et al., 2014](http://arxiv.org/abs/1412.6980),
+    the method is "*computationally
+    efficient, has little memory requirement, invariant to diagonal rescaling of
+    gradients, and is well suited for problems that are large in terms of
+    data/parameters*".
+
+    Args:
+        learning_rate: A float, a
+            `keras.optimizers.schedules.LearningRateSchedule` instance, or
+            a callable that takes no arguments and returns the actual value to
+            use. The learning rate. Defaults to `0.001`.
+        beta_1: A float value or a constant float tensor, or a callable
+            that takes no arguments and returns the actual value to use. The
+            exponential decay rate for the 1st moment estimates. Defaults to
+            `0.9`.
+        beta_2: A float value or a constant float tensor, or a callable
+            that takes no arguments and returns the actual value to use. The
+            exponential decay rate for the 2nd moment estimates. Defaults to
+            `0.999`.
+        epsilon: A small constant for numerical stability. This epsilon is
+            "epsilon hat" in the Kingma and Ba paper (in the formula just before
+            Section 2.1), not the epsilon in Algorithm 1 of the paper. Defaults
+            to `1e-7`.
+        amsgrad: Boolean. Whether to apply AMSGrad variant of this algorithm
+            from the paper "On the Convergence of Adam and beyond". Defaults
+            to `False`.
+        {{base_optimizer_keyword_args}}
+    """
+
+    def __init__(
+        self,
+        learning_rate=0.001,
+        beta_1=0.9,
+        beta_2=0.999,
+        epsilon=1e-7,
+        amsgrad=False,
+        weight_decay=None,
+        clipnorm=None,
+        clipvalue=None,
+        global_clipnorm=None,
+        use_ema=False,
+        ema_momentum=0.99,
+        ema_overwrite_frequency=None,
+        loss_scale_factor=None,
+        gradient_accumulation_steps=None,
+        name="adam",
+        **kwargs,
+    ):
+        super().__init__(
+            learning_rate=learning_rate,
+            name=name,
+            weight_decay=weight_decay,
+            clipnorm=clipnorm,
+            clipvalue=clipvalue,
+            global_clipnorm=global_clipnorm,
+            use_ema=use_ema,
+            ema_momentum=ema_momentum,
+            ema_overwrite_frequency=ema_overwrite_frequency,
+            loss_scale_factor=loss_scale_factor,
+            gradient_accumulation_steps=gradient_accumulation_steps,
+            **kwargs,
+        )
+        self.beta_1 = beta_1
+        self.beta_2 = beta_2
+        self.epsilon = epsilon
+        self.amsgrad = amsgrad
+
+    def build(self, var_list):
+        """Initialize optimizer variables.
+
+        Adam optimizer has 3 types of variables: momentums, velocities and
+        velocity_hat (only set when amsgrad is applied),
+
+        Args:
+            var_list: list of model variables to build Adam variables on.
+        """
+        if self.built:
+            return
+        super().build(var_list)
+        self._momentums = []
+        self._velocities = []
+        for var in var_list:
+            self._momentums.append(
+                self.add_variable_from_reference(
+                    reference_variable=var, name="momentum"
+                )
+            )
+            self._velocities.append(
+                self.add_variable_from_reference(
+                    reference_variable=var, name="velocity"
+                )
+            )
+        if self.amsgrad:
+            self._velocity_hats = []
+            for var in var_list:
+                self._velocity_hats.append(
+                    self.add_variable_from_reference(
+                        reference_variable=var, name="velocity_hat"
+                    )
+                )
+
+    def update_step(self, gradient, variable, learning_rate):
+        """Update step given gradient and the associated model variable."""
+        lr = ops.cast(learning_rate, variable.dtype)
+        gradient = ops.cast(gradient, variable.dtype)
+        local_step = ops.cast(self.iterations + 1, variable.dtype)
+        beta_1_power = ops.power(
+            ops.cast(self.beta_1, variable.dtype), local_step
+        )
+        beta_2_power = ops.power(
+            ops.cast(self.beta_2, variable.dtype), local_step
+        )
+
+        m = self._momentums[self._get_variable_index(variable)]
+        v = self._velocities[self._get_variable_index(variable)]
+
+        alpha = lr * ops.sqrt(1 - beta_2_power) / (1 - beta_1_power)
+
+        self.assign_add(
+            m, ops.multiply(ops.subtract(gradient, m), 1 - self.beta_1)
+        )
+        self.assign_add(
+            v,
+            ops.multiply(
+                ops.subtract(ops.square(gradient), v), 1 - self.beta_2
+            ),
+        )
+        if self.amsgrad:
+            v_hat = self._velocity_hats[self._get_variable_index(variable)]
+            self.assign(v_hat, ops.maximum(v_hat, v))
+            v = v_hat
+        self.assign_sub(
+            variable,
+            ops.divide(
+                ops.multiply(m, alpha), ops.add(ops.sqrt(v), self.epsilon)
+            ),
+        )
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "beta_1": self.beta_1,
+                "beta_2": self.beta_2,
+                "epsilon": self.epsilon,
+                "amsgrad": self.amsgrad,
+            }
+        )
+        return config
+
+
+Adam.__doc__ = Adam.__doc__.replace(
+    "{{base_optimizer_keyword_args}}", optimizer.base_optimizer_keyword_args
+)

SwarmUI/dlbackend/ComfyUI/venv/lib/python3.10/site-packages/keras/src/optimizers/adamax.py

@@ -0,0 +1,156 @@
+from keras.src import ops
+from keras.src.api_export import keras_export
+from keras.src.optimizers import optimizer
+
+
+@keras_export(["keras.optimizers.Adamax"])
+class Adamax(optimizer.Optimizer):
+    """Optimizer that implements the Adamax algorithm.
+
+    Adamax, a variant of Adam based on the infinity norm, is a first-order
+    gradient-based optimization method. Due to its capability of adjusting the
+    learning rate based on data characteristics, it is suited to learn
+    time-variant process, e.g., speech data with dynamically changed noise
+    conditions. Default parameters follow those provided in the paper (see
+    references below).
+
+    Initialization:
+
+    ```python
+    m = 0  # Initialize initial 1st moment vector
+    u = 0  # Initialize the exponentially weighted infinity norm
+    t = 0  # Initialize timestep
+    ```
+
+    The update rule for parameter `w` with gradient `g` is described at the end
+    of section 7.1 of the paper (see the reference section):
+
+    ```python
+    t += 1
+    m = beta1 * m + (1 - beta) * g
+    u = max(beta2 * u, abs(g))
+    current_lr = learning_rate / (1 - beta1 ** t)
+    w = w - current_lr * m / (u + epsilon)
+    ```
+
+    Args:
+        learning_rate: A float, a
+            `keras.optimizers.schedules.LearningRateSchedule` instance, or
+            a callable that takes no arguments and returns the actual value to
+            use. The learning rate. Defaults to `0.001`.
+        beta_1: A float value or a constant float tensor. The exponential decay
+            rate for the 1st moment estimates.
+        beta_2: A float value or a constant float tensor. The exponential decay
+            rate for the exponentially weighted infinity norm.
+        epsilon: A small constant for numerical stability.
+            {{base_optimizer_keyword_args}}
+
+    Reference:
+
+    - [Kingma et al., 2014](http://arxiv.org/abs/1412.6980)
+    """
+
+    def __init__(
+        self,
+        learning_rate=0.001,
+        beta_1=0.9,
+        beta_2=0.999,
+        epsilon=1e-7,
+        weight_decay=None,
+        clipnorm=None,
+        clipvalue=None,
+        global_clipnorm=None,
+        use_ema=False,
+        ema_momentum=0.99,
+        ema_overwrite_frequency=None,
+        loss_scale_factor=None,
+        gradient_accumulation_steps=None,
+        name="adamax",
+        **kwargs,
+    ):
+        super().__init__(
+            learning_rate=learning_rate,
+            name=name,
+            weight_decay=weight_decay,
+            clipnorm=clipnorm,
+            clipvalue=clipvalue,
+            global_clipnorm=global_clipnorm,
+            use_ema=use_ema,
+            ema_momentum=ema_momentum,
+            ema_overwrite_frequency=ema_overwrite_frequency,
+            loss_scale_factor=loss_scale_factor,
+            gradient_accumulation_steps=gradient_accumulation_steps,
+            **kwargs,
+        )
+        self.beta_1 = beta_1
+        self.beta_2 = beta_2
+        self.epsilon = epsilon
+
+    def build(self, var_list):
+        """Initialize optimizer variables.
+
+        Adamax optimizer has 2 types of variables: momentums (denoted as m),
+        exponentially weighted infinity norm (denoted as u).
+
+        Args:
+            var_list: list of model variables to build Adamax variables on.
+        """
+        if self.built:
+            return
+        super().build(var_list)
+        self._m = []
+        self._u = []
+        for var in var_list:
+            self._m.append(
+                self.add_variable_from_reference(
+                    reference_variable=var, name="momentum"
+                )
+            )
+            self._u.append(
+                self.add_variable_from_reference(
+                    reference_variable=var, name="norm"
+                )
+            )
+
+    def update_step(self, gradient, variable, learning_rate):
+        """Update step given gradient and the associated model variable."""
+        lr = ops.cast(learning_rate, variable.dtype)
+        gradient = ops.cast(gradient, variable.dtype)
+        local_step = ops.cast(self.iterations + 1, variable.dtype)
+        beta_1_power = ops.power(
+            ops.cast(self.beta_1, variable.dtype), local_step
+        )
+
+        m = self._m[self._get_variable_index(variable)]
+        u = self._u[self._get_variable_index(variable)]
+
+        self.assign_add(
+            m, ops.multiply(ops.subtract(gradient, m), (1 - self.beta_1))
+        )
+        self.assign(
+            u, ops.maximum(ops.multiply(self.beta_2, u), ops.abs(gradient))
+        )
+        self.assign_sub(
+            variable,
+            ops.divide(
+                ops.multiply(lr, m),
+                ops.multiply((1 - beta_1_power), ops.add(u, self.epsilon)),
+            ),
+        )
+
+    def get_config(self):
+        config = super().get_config()
+
+        config.update(
+            {
+                "beta_1": self.beta_1,
+                "beta_2": self.beta_2,
+                "epsilon": self.epsilon,
+            }
+        )
+        return config
+
+
+Adamax.__doc__ = Adamax.__doc__.replace(
+    "{{base_optimizer_keyword_args}}", optimizer.base_optimizer_keyword_args
+)

SwarmUI/dlbackend/ComfyUI/venv/lib/python3.10/site-packages/keras/src/optimizers/adamw.py

@@ -0,0 +1,100 @@
+from keras.src.api_export import keras_export
+from keras.src.optimizers import adam
+from keras.src.optimizers import optimizer
+
+
+@keras_export(["keras.optimizers.AdamW"])
+class AdamW(adam.Adam):
+    """Optimizer that implements the AdamW algorithm.
+
+    AdamW optimization is a stochastic gradient descent method that is based on
+    adaptive estimation of first-order and second-order moments with an added
+    method to decay weights per the techniques discussed in the paper,
+    'Decoupled Weight Decay Regularization' by
+    [Loshchilov, Hutter et al., 2019](https://arxiv.org/abs/1711.05101).
+
+    According to
+    [Kingma et al., 2014](http://arxiv.org/abs/1412.6980),
+    the underlying Adam method is "*computationally
+    efficient, has little memory requirement, invariant to diagonal rescaling of
+    gradients, and is well suited for problems that are large in terms of
+    data/parameters*".
+
+    Args:
+        learning_rate: A float, a
+            `keras.optimizers.schedules.LearningRateSchedule` instance, or
+            a callable that takes no arguments and returns the actual value to
+            use. The learning rate. Defaults to `0.001`.
+        beta_1: A float value or a constant float tensor, or a callable
+            that takes no arguments and returns the actual value to use. The
+            exponential decay rate for the 1st moment estimates.
+            Defaults to `0.9`.
+        beta_2: A float value or a constant float tensor, or a callable
+            that takes no arguments and returns the actual value to use. The
+            exponential decay rate for the 2nd moment estimates.
+            Defaults to `0.999`.
+        epsilon: A small constant for numerical stability. This epsilon is
+            "epsilon hat" in the Kingma and Ba paper (in the formula just
+            before Section 2.1), not the epsilon in Algorithm 1 of the paper.
+            Defaults to 1e-7.
+        amsgrad: Boolean. Whether to apply AMSGrad variant of this algorithm
+            from the paper "On the Convergence of Adam and beyond".
+            Defaults to `False`.
+        {{base_optimizer_keyword_args}}
+
+    References:
+
+    - [Loshchilov et al., 2019](https://arxiv.org/abs/1711.05101)
+    - [Kingma et al., 2014](http://arxiv.org/abs/1412.6980) for `adam`
+    - [Reddi et al., 2018](
+        https://openreview.net/pdf?id=ryQu7f-RZ) for `amsgrad`.
+    """
+
+    def __init__(
+        self,
+        learning_rate=0.001,
+        weight_decay=0.004,
+        beta_1=0.9,
+        beta_2=0.999,
+        epsilon=1e-7,
+        amsgrad=False,
+        clipnorm=None,
+        clipvalue=None,
+        global_clipnorm=None,
+        use_ema=False,
+        ema_momentum=0.99,
+        ema_overwrite_frequency=None,
+        loss_scale_factor=None,
+        gradient_accumulation_steps=None,
+        name="adamw",
+        **kwargs,
+    ):
+        super().__init__(
+            learning_rate=learning_rate,
+            beta_1=beta_1,
+            beta_2=beta_2,
+            epsilon=epsilon,
+            amsgrad=amsgrad,
+            name=name,
+            weight_decay=weight_decay,
+            clipnorm=clipnorm,
+            clipvalue=clipvalue,
+            global_clipnorm=global_clipnorm,
+            use_ema=use_ema,
+            ema_momentum=ema_momentum,
+            ema_overwrite_frequency=ema_overwrite_frequency,
+            loss_scale_factor=loss_scale_factor,
+            gradient_accumulation_steps=gradient_accumulation_steps,
+            **kwargs,
+        )
+
+        if self.weight_decay is None:
+            raise ValueError(
+                "Argument `weight_decay` must be a float. Received: "
+                "weight_decay=None"
+            )
+
+
+AdamW.__doc__ = AdamW.__doc__.replace(
+    "{{base_optimizer_keyword_args}}", optimizer.base_optimizer_keyword_args
+)

SwarmUI/dlbackend/ComfyUI/venv/lib/python3.10/site-packages/keras/src/optimizers/base_optimizer.py

@@ -0,0 +1,1102 @@
+import re
+import warnings
+
+from keras.src import backend
+from keras.src import initializers
+from keras.src import ops
+from keras.src.optimizers.schedules import learning_rate_schedule
+from keras.src.saving import serialization_lib
+from keras.src.saving.keras_saveable import KerasSaveable
+from keras.src.utils import tracking
+from keras.src.utils.naming import auto_name
+
+
+class BaseOptimizer(KerasSaveable):
+    """Abstract optimizer base class.
+
+    If you intend to create your own optimization algorithm, please inherit from
+    this class and override the following methods:
+
+    - `build`: Create your optimizer-related variables, such as momentum
+        variables in the SGD optimizer.
+    - `update_step`: Implement your optimizer's variable updating logic.
+    - `get_config`: serialization of the optimizer.
+
+    Example:
+
+    ```python
+    class SGD(Optimizer):
+        def __init__(self, **kwargs):
+            super().__init__(**kwargs)
+            self.momentum = 0.9
+
+        def build(self, variables):
+            super().build(variables)
+            self.momentums = []
+            for variable in variables:
+                self.momentums.append(
+                    self.add_variable_from_reference(
+                        reference_variable=variable, name="momentum"
+                    )
+                )
+
+        def update_step(self, gradient, variable, learning_rate):
+            learning_rate = ops.cast(learning_rate, variable.dtype)
+            gradient = ops.cast(gradient, variable.dtype)
+            m = self.momentums[self._get_variable_index(variable)]
+            self.assign(
+                m,
+                ops.subtract(
+                    ops.multiply(m, ops.cast(self.momentum, variable.dtype)),
+                    ops.multiply(gradient, learning_rate),
+                ),
+            )
+            self.assign_add(variable, m)
+
+        def get_config(self):
+            config = super().get_config()
+            config.update(
+                {
+                    "momentum": self.momentum,
+                    "nesterov": self.nesterov,
+                }
+            )
+            return config
+    ```
+    """
+
+    def __init__(
+        self,
+        learning_rate,
+        weight_decay=None,
+        clipnorm=None,
+        clipvalue=None,
+        global_clipnorm=None,
+        use_ema=False,
+        ema_momentum=0.99,
+        ema_overwrite_frequency=None,
+        loss_scale_factor=None,
+        gradient_accumulation_steps=None,
+        name=None,
+        **kwargs,
+    ):
+        self._lock = False
+
+        if kwargs.pop("decay", None) is not None:
+            warnings.warn(
+                "Argument `decay` is no longer supported and will be ignored."
+            )
+        if kwargs:
+            raise ValueError(f"Argument(s) not recognized: {kwargs}")
+
+        if name is None:
+            name = auto_name(self.__class__.__name__)
+        self.name = name
+        self.weight_decay = weight_decay
+        self.clipnorm = clipnorm
+        self.global_clipnorm = global_clipnorm
+        self.clipvalue = clipvalue
+        self.use_ema = use_ema
+        self.loss_scale_factor = loss_scale_factor
+        self.gradient_accumulation_steps = gradient_accumulation_steps
+
+        if gradient_accumulation_steps:
+            if not gradient_accumulation_steps >= 2:
+                raise ValueError(
+                    "`gradient_accumulation_steps` must be an integer >= 2. "
+                    "Received: gradient_accumulation_steps="
+                    f"{gradient_accumulation_steps}"
+                )
+
+        if use_ema:
+            # Verify the arguments related to EMA.
+            if ema_momentum > 1 or ema_momentum < 0:
+                raise ValueError(
+                    "`ema_momentum` must be in the range [0, 1]. "
+                    f"Received: ema_momentum={ema_momentum}"
+                )
+            if ema_overwrite_frequency and (
+                not isinstance(ema_overwrite_frequency, int)
+                or ema_overwrite_frequency < 1
+            ):
+                raise ValueError(
+                    "`ema_overwrite_frequency` must be an integer >= 1 or "
+                    "None. Received: ema_overwrite_frequency="
+                    f"{ema_overwrite_frequency}"
+                )
+        self.ema_momentum = ema_momentum
+        self.ema_overwrite_frequency = ema_overwrite_frequency
+
+        clip_args_sum = sum(
+            a is not None for a in [clipnorm, clipvalue, global_clipnorm]
+        )
+        if clip_args_sum > 1:
+            raise ValueError(
+                "Only one of `clipnorm`, `clipvalue` and `global_clipnorm` can "
+                f"be set. Received: clipnorm={clipnorm}, "
+                f"clipvalue={clipvalue}, global_clipnorm={global_clipnorm}"
+            )
+        self.built = False
+
+        # Set up variable tracking.
+        self._variables = []
+        self._trainable_variables = []
+        self._tracker = tracking.Tracker(
+            {
+                "variables": (
+                    lambda x: isinstance(x, backend.Variable),
+                    self._variables,
+                ),
+            }
+        )
+        self._trainable_variables_indices = {}
+
+        # Create iteration variable
+        # Note: dtype="int" will resolve to int32 in JAX
+        # (since int64 is disallowed in JAX) and to int64 in TF.
+        with backend.name_scope(self.name, caller=self):
+            iterations = backend.Variable(
+                0,
+                name="iteration",
+                dtype="int",
+                trainable=False,
+                aggregation="only_first_replica",
+            )
+        self._track_variable(iterations)
+        self._iterations = iterations
+
+        # Create learning rate (schedule or variable)
+        if isinstance(
+            learning_rate, learning_rate_schedule.LearningRateSchedule
+        ):
+            self._learning_rate = learning_rate
+        elif callable(learning_rate):
+            self._learning_rate = learning_rate
+        else:
+            if not isinstance(learning_rate, float):
+                raise ValueError(
+                    "Argument `learning_rate` should be float, or an instance "
+                    "of LearningRateSchedule, or a callable "
+                    "(that takes in the current iteration value "
+                    "and returns the corresponding learning rate value). "
+                    f"Received instead: learning_rate={learning_rate}"
+                )
+            with backend.name_scope(self.name, caller=self):
+                learning_rate = backend.Variable(
+                    learning_rate,
+                    name="learning_rate",
+                    dtype=backend.floatx(),
+                    trainable=False,
+                    aggregation="only_first_replica",
+                )
+            self._track_variable(learning_rate)
+            self._learning_rate = learning_rate
+
+    @property
+    def iterations(self):
+        if self.gradient_accumulation_steps:
+            return ops.floor_divide(
+                self._iterations, self.gradient_accumulation_steps
+            )
+
+        return self._iterations
+
+    def _track_variable(self, variable):
+        self._tracker.add_to_store("variables", variable)
+
+    @tracking.no_automatic_dependency_tracking
+    def build(self, variables):
+        if self.use_ema:
+            self._model_variables_moving_average = []
+        if self.gradient_accumulation_steps:
+            self._accumulated_gradients = []
+        for i, variable in enumerate(variables):
+            self._trainable_variables_indices[self._var_key(variable)] = i
+            if self.use_ema:
+                self._model_variables_moving_average.append(
+                    self.add_variable_from_reference(
+                        variable,
+                        name="average",
+                    )
+                )
+            if self.gradient_accumulation_steps:
+                self._accumulated_gradients.append(
+                    self.add_variable_from_reference(
+                        variable,
+                        name="gradient_accumulator",
+                    )
+                )
+        self._trainable_variables = variables[:]
+        self.built = True
+
+    def _var_key(self, variable):
+        # Helper function to get a stable ID and the variable instance mapping.
+        return id(variable)
+
+    @property
+    def variables(self):
+        return self._variables[:]
+
+    def _get_variable_index(self, variable):
+        return self._trainable_variables_indices[self._var_key(variable)]
+
+    def add_variable(
+        self,
+        shape,
+        initializer="zeros",
+        dtype=None,
+        aggregation="none",
+        name=None,
+    ):
+        """Add a variable to the optimizer.
+
+        Args:
+            shape: Shape tuple for the variable. Must be fully-defined
+                (no `None` entries).
+            initializer: Initializer object to use to populate the initial
+                variable value, or string name of a built-in initializer
+                (e.g. `"random_normal"`). Defaults to `"zeros"`.
+            dtype: Dtype of the variable to create, e.g. `"float32"`. If
+                unspecified, defaults to the `keras.backend.floatx()`.
+            aggregation: Optional string, one of `None`, `"none"`, `"mean"`,
+                `"sum"` or `"only_first_replica"`. Annotates the variable with
+                the type of multi-replica aggregation to be used for this
+                variable when writing custom data parallel training loops.
+                Defaults to `"none"`.
+            name: String name of the variable. Useful for debugging purposes.
+
+        Returns:
+            An optimizer variable, in the format of `keras.Variable`.
+        """
+        self._check_super_called()
+        initializer = initializers.get(initializer)
+        with backend.name_scope(self.name, caller=self):
+            variable = backend.Variable(
+                initializer=initializer,
+                shape=shape,
+                dtype=dtype,
+                trainable=False,
+                aggregation=aggregation,
+                name=name,
+            )
+        self._track_variable(variable)
+        return variable
+
+    def add_variable_from_reference(
+        self, reference_variable, name=None, initializer="zeros"
+    ):
+        """Add an optimizer variable from the model variable.
+
+        Create an optimizer variable based on the information of model variable.
+        For example, in SGD optimizer momemtum, for each model variable, a
+        corresponding momemtum variable is created of the same shape and dtype.
+
+        Args:
+            reference_variable: `keras.Variable`. The corresponding model
+                variable to the optimizer variable to be created.
+            name: Optional string. The name prefix of the optimizer variable to
+                be created. If not provided, it will be set to `"var"`. The
+                variable name will follow the pattern
+                `{variable_name}_{reference_variable.name}`,
+                e.g., `momemtum/dense_1`. Defaults to `None`.
+            initializer: Initializer object to use to populate the initial
+                variable value, or string name of a built-in initializer
+                (e.g. `"random_normal"`). If unspecified, defaults to
+                `"zeros"`.
+
+        Returns:
+            An optimizer variable, in the format of `keras.Variable`.
+        """
+        name = name or "var"
+        if hasattr(reference_variable, "path"):
+            name = reference_variable.path.replace("/", "_") + "_" + name
+        else:
+            name = (
+                str(reference_variable.name).replace("/", "_").replace(":", "_")
+                + "_"
+                + name
+            )
+        return self.add_variable(
+            shape=reference_variable.shape,
+            initializer=initializer,
+            dtype=reference_variable.dtype,
+            name=name,
+        )
+
+    def _check_variables_are_known(self, variables):
+        for v in variables:
+            if self._var_key(v) not in self._trainable_variables_indices:
+                raise ValueError(
+                    f"Unknown variable: {v}. This optimizer can only "
+                    "be called for the variables it was originally built with. "
+                    "When working with a new set of variables, you should "
+                    "recreate a new optimizer instance."
+                )
+
+    def assign(self, variable, value):
+        """Assign a value to a variable.
+
+        This should be used in optimizers instead of `variable.assign(value)` to
+        support backend specific optimizations.
+        Note that the variable can be a model variable or an optimizer variable;
+        it can be a backend native variable or a Keras variable.
+
+        Args:
+            variable: The variable to update.
+            value: The value to add to the variable.
+        """
+        variable.assign(value)
+
+    def assign_add(self, variable, value):
+        """Add a value to a variable.
+
+        This should be used in optimizers instead of
+        `variable.assign_add(value)` to support backend specific optimizations.
+        Note that the variable can be a model variable or an optimizer variable;
+        it can be a backend native variable or a Keras variable.
+
+        Args:
+            variable: The variable to update.
+            value: The value to add to the variable.
+        """
+        variable.assign_add(value)
+
+    def assign_sub(self, variable, value):
+        """Subtract a value from a variable.
+
+        This should be used in optimizers instead of
+        `variable.assign_sub(value)` to support backend specific optimizations.
+        Note that the variable can be a model variable or an optimizer variable;
+        it can be a backend native variable or a Keras variable.
+
+        Args:
+            variable: The variable to update.
+            value: The value to add to the variable.
+        """
+        variable.assign_sub(value)
+
+    def update_step(self, gradient, variable, learning_rate):
+        raise NotImplementedError
+
+    def apply_gradients(self, grads_and_vars):
+        grads, trainable_variables = zip(*grads_and_vars)
+        self.apply(grads, trainable_variables)
+        # Return iterations for compat with tf.keras.
+        return self._iterations
+
+    def apply(self, grads, trainable_variables=None):
+        """Update traininable variables according to provided gradient values.
+
+        `grads` should be a list of gradient tensors
+        with 1:1 mapping to the list of variables the optimizer was built with.
+
+        `trainable_variables` can be provided
+        on the first call to build the optimizer.
+        """
+        if len(grads) == 0:
+            # It is possible that the grad is empty. In this case,
+            # `apply_gradients` is a no-op.
+            return
+
+        if trainable_variables is None:
+            if not self.built:
+                raise ValueError(
+                    "When passing `grads` without `variables`, the optimizer "
+                    "must already be built on a list of variables. "
+                    "Call `optimizer.build(trainable_variables)` first. "
+                )
+            if len(grads) != len(self._trainable_variables_indices):
+                raise ValueError(
+                    "When passing `grads` as a list of gradient tensors, the "
+                    f"gradients must match `optimizer.variables` one-to-on. "
+                    f"Received a list of {len(grads)} gradients, but the "
+                    f"optimizer is tracking {len(self._trainable_variables)} "
+                    "trainable variables."
+                )
+            trainable_variables = self._trainable_variables
+        else:
+            trainable_variables = list(trainable_variables)
+            # Optionally build optimizer.
+            if not self.built:
+                with backend.name_scope(self.name, caller=self):
+                    self.build(trainable_variables)
+                self.built = True
+            self._check_variables_are_known(trainable_variables)
+
+        with backend.name_scope(self.name, caller=self):
+            # Overwrite targeted variables directly with their gradients if
+            # their `overwrite_with_gradient` is set.
+            grads, trainable_variables = (
+                self._overwrite_variables_directly_with_gradients(
+                    grads, trainable_variables
+                )
+            )
+
+            # Filter empty gradients.
+            grads, trainable_variables = self._filter_empty_gradients(
+                grads, trainable_variables
+            )
+            if len(list(grads)) == 0:
+                return
+
+            # Unscale gradients.
+            scale = self.loss_scale_factor
+            if scale is not None:
+                grads = [g if g is None else g / scale for g in grads]
+
+            # Apply gradient updates.
+            self._backend_apply_gradients(grads, trainable_variables)
+            # Apply variable constraints after applying gradients.
+            for variable in trainable_variables:
+                if variable.constraint is not None:
+                    variable.assign(variable.constraint(variable))
+
+    def _backend_apply_gradients(self, grads, trainable_variables):
+        """Apply method that can be overridden by different backends.
+
+        JAX overrides it in order to deal with statelessness in gradient
+        accumulation and EMA handling.
+
+        The below implementation is intended to be generally backend-agnostic,
+        but may not work with all backends.
+
+        This method does 4 things:
+        - Call the optimizer's update_step() to update trainable variables
+            and optimizer variables.
+        - Update EMA variables, if EMA is configured.
+        - Update gradient accumulators, if gradient accumulation is configured.
+        - Update the iteration counter.
+        """
+        if self.gradient_accumulation_steps:
+            is_update_step = (
+                self._iterations + 1
+            ) % self.gradient_accumulation_steps == 0
+            # `trainable_variables` might have been filtered in previous
+            # processing steps, so we need to ensure the correct mapping between
+            # `self._accumulated_gradients` and `trainable_variables`
+            acc_grads = [
+                self._accumulated_gradients[self._get_variable_index(v)]
+                for v in trainable_variables
+            ]
+
+            def _update_step_fn(grads, trainable_variables):
+                # Run update step with accumulated grads + reset accumulators
+                steps = self.gradient_accumulation_steps
+                grads = [
+                    (g + acc_g) / steps for g, acc_g in zip(grads, acc_grads)
+                ]
+
+                # Apply clipping and weight decay.
+                grads = self._clip_gradients(grads)
+                self._apply_weight_decay(trainable_variables)
+
+                self._backend_update_step(
+                    grads, trainable_variables, self.learning_rate
+                )
+                self._backend_reset_gradient_accumulators()
+
+            ops.cond(
+                is_update_step,
+                lambda: _update_step_fn(grads, trainable_variables),
+                lambda: self._backend_increment_gradient_accumulators(
+                    grads, acc_grads
+                ),
+            )
+        else:
+            # Apply clipping and weight decay.
+            grads = self._clip_gradients(grads)
+            self._apply_weight_decay(trainable_variables)
+
+            # Run update step.
+            self._backend_update_step(
+                grads, trainable_variables, self.learning_rate
+            )
+
+        if self.use_ema:
+            self._update_model_variables_moving_average(
+                self._trainable_variables
+            )
+            if self.ema_overwrite_frequency:
+                # Only when self.ema_overwrite_frequency is not None, we
+                # overwrite the model variables.
+                should_overwrite_model_vars = (
+                    self.iterations + 1
+                ) % self.ema_overwrite_frequency == 0
+                ops.cond(
+                    should_overwrite_model_vars,
+                    lambda: self._overwrite_model_variables_with_average_value(
+                        self._trainable_variables
+                    ),
+                    lambda: None,
+                )
+        # Update iteration counter.
+        self._iterations.assign_add(1)
+
+    def _backend_update_step(self, grads, trainable_variables, learning_rate):
+        """Collective update_step that can be overridden by the backend.
+
+        It is overridden by torch for performance reasons, and
+        by TF to support tf.distribute.
+        """
+        for grad, var in zip(grads, trainable_variables):
+            self.update_step(grad, var, learning_rate)
+
+    def _backend_reset_gradient_accumulators(self):
+        for g_acc in self._accumulated_gradients:
+            g_acc.assign(ops.zeros(g_acc.shape, dtype=g_acc.dtype))
+
+    def _backend_increment_gradient_accumulators(self, grads, acc_grads):
+        new_g_accs = [(g + acc_g) for g, acc_g in zip(grads, acc_grads)]
+        for n_g_acc, g_acc in zip(new_g_accs, acc_grads):
+            g_acc.assign(n_g_acc)
+
+    def stateless_apply(self, optimizer_variables, grads, trainable_variables):
+        self._check_super_called()
+
+        if not self.built:
+            raise ValueError(
+                f"To call `stateless_apply`, {self.__class__.__name__} "
+                "must be built (i.e. its variables must have been created). "
+                "You can build it via `optimizer.build(trainable_variables)`."
+            )
+        if len(optimizer_variables) != len(self.variables):
+            raise ValueError(
+                "Argument `optimizer_variables` must be a list of tensors "
+                f"corresponding 1:1 to {self.__class__.__name__}().variables. "
+                f"Received list with length {len(optimizer_variables)}, but "
+                f"expected {len(self.variables)} variables."
+            )
+        if len(trainable_variables) != len(self._trainable_variables):
+            raise ValueError(
+                "Argument `optimizer_variables` must be a list of tensors "
+                "corresponding 1:1 to the trainable variables list that "
+                "the optimizer was built with. Received "
+                f"len(trainable_variables) == {len(trainable_variables)} "
+                "whereas the optimizer was built with "
+                f"{len(self._trainable_variables)} variables."
+            )
+
+        # Gather variable mapping
+        mapping = list(
+            zip(self._trainable_variables, trainable_variables)
+        ) + list(zip(self.variables, optimizer_variables))
+
+        # Call in stateless scope
+        with backend.StatelessScope(state_mapping=mapping) as scope:
+            self.apply(grads)
+
+        # Gather updated variables
+        trainable_variables = []
+        for v in self._trainable_variables:
+            new_v = scope.get_current_value(v)
+            if new_v is not None:
+                trainable_variables.append(new_v)
+            else:
+                trainable_variables.append(v)
+        optimizer_variables = []
+        for v in self.variables:
+            new_v = scope.get_current_value(v)
+            if new_v is not None:
+                optimizer_variables.append(new_v)
+            else:
+                optimizer_variables.append(v)
+        return trainable_variables, optimizer_variables
+
+    def scale_loss(self, loss):
+        """Scale the loss before computing gradients.
+
+        Scales the loss before gradients are computed in a `train_step`. This
+        is primarily useful during mixed precision training to prevent numeric
+        underflow.
+        """
+        if self.loss_scale_factor is not None:
+            return loss * self.loss_scale_factor
+        return loss
+
+    @property
+    def learning_rate(self):
+        return self._get_current_learning_rate()
+
+    @learning_rate.setter
+    def learning_rate(self, learning_rate):
+        if isinstance(self._learning_rate, backend.Variable):
+            prev_lr_var = self._learning_rate
+        else:
+            prev_lr_var = None
+        if isinstance(
+            learning_rate, learning_rate_schedule.LearningRateSchedule
+        ):
+            self._learning_rate = learning_rate
+        elif callable(learning_rate):
+            self._learning_rate = learning_rate
+        else:
+            if isinstance(
+                self._learning_rate, learning_rate_schedule.LearningRateSchedule
+            ):
+                raise TypeError(
+                    "This optimizer was created with a `LearningRateSchedule`"
+                    " object as its `learning_rate` constructor argument, "
+                    "hence its learning rate is not settable. If you need the"
+                    " learning rate to be settable, you should instantiate "
+                    "the optimizer with a float `learning_rate` argument."
+                )
+            self._learning_rate.assign(learning_rate)
+        if prev_lr_var is not None and not isinstance(
+            self._learning_rate, backend.Variable
+        ):
+            # Untrack learning rate variable
+            self._untrack_variable(prev_lr_var)
+
+    def set_weights(self, weights):
+        """Set the weights of the optimizer."""
+        if not self.built:
+            raise ValueError(
+                "You are calling `set_weights()` on an optimizer that has not "
+                "yet been built. Please call "
+                "`optimizer.build(trainable_variables)` to create the "
+                "optimizer weights before calling `set_weights()`."
+            )
+        for variable, weight in zip(self._variables, weights):
+            if variable.shape != weight.shape:
+                raise ValueError(
+                    f"Optimizer variable {self._var_key(variable)} has shape "
+                    f"{str(variable.shape)} not compatible with provided "
+                    f"weight shape {str(weight.shape)}."
+                )
+            variable.assign(weight)
+
+    def save_own_variables(self, store):
+        """Get the state of this optimizer object."""
+        for i, variable in enumerate(self.variables):
+            store[str(i)] = variable.numpy()
+
+    def load_own_variables(self, store):
+        """Set the state of this optimizer object."""
+        if len(store.keys()) != len(self.variables):
+            msg = (
+                f"Skipping variable loading for optimizer '{self.name}', "
+                f"because it has {len(self.variables)} variables whereas "
+                f"the saved optimizer has {len(store.keys())} variables. "
+            )
+            if len(self.variables) == 0:
+                msg += (
+                    "This is likely because the optimizer has not been "
+                    "called/built yet."
+                )
+            warnings.warn(msg, stacklevel=2)
+            return
+        for i, variable in enumerate(self.variables):
+            variable.assign(store[str(i)])
+
+    def _get_current_learning_rate(self):
+        if isinstance(
+            self._learning_rate, learning_rate_schedule.LearningRateSchedule
+        ):
+            return self._learning_rate(self._iterations)
+        elif callable(self._learning_rate):
+            return self._learning_rate()
+        return self._learning_rate
+
+    def _overwrite_variables_directly_with_gradients(self, grads, vars):
+        """Overwrite the variables directly by their gradients.
+
+        This method is designed for a special case where we want to overwrite
+        the variable directly with its computed gradient. For example, in float8
+        training, new `scale` and `amax_history` are computed as gradients, and
+        we want to overwrite them directly instead of following the typical
+        procedure such as gradient descent with a learning rate, gradient
+        clipping and weight decaying.
+
+        After the update, the processed pairs will be filtered out.
+        """
+        # Shortcut for `tf.Variable` because it doesn't have a
+        # `overwrite_with_gradient` attr
+        if any(not hasattr(v, "overwrite_with_gradient") for v in vars):
+            return grads, vars
+
+        # Shallow copies
+        filtered_grads = list(grads)
+        filtered_vars = list(vars)
+
+        # Iterate from right to left for safe popping
+        for i in range(len(filtered_grads) - 1, -1, -1):
+            g, v = filtered_grads[i], filtered_vars[i]
+            if v.overwrite_with_gradient:
+                if self.gradient_accumulation_steps:
+                    # Utilize a stateless manner for JAX compatibility
+                    steps = self.gradient_accumulation_steps
+                    is_update_step = (self._iterations + 1) % steps == 0
+                    acc_g = self._accumulated_gradients[
+                        self._get_variable_index(v)
+                    ]
+                    # `ops.maximum` is utilized for gradient accumulation for
+                    # `overwrite_with_gradient=True` variables
+                    new_g_acc = ops.cond(
+                        is_update_step,
+                        lambda: ops.zeros(g.shape, dtype=g.dtype),
+                        lambda: ops.maximum(g, acc_g),
+                    )
+                    new_g = ops.cond(
+                        is_update_step,
+                        lambda: ops.maximum(g, acc_g),
+                        lambda: g,
+                    )
+                    new_v = ops.cond(
+                        is_update_step, lambda: new_g, lambda: v.value
+                    )
+                    v.assign(new_v)
+                    acc_g.assign(new_g_acc)
+                else:
+                    v.assign(g)
+                filtered_grads.pop(i)
+                filtered_vars.pop(i)
+        return filtered_grads, filtered_vars
+
+    def _filter_empty_gradients(self, grads, vars):
+        filtered_grads = list(grads)
+        filtered_vars = list(vars)
+        missing_grad_vars = []
+
+        # Iterate from right to left for safe popping
+        for i in range(len(filtered_grads) - 1, -1, -1):
+            if filtered_grads[i] is None:
+                filtered_grads.pop(i)
+                v = filtered_vars.pop(i)
+                try:
+                    missing_grad_vars.append(v.path)
+                except AttributeError:
+                    # `tf.Variable` doesn't have `path` attr.
+                    missing_grad_vars.append(v.name)
+
+        if not filtered_grads:
+            raise ValueError("No gradients provided for any variable.")
+        if missing_grad_vars:
+            warnings.warn(
+                "Gradients do not exist for variables "
+                f"{list(reversed(missing_grad_vars))} when minimizing the loss."
+                " If using `model.compile()`, did you forget to provide a "
+                "`loss` argument?"
+            )
+        return filtered_grads, filtered_vars
+
+    def _clip_gradients(self, grads):
+        if self.clipnorm and self.clipnorm > 0:
+            return [
+                self._clip_by_norm(g) if g is not None else g for g in grads
+            ]
+        elif self.global_clipnorm and self.global_clipnorm > 0:
+            return clip_by_global_norm(grads, self.global_clipnorm)
+        elif self.clipvalue and self.clipvalue > 0:
+            v = self.clipvalue
+            return [ops.clip(g, -v, v) if g is not None else g for g in grads]
+        else:
+            return grads
+
+    def exclude_from_weight_decay(self, var_list=None, var_names=None):
+        """Exclude variables from weight decay.
+
+        This method must be called before the optimizer's `build` method is
+        called. You can set specific variables to exclude out, or set a list of
+        strings as the anchor words, if any of which appear in a variable's
+        name, then the variable is excluded.
+
+        Args:
+            var_list: A list of `Variable`s to exclude from weight decay.
+            var_names: A list of strings. If any string in `var_names` appear
+                in the model variable's name, then this model variable is
+                excluded from weight decay. For example, `var_names=['bias']`
+                excludes all bias variables from weight decay.
+        """
+        if hasattr(self, "_built") and self._built:
+            raise ValueError(
+                "`exclude_from_weight_decay()` can only be configured before "
+                "the optimizer is built."
+            )
+
+        # Use a `set` for the ids of `var_list` to speed up the searching
+        if var_list:
+            self._exclude_from_weight_decay = set(
+                self._var_key(variable) for variable in var_list
+            )
+        else:
+            self._exclude_from_weight_decay = set()
+
+        # Precompile the pattern for `var_names` to speed up the searching
+        if var_names and len(var_names) > 0:
+            self._exclude_from_weight_decay_pattern = re.compile(
+                "|".join(set(var_names))
+            )
+        else:
+            self._exclude_from_weight_decay_pattern = None
+
+        # Reset cache
+        self._exclude_from_weight_decay_cache = dict()
+
+    def _use_weight_decay(self, variable):
+        variable_id = self._var_key(variable)
+
+        # Immediately return the value if `variable_id` hits the cache
+        if not hasattr(self, "_exclude_from_weight_decay_cache"):
+            self._exclude_from_weight_decay_cache = dict()
+        if variable_id in self._exclude_from_weight_decay_cache:
+            return self._exclude_from_weight_decay_cache[variable_id]
+
+        # Determine whether the variable should apply weight decay or not
+        exclude_from_weight_decay = getattr(
+            self, "_exclude_from_weight_decay", set()
+        )
+        exclude_from_weight_decay_pattern = getattr(
+            self, "_exclude_from_weight_decay_pattern", None
+        )
+        if variable_id in exclude_from_weight_decay:
+            self._exclude_from_weight_decay_cache[variable_id] = False
+            return False
+        if exclude_from_weight_decay_pattern is not None:
+            if (
+                re.search(exclude_from_weight_decay_pattern, variable.name)
+                is not None
+            ):
+                self._exclude_from_weight_decay_cache[variable_id] = False
+                return False
+        self._exclude_from_weight_decay_cache[variable_id] = True
+        return True
+
+    def _apply_weight_decay(self, variables):
+        if self.weight_decay is None:
+            return
+        for variable in variables:
+            if self._use_weight_decay(variable):
+                lr = ops.cast(self.learning_rate, variable.dtype)
+                wd = ops.cast(self.weight_decay, variable.dtype)
+                variable.assign(variable - variable * wd * lr)
+
+    def _check_super_called(self):
+        if not hasattr(self, "_lock"):
+            raise RuntimeError(
+                f"In optimizer '{self.__class__.__name__}', you forgot to call "
+                "`super().__init__()` as the first statement "
+                "in the `__init__()` method. "
+                "Go add it!"
+            )
+
+    def _update_model_variables_moving_average(self, trainable_variables):
+        """Update the stored moving average using the latest value."""
+        if self.use_ema:
+            for var, average in zip(
+                trainable_variables, self._model_variables_moving_average
+            ):
+                not_first_step = ops.not_equal(self.iterations, 0)
+                momentum = (
+                    ops.cast(not_first_step, var.dtype) * self.ema_momentum
+                )
+                average.assign(momentum * average + (1 - momentum) * var)
+
+    def _overwrite_model_variables_with_average_value(
+        self, trainable_variables
+    ):
+        """Overwrite model variables with its moving average."""
+        if len(trainable_variables) != len(
+            self._model_variables_moving_average
+        ):
+            raise ValueError(
+                f"The length of model variables ({len(trainable_variables)}) "
+                "to override does not match the length of model variables "
+                "stored in the optimizer "
+                f"({len(self._model_variables_moving_average)}). Please "
+                "check if the optimizer was called on your model."
+            )
+        for var, average_var in zip(
+            trainable_variables, self._model_variables_moving_average
+        ):
+            var.assign(average_var)
+
+    def finalize_variable_values(self, var_list):
+        """Set the final value of model's trainable variables.
+
+        Sometimes there are some extra steps before ending the variable updates,
+        such as overriding the model variables with its average value.
+
+        Args:
+          var_list: list of model variables.
+        """
+        if self.use_ema:
+            # If the optimizer uses EMA, then when finalizing, we replace the
+            # model variable value with its moving average stored inside
+            # optimizer.
+            self._overwrite_model_variables_with_average_value(var_list)
+
+    def _obj_type(self):
+        return "Optimizer"
+
+    def get_config(self):
+        """Returns the config of the optimizer.
+
+        An optimizer config is a Python dictionary (serializable)
+        containing the configuration of an optimizer.
+        The same optimizer can be reinstantiated later
+        (without any saved state) from this configuration.
+
+        Subclass optimizer should override this method to include other
+        hyperparameters.
+
+        Returns:
+            Python dictionary.
+        """
+
+        if isinstance(
+            self._learning_rate, learning_rate_schedule.LearningRateSchedule
+        ):
+            learning_rate = learning_rate_schedule.serialize(
+                self._learning_rate
+            )
+        elif isinstance(self._learning_rate, backend.Variable):
+            learning_rate = float(self._learning_rate.numpy())
+        elif ops.is_tensor(self._learning_rate):
+            learning_rate = float(self._learning_rate)
+        elif callable(self._learning_rate):
+            learning_rate = serialization_lib.serialize_keras_object(
+                self._learning_rate
+            )
+        else:
+            learning_rate = 0.5
+
+        config = {
+            "name": self.name,
+            "learning_rate": learning_rate,
+            "weight_decay": self.weight_decay,
+            "clipnorm": self.clipnorm,
+            "global_clipnorm": self.global_clipnorm,
+            "clipvalue": self.clipvalue,
+            "use_ema": self.use_ema,
+            "ema_momentum": self.ema_momentum,
+            "ema_overwrite_frequency": self.ema_overwrite_frequency,
+            "loss_scale_factor": self.loss_scale_factor,
+            "gradient_accumulation_steps": self.gradient_accumulation_steps,
+        }
+        return config
+
+    @classmethod
+    def from_config(cls, config, custom_objects=None):
+        """Creates an optimizer from its config.
+
+        This method is the reverse of `get_config`, capable of instantiating the
+        same optimizer from the config dictionary.
+
+        Args:
+            config: A Python dictionary, typically the output of get_config.
+            custom_objects: A Python dictionary mapping names to additional
+              user-defined Python objects needed to recreate this optimizer.
+
+        Returns:
+            An optimizer instance.
+        """
+        if "learning_rate" in config:
+            if isinstance(config["learning_rate"], dict):
+                config["learning_rate"] = (
+                    serialization_lib.deserialize_keras_object(
+                        config["learning_rate"], custom_objects=custom_objects
+                    )
+                )
+        return cls(**config)
+
+    def __setattr__(self, name, value):
+        # Prevent users from attaching state to the
+        # layer before `super()` is called -- since that
+        # state would silently not be tracked.
+        if name != "_lock":
+            self._check_super_called()
+        # Track Variables.
+        if hasattr(self, "_tracker"):
+            value = self._tracker.track(value)
+        return super().__setattr__(name, value)
+
+    def _clip_by_norm(self, values, axes=None):
+        # Calculate L2-norm, clip elements by ratio of clip_norm to L2-norm
+        l2sum = ops.sum(ops.square(values), axes, keepdims=True)
+        pred = l2sum > 0
+        # Two-tap tf.where trick to bypass NaN gradients
+        l2sum_safe = ops.where(pred, l2sum, ops.ones_like(l2sum))
+        l2norm = ops.where(pred, ops.sqrt(l2sum_safe), l2sum)
+        intermediate = ops.multiply(values, self.clipnorm)
+        values_clip = ops.convert_to_tensor(intermediate) / ops.maximum(
+            l2norm, self.clipnorm
+        )
+        return values_clip
+
+    def _untrack_variable(self, variable):
+        previous_lock_state = self._tracker.locked
+        self._tracker.unlock()
+        self._tracker.untrack(variable)
+        if previous_lock_state is True:
+            self._tracker.lock()
+
+
+base_optimizer_keyword_args = """name: String. The name to use
+            for momentum accumulator weights created by
+            the optimizer.
+        weight_decay: Float. If set, weight decay is applied.
+        clipnorm: Float. If set, the gradient of each weight is individually
+            clipped so that its norm is no higher than this value.
+        clipvalue: Float. If set, the gradient of each weight is clipped to be
+            no higher than this value.
+        global_clipnorm: Float. If set, the gradient of all weights is clipped
+            so that their global norm is no higher than this value.
+        use_ema: Boolean, defaults to `False`.
+            If `True`, exponential moving average
+            (EMA) is applied. EMA consists of computing an exponential moving
+            average of the weights of the model (as the weight values change
+            after each training batch), and periodically overwriting the
+            weights with their moving average.
+        ema_momentum: Float, defaults to 0.99. Only used if `use_ema=True`.
+            This is the momentum to use when computing
+            the EMA of the model's weights:
+            `new_average = ema_momentum * old_average + (1 - ema_momentum) *
+            current_variable_value`.
+        ema_overwrite_frequency: Int or None, defaults to None. Only used if
+            `use_ema=True`. Every `ema_overwrite_frequency` steps of iterations,
+            we overwrite the model variable by its moving average.
+            If None, the optimizer
+            does not overwrite model variables in the middle of training,
+            and you need to explicitly overwrite the variables
+            at the end of training by calling
+            `optimizer.finalize_variable_values()` (which updates the model
+            variables in-place). When using the built-in `fit()` training loop,
+            this happens automatically after the last epoch,
+            and you don't need to do anything.
+        loss_scale_factor: Float or `None`. If a float, the scale factor will
+            be multiplied the loss before computing gradients, and the inverse
+            of the scale factor will be multiplied by the gradients before
+            updating variables. Useful for preventing underflow during
+            mixed precision training. Alternately,
+            `keras.optimizers.LossScaleOptimizer` will
+            automatically set a loss scale factor.
+        gradient_accumulation_steps: Int or `None`. If an int, model & optimizer
+            variables will not be updated at every step; instead they will be
+            updated every `gradient_accumulation_steps` steps, using the average
+            value of the gradients since the last update. This is known as
+            "gradient accumulation". This can be useful
+            when your batch size is very small, in order to reduce gradient
+            noise at each update step. EMA frequency will look at "accumulated"
+            iterations value (optimizer steps // gradient_accumulation_steps).
+            Learning rate schedules will look at "real" iterations value
+            (optimizer steps).
+"""
+
+
+def global_norm(value_list):
+    """Computes the global norm of multiple tensors."""
+    squared_norms = [
+        ops.sum(ops.square(v)) for v in value_list if v is not None
+    ]
+    squared_norm = ops.sum(ops.stack(squared_norms))
+    return ops.sqrt(squared_norm)
+
+
+def clip_by_global_norm(value_list, clip_norm):
+    use_norm = global_norm(value_list)
+    # Calculate L2-norm, clip elements by ratio of clip_norm to L2-norm
+    scale_for_finite = clip_norm * ops.minimum(1.0 / use_norm, 1.0 / clip_norm)
+    # If use_norm is any finite number, this is a no-op. For inf/-inf/NaN,
+    # this will make scale NaN.
+    scale = scale_for_finite + (use_norm - use_norm)
+    return [v * scale if v is not None else v for v in value_list]

SwarmUI/dlbackend/ComfyUI/venv/lib/python3.10/site-packages/keras/src/optimizers/ftrl.py

@@ -0,0 +1,249 @@
+from keras.src import initializers
+from keras.src import ops
+from keras.src.api_export import keras_export
+from keras.src.optimizers import optimizer
+
+
+@keras_export(["keras.optimizers.Ftrl"])
+class Ftrl(optimizer.Optimizer):
+    r"""Optimizer that implements the FTRL algorithm.
+
+    "Follow The Regularized Leader" (FTRL) is an optimization algorithm
+    developed at Google for click-through rate prediction in the early 2010s. It
+    is most suitable for shallow models with large and sparse feature spaces.
+    The algorithm is described by
+    [McMahan et al., 2013](https://research.google.com/pubs/archive/41159.pdf).
+    The Keras version has support for both online L2 regularization
+    (the L2 regularization described in the paper
+    above) and shrinkage-type L2 regularization
+    (which is the addition of an L2 penalty to the loss function).
+
+    Initialization:
+
+    ```python
+    n = 0
+    sigma = 0
+    z = 0
+    ```
+
+    Update rule for one variable `w`:
+
+    ```python
+    prev_n = n
+    n = n + g ** 2
+    sigma = (n ** -lr_power - prev_n ** -lr_power) / lr
+    z = z + g - sigma * w
+    if abs(z) < lambda_1:
+      w = 0
+    else:
+      w = (sgn(z) * lambda_1 - z) / ((beta + sqrt(n)) / alpha + lambda_2)
+    ```
+
+    Notation:
+
+    - `lr` is the learning rate
+    - `g` is the gradient for the variable
+    - `lambda_1` is the L1 regularization strength
+    - `lambda_2` is the L2 regularization strength
+    - `lr_power` is the power to scale n.
+
+    Check the documentation for the `l2_shrinkage_regularization_strength`
+    parameter for more details when shrinkage is enabled, in which case gradient
+    is replaced with a gradient with shrinkage.
+
+    Args:
+        learning_rate: A float, a
+            `keras.optimizers.schedules.LearningRateSchedule` instance, or
+            a callable that takes no arguments and returns the actual value to
+            use. The learning rate. Defaults to `0.001`.
+        learning_rate_power: A float value, must be less or equal to zero.
+            Controls how the learning rate decreases during training. Use zero
+            for a fixed learning rate.
+        initial_accumulator_value: The starting value for accumulators. Only
+            zero or positive values are allowed.
+        l1_regularization_strength: A float value, must be greater than or equal
+            to zero. Defaults to `0.0`.
+        l2_regularization_strength: A float value, must be greater than or equal
+            to zero. Defaults to `0.0`.
+        l2_shrinkage_regularization_strength: A float value, must be greater
+            than or equal to zero. This differs from L2 above in that the L2
+            above is a stabilization penalty, whereas this L2 shrinkage is a
+            magnitude penalty. When input is sparse shrinkage will only happen
+            on the active weights.
+        beta: A float value, representing the beta value from the paper.
+            Defaults to `0.0`.
+        {{base_optimizer_keyword_args}}
+    """
+
+    def __init__(
+        self,
+        learning_rate=0.001,
+        learning_rate_power=-0.5,
+        initial_accumulator_value=0.1,
+        l1_regularization_strength=0.0,
+        l2_regularization_strength=0.0,
+        l2_shrinkage_regularization_strength=0.0,
+        beta=0.0,
+        weight_decay=None,
+        clipnorm=None,
+        clipvalue=None,
+        global_clipnorm=None,
+        use_ema=False,
+        ema_momentum=0.99,
+        ema_overwrite_frequency=None,
+        loss_scale_factor=None,
+        gradient_accumulation_steps=None,
+        name="ftrl",
+        **kwargs,
+    ):
+        super().__init__(
+            learning_rate=learning_rate,
+            name=name,
+            weight_decay=weight_decay,
+            clipnorm=clipnorm,
+            clipvalue=clipvalue,
+            global_clipnorm=global_clipnorm,
+            use_ema=use_ema,
+            ema_momentum=ema_momentum,
+            ema_overwrite_frequency=ema_overwrite_frequency,
+            loss_scale_factor=loss_scale_factor,
+            gradient_accumulation_steps=gradient_accumulation_steps,
+            **kwargs,
+        )
+
+        if initial_accumulator_value < 0.0:
+            raise ValueError(
+                "`initial_accumulator_value` needs to be positive or zero. "
+                "Received: initial_accumulator_value="
+                f"{initial_accumulator_value}."
+            )
+        if learning_rate_power > 0.0:
+            raise ValueError(
+                "`learning_rate_power` needs to be negative or zero. Received: "
+                f"learning_rate_power={learning_rate_power}."
+            )
+        if l1_regularization_strength < 0.0:
+            raise ValueError(
+                "`l1_regularization_strength` needs to be positive or zero. "
+                "Received: l1_regularization_strength="
+                f"{l1_regularization_strength}."
+            )
+        if l2_regularization_strength < 0.0:
+            raise ValueError(
+                "`l2_regularization_strength` needs to be positive or zero. "
+                "Received: l2_regularization_strength="
+                f"{l2_regularization_strength}."
+            )
+        if l2_shrinkage_regularization_strength < 0.0:
+            raise ValueError(
+                "`l2_shrinkage_regularization_strength` needs to be positive "
+                "or zero. Received: l2_shrinkage_regularization_strength"
+                f"={l2_shrinkage_regularization_strength}."
+            )
+
+        self.learning_rate_power = learning_rate_power
+        self.initial_accumulator_value = initial_accumulator_value
+        self.l1_regularization_strength = l1_regularization_strength
+        self.l2_regularization_strength = l2_regularization_strength
+        self.l2_shrinkage_regularization_strength = (
+            l2_shrinkage_regularization_strength
+        )
+        self.beta = beta
+
+    def build(self, var_list):
+        """Initialize optimizer variables.
+
+        Args:
+            var_list: list of model variables to build Ftrl variables on.
+        """
+        if self.built:
+            return
+        super().build(var_list)
+        self._accumulators = []
+        self._linears = []
+        for var in var_list:
+            self._accumulators.append(
+                self.add_variable(
+                    shape=var.shape,
+                    dtype=var.dtype,
+                    name="accumulator",
+                    initializer=initializers.Constant(
+                        self.initial_accumulator_value,
+                    ),
+                )
+            )
+            self._linears.append(
+                self.add_variable_from_reference(
+                    reference_variable=var, name="linear"
+                )
+            )
+
+    def update_step(self, gradient, variable, learning_rate):
+        """Update step given gradient and the associated model variable."""
+
+        lr = ops.cast(learning_rate, variable.dtype)
+        gradient = ops.cast(gradient, variable.dtype)
+
+        accum = self._accumulators[self._get_variable_index(variable)]
+        linear = self._linears[self._get_variable_index(variable)]
+
+        lr_power = self.learning_rate_power
+        l2_reg = self.l2_regularization_strength
+        l2_reg = l2_reg + self.beta / (2.0 * lr)
+
+        grad_to_use = ops.add(
+            gradient,
+            ops.multiply(
+                2 * self.l2_shrinkage_regularization_strength, variable
+            ),
+        )
+        new_accum = ops.add(accum, ops.square(gradient))
+        self.assign_add(
+            linear,
+            ops.subtract(
+                grad_to_use,
+                ops.multiply(
+                    ops.divide(
+                        ops.subtract(
+                            ops.power(new_accum, -lr_power),
+                            ops.power(accum, -lr_power),
+                        ),
+                        lr,
+                    ),
+                    variable,
+                ),
+            ),
+        )
+        quadratic = ops.add(
+            ops.divide(ops.power(new_accum, (-lr_power)), lr), 2 * l2_reg
+        )
+        linear_clipped = ops.clip(
+            linear,
+            -self.l1_regularization_strength,
+            self.l1_regularization_strength,
+        )
+        self.assign(
+            variable,
+            ops.divide(ops.subtract(linear_clipped, linear), quadratic),
+        )
+        self.assign(accum, new_accum)
+
+    def get_config(self):
+        config = super().get_config()
+
+        config.update(
+            {
+                "learning_rate_power": self.learning_rate_power,
+                "initial_accumulator_value": self.initial_accumulator_value,
+                "l1_regularization_strength": self.l1_regularization_strength,
+                "l2_regularization_strength": self.l2_regularization_strength,
+                "l2_shrinkage_regularization_strength": self.l2_shrinkage_regularization_strength,  # noqa: E501
+                "beta": self.beta,
+            }
+        )
+        return config
+
+
+Ftrl.__doc__ = Ftrl.__doc__.replace(
+    "{{base_optimizer_keyword_args}}", optimizer.base_optimizer_keyword_args
+)

SwarmUI/dlbackend/ComfyUI/venv/lib/python3.10/site-packages/keras/src/optimizers/lamb.py

@@ -0,0 +1,158 @@
+from keras.src import ops
+from keras.src.api_export import keras_export
+from keras.src.optimizers import optimizer
+
+
+@keras_export("keras.optimizers.Lamb")
+class Lamb(optimizer.Optimizer):
+    """Optimizer that implements the Lamb algorithm.
+
+    Lamb is a stochastic gradient descent method that
+    uses layer-wise adaptive moments to adjusts the
+    learning rate for each parameter based on the ratio of the
+    norm of the weight to the norm of the gradient
+    This helps to stabilize the training process and improves convergence
+    especially for large batch sizes.
+
+    Args:
+        learning_rate: A float, a
+            `keras.optimizers.schedules.LearningRateSchedule` instance, or
+            a callable that takes no arguments and returns the actual value to
+            use. The learning rate. Defaults to `0.001`.
+        beta_1: A float value or a constant float tensor, or a callable
+            that takes no arguments and returns the actual value to use. The
+            exponential decay rate for the 1st moment estimates. Defaults to
+            `0.9`.
+        beta_2: A float value or a constant float tensor, or a callable
+            that takes no arguments and returns the actual value to use. The
+            exponential decay rate for the 2nd moment estimates. Defaults to
+            `0.999`.
+        epsilon: A small constant for numerical stability.
+            Defaults to `1e-7`.
+        {{base_optimizer_keyword_args}}
+
+    References:
+        - [Yang et al.](https://arxiv.org/pdf/1904.00962)
+    """
+
+    def __init__(
+        self,
+        learning_rate=0.001,
+        beta_1=0.9,
+        beta_2=0.999,
+        epsilon=1e-7,
+        weight_decay=None,
+        clipnorm=None,
+        clipvalue=None,
+        global_clipnorm=None,
+        use_ema=False,
+        ema_momentum=0.99,
+        ema_overwrite_frequency=None,
+        loss_scale_factor=None,
+        gradient_accumulation_steps=None,
+        name="lamb",
+        **kwargs,
+    ):
+        super().__init__(
+            learning_rate=learning_rate,
+            name=name,
+            weight_decay=weight_decay,
+            clipnorm=clipnorm,
+            clipvalue=clipvalue,
+            global_clipnorm=global_clipnorm,
+            use_ema=use_ema,
+            ema_momentum=ema_momentum,
+            ema_overwrite_frequency=ema_overwrite_frequency,
+            loss_scale_factor=loss_scale_factor,
+            gradient_accumulation_steps=gradient_accumulation_steps,
+            **kwargs,
+        )
+        self.beta_1 = beta_1
+        self.beta_2 = beta_2
+        self.epsilon = epsilon
+
+    def build(self, var_list):
+        """Initialize optimizer variables.
+
+        Lamb optimizer has 2 types of variables: momentums and velocities
+
+        Args:
+            var_list: list of model variables to build Lamb variables on.
+        """
+        if self.built:
+            return
+        super().build(var_list)
+        self._momentums = []
+        self._velocities = []
+        for var in var_list:
+            self._momentums.append(
+                self.add_variable_from_reference(
+                    reference_variable=var, name="momentum"
+                )
+            )
+            self._velocities.append(
+                self.add_variable_from_reference(
+                    reference_variable=var, name="velocity"
+                )
+            )
+
+    def update_step(self, gradient, variable, learning_rate):
+        """Update step given gradient and the associated model variable."""
+        lr = ops.cast(learning_rate, variable.dtype)
+        gradient = ops.cast(gradient, variable.dtype)
+        local_step = ops.cast(self.iterations + 1, variable.dtype)
+
+        beta_1_power = ops.power(
+            ops.cast(self.beta_1, variable.dtype), local_step
+        )
+        beta_2_power = ops.power(
+            ops.cast(self.beta_2, variable.dtype), local_step
+        )
+
+        m = self._momentums[self._get_variable_index(variable)]
+        v = self._velocities[self._get_variable_index(variable)]
+
+        self.assign_add(
+            m, ops.multiply(ops.subtract(gradient, m), 1 - self.beta_1)
+        )
+
+        self.assign_add(
+            v,
+            ops.multiply(
+                ops.subtract(ops.square(gradient), v), 1 - self.beta_2
+            ),
+        )
+
+        m_t_hat = ops.divide(m, (1.0 - beta_1_power))
+        v_sqrt = ops.add(
+            ops.sqrt(ops.divide(v, (1.0 - beta_2_power))), self.epsilon
+        )
+
+        update = ops.divide(m_t_hat, v_sqrt)
+        w_norm = ops.sqrt(ops.sum(ops.power(variable, 2)))
+        g_norm = ops.sqrt(ops.sum(ops.power(update, 2)))
+
+        # ratio = w_norm / g_norm if w_norm > 0 and g_norm > 0 else 1
+        ratio = ops.where(
+            ops.greater(w_norm, 0),
+            ops.where(ops.greater(g_norm, 0), (w_norm / g_norm), 1.0),
+            1.0,
+        )
+
+        self.assign_sub(variable, ratio * lr * update)
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "beta_1": self.beta_1,
+                "beta_2": self.beta_2,
+                "epsilon": self.epsilon,
+            }
+        )
+        return config
+
+
+Lamb.__doc__ = Lamb.__doc__.replace(
+    "{{base_optimizer_keyword_args}}", optimizer.base_optimizer_keyword_args
+)

SwarmUI/dlbackend/ComfyUI/venv/lib/python3.10/site-packages/keras/src/optimizers/lion.py

@@ -0,0 +1,142 @@
+from keras.src import ops
+from keras.src.api_export import keras_export
+from keras.src.optimizers import optimizer
+
+
+@keras_export(["keras.optimizers.Lion"])
+class Lion(optimizer.Optimizer):
+    """Optimizer that implements the Lion algorithm.
+
+    The Lion optimizer is a stochastic-gradient-descent method that uses the
+    sign operator to control the magnitude of the update, unlike other adaptive
+    optimizers such as Adam that rely on second-order moments. This make
+    Lion more memory-efficient as it only keeps track of the momentum. According
+    to the authors (see reference), its performance gain over Adam grows with
+    the batch size. Because the update of Lion is produced through the sign
+    operation, resulting in a larger norm, a suitable learning rate for Lion is
+    typically 3-10x smaller than that for AdamW. The weight decay for Lion
+    should be in turn 3-10x larger than that for AdamW to maintain a
+    similar strength (lr * wd).
+
+    Args:
+        learning_rate: A float, a
+            `keras.optimizers.schedules.LearningRateSchedule` instance, or
+            a callable that takes no arguments and returns the actual value to
+            use. The learning rate. Defaults to `0.001`.
+        beta_1: A float value or a constant float tensor, or a callable
+            that takes no arguments and returns the actual value to use. The
+            rate to combine the current gradient and the 1st moment estimate.
+            Defaults to `0.9`.
+        beta_2: A float value or a constant float tensor, or a callable
+            that takes no arguments and returns the actual value to use. The
+            exponential decay rate for the 1st moment estimate. Defaults to
+            `0.99`.
+        {{base_optimizer_keyword_args}}
+
+    References:
+
+    - [Chen et al., 2023](http://arxiv.org/abs/2302.06675)
+    - [Authors' implementation](
+        http://github.com/google/automl/tree/master/lion)
+
+    """
+
+    def __init__(
+        self,
+        learning_rate=0.001,
+        beta_1=0.9,
+        beta_2=0.99,
+        weight_decay=None,
+        clipnorm=None,
+        clipvalue=None,
+        global_clipnorm=None,
+        use_ema=False,
+        ema_momentum=0.99,
+        ema_overwrite_frequency=None,
+        loss_scale_factor=None,
+        gradient_accumulation_steps=None,
+        name="lion",
+        **kwargs,
+    ):
+        super().__init__(
+            learning_rate=learning_rate,
+            name=name,
+            weight_decay=weight_decay,
+            clipnorm=clipnorm,
+            clipvalue=clipvalue,
+            global_clipnorm=global_clipnorm,
+            use_ema=use_ema,
+            ema_momentum=ema_momentum,
+            ema_overwrite_frequency=ema_overwrite_frequency,
+            loss_scale_factor=loss_scale_factor,
+            gradient_accumulation_steps=gradient_accumulation_steps,
+            **kwargs,
+        )
+        self.beta_1 = beta_1
+        self.beta_2 = beta_2
+        if beta_1 <= 0 or beta_1 > 1:
+            raise ValueError(
+                "Argument `beta_1` must be in the [0, 1] range. Otherwise, the "
+                f"optimizer degenerates to SignSGD. Received: beta_1={beta_1}."
+            )
+
+    def build(self, var_list):
+        """Initialize optimizer variables.
+
+        Lion optimizer has one variable `momentums`.
+
+        Args:
+            var_list: list of model variables to build Lion variables on.
+        """
+        if self.built:
+            return
+        super().build(var_list)
+        self._momentums = []
+        for var in var_list:
+            self._momentums.append(
+                self.add_variable_from_reference(
+                    reference_variable=var, name="momentum"
+                )
+            )
+
+    def update_step(self, gradient, variable, learning_rate):
+        """Update step given gradient and the associated model variable."""
+        lr = ops.cast(learning_rate, variable.dtype)
+        gradient = ops.cast(gradient, variable.dtype)
+        beta_1 = ops.cast(self.beta_1, variable.dtype)
+        beta_2 = ops.cast(self.beta_2, variable.dtype)
+        m = self._momentums[self._get_variable_index(variable)]
+
+        self.assign_sub(
+            variable,
+            ops.multiply(
+                lr,
+                ops.sign(
+                    ops.add(
+                        ops.multiply(m, beta_1),
+                        ops.multiply(gradient, (1.0 - beta_1)),
+                    )
+                ),
+            ),
+        )
+        self.assign(
+            m,
+            ops.add(
+                ops.multiply(m, beta_2), ops.multiply(gradient, (1.0 - beta_2))
+            ),
+        )
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "beta_1": self.beta_1,
+                "beta_2": self.beta_2,
+            }
+        )
+        return config
+
+
+Lion.__doc__ = Lion.__doc__.replace(
+    "{{base_optimizer_keyword_args}}", optimizer.base_optimizer_keyword_args
+)

SwarmUI/dlbackend/ComfyUI/venv/lib/python3.10/site-packages/keras/src/optimizers/loss_scale_optimizer.py

@@ -0,0 +1,298 @@
+from keras.src import backend
+from keras.src import initializers
+from keras.src import ops
+from keras.src.api_export import keras_export
+from keras.src.optimizers import optimizer
+from keras.src.saving import serialization_lib
+from keras.src.utils import tracking
+
+
+@keras_export(
+    [
+        "keras.optimizers.LossScaleOptimizer",
+        "keras.mixed_precision.LossScaleOptimizer",
+    ]
+)
+class LossScaleOptimizer(optimizer.Optimizer):
+    """An optimizer that dynamically scales the loss to prevent underflow.
+
+    Loss scaling is a technique to prevent numeric underflow in intermediate
+    gradients when float16 is used. To prevent underflow, the loss is multiplied
+    (or "scaled") by a certain factor called the "loss scale", which causes
+    intermediate gradients to be scaled by the loss scale as well. The final
+    gradients are divided (or "unscaled") by the loss scale to bring them back
+    to their original value.
+
+    `LossScaleOptimizer` wraps another optimizer and applies dynamic loss
+    scaling to it. This loss scale is dynamically updated over time as follows:
+    - On any train step, if a nonfinite gradient is encountered, the loss scale
+      is halved, and the train step is skipped.
+    - If `dynamic_growth_steps` have occurred since the last time the loss scale
+      was updated, and no nonfinite gradients have occurred, the loss scale
+      is doubled.
+
+    Args:
+        inner_optimizer: The `keras.optimizers.Optimizer` instance to wrap.
+        initial_scale: Float. The initial loss scale. This scale will be updated
+            during training. It is recommended for this to be a very high
+            number, because a loss scale that is too high gets lowered far more
+            quickly than a loss scale that is too low gets raised.
+        dynamic_growth_steps: Int. How often to update the scale upwards. After
+            every `dynamic_growth_steps` steps with finite gradients, the
+            loss scale is doubled.
+        {{base_optimizer_keyword_args}}
+    """
+
+    def __init__(
+        self,
+        inner_optimizer,
+        initial_scale=2.0**15,
+        dynamic_growth_steps=2000,
+        **kwargs,
+    ):
+        if not kwargs.pop("dynamic", True):
+            raise ValueError(
+                "LossScaleOptimizer no longer supports `dynamic=False`. "
+                "Instead, simply set `loss_scale_factor` directly on the "
+                "`inner_optimizer`."
+            )
+        super().__init__(learning_rate=0.0, **kwargs)
+        self.inner_optimizer = inner_optimizer
+        self.initial_scale = initial_scale
+        self.dynamic_growth_steps = dynamic_growth_steps
+
+    @tracking.no_automatic_dependency_tracking
+    def build(self, var_list):
+        self.step_counter = self.add_variable(
+            shape=(),
+            dtype="int",
+            initializer=initializers.Zeros(),
+            aggregation="none",
+            name="step_counter",
+        )
+        self.dynamic_scale = self.add_variable(
+            shape=(),
+            dtype="float32",
+            initializer=initializers.Constant(self.initial_scale),
+            aggregation="none",
+            name="dynamic_scale",
+        )
+        self.inner_optimizer.build(var_list)
+        self.built = True
+
+    @property
+    def variables(self):
+        return self._variables + self.inner_optimizer.variables
+
+    def stateless_apply(self, optimizer_variables, grads, trainable_variables):
+        if not self.built:
+            raise ValueError(
+                f"To call `stateless_apply`, {self.__class__.__name__} "
+                "must be built (i.e. its variables must have been created). "
+                "You can build it via `optimizer.build(trainable_variables)`."
+            )
+        finite = self.check_finite(grads)
+        return ops.cond(
+            finite,
+            lambda: self._stateless_handle_finite_grads(
+                optimizer_variables, grads, trainable_variables
+            ),
+            lambda: self._stateless_handle_non_finite_grads(
+                optimizer_variables, trainable_variables
+            ),
+        )
+
+    def _stateless_handle_finite_grads(
+        self, optimizer_variables, grads, trainable_variables
+    ):
+        def upscale():
+            mapping = list(zip(self.variables, optimizer_variables))
+            with backend.StatelessScope(state_mapping=mapping) as scope:
+                self.step_counter.assign(0)
+                self.dynamic_scale.assign(self.dynamic_scale * 2.0)
+            return [scope.get_current_value(v) for v in self._variables]
+
+        def increment():
+            mapping = list(zip(self.variables, optimizer_variables))
+            with backend.StatelessScope(state_mapping=mapping) as scope:
+                self.step_counter.assign_add(1)
+            return [scope.get_current_value(v) for v in self._variables]
+
+        mapping = list(zip(self.variables, optimizer_variables))
+        with backend.StatelessScope(state_mapping=mapping):
+            # Potentially upscale loss and reset counter.
+            own_variables = ops.cond(
+                ops.equal(self.step_counter, self.dynamic_growth_steps - 1),
+                upscale,
+                increment,
+            )
+
+            # Unscale gradients.
+            scale = self.dynamic_scale
+            unscaled_grads = [
+                g if g is None else ops.divide(g, scale) for g in grads
+            ]
+            (
+                new_trainable_variables,
+                new_inner_variables,
+            ) = self.inner_optimizer.stateless_apply(
+                self.inner_optimizer.variables,
+                unscaled_grads,
+                trainable_variables,
+            )
+
+        new_optimizer_variables = own_variables + new_inner_variables
+        return new_trainable_variables, new_optimizer_variables
+
+    def _stateless_handle_non_finite_grads(
+        self, optimizer_variables, trainable_variables
+    ):
+        mapping = list(zip(self.variables, optimizer_variables))
+        with backend.StatelessScope(state_mapping=mapping) as scope:
+            self.step_counter.assign(0)
+            self.dynamic_scale.assign(self.dynamic_scale / 2.0)
+        new_optimizer_variables = []
+        for v in self.variables:
+            new_optimizer_variables.append(scope.get_current_value(v))
+        return trainable_variables, new_optimizer_variables
+
+    def apply(self, grads, trainable_variables=None):
+        # Optionally build optimizer.
+        if not self.built:
+            with backend.name_scope(self.name, caller=self):
+                self.build(trainable_variables)
+            self.built = True
+
+        if backend.backend() == "tensorflow":
+            self._tf_apply(grads, trainable_variables)
+        else:
+            self._common_apply(grads, trainable_variables)
+
+    def _stateful_handle_finite_grads(self, grads, trainable_variables):
+        scale = self.dynamic_scale
+        # Unscale gradients.
+        unscaled_grads = [
+            g if g is None else ops.divide(g, scale) for g in grads
+        ]
+        self.inner_optimizer.apply(
+            unscaled_grads, trainable_variables=trainable_variables
+        )
+
+        def upscale():
+            self.step_counter.assign(0)
+            self.dynamic_scale.assign(self.dynamic_scale * 2.0)
+
+        def increment():
+            self.step_counter.assign_add(1)
+
+        # Potentially upscale loss and reset counter.
+        ops.cond(
+            ops.equal(self.step_counter, self.dynamic_growth_steps - 1),
+            upscale,
+            increment,
+        )
+
+    def _stateful_handle_non_finite_grads(self):
+        # If any inf or nan in grads, downscale loss and reset counter.
+        self.step_counter.assign(0)
+        self.dynamic_scale.assign(self.dynamic_scale / 2.0)
+
+    def _common_apply(self, grads, trainable_variables=None):
+        finite = self.check_finite(grads)
+        ops.cond(
+            finite,
+            lambda: self._stateful_handle_finite_grads(
+                grads, trainable_variables
+            ),
+            self._stateful_handle_non_finite_grads,
+        )
+
+    def _tf_apply(self, grads, trainable_variables=None):
+        """Tensorflow specific logic for apply, which handles distribution."""
+        from keras.src.utils.module_utils import tensorflow as tf
+
+        if tf.distribute.in_cross_replica_context():
+            raise ValueError("apply() must be called in a replica context.")
+
+        if tf.__internal__.distribute.strategy_supports_no_merge_call():
+            self._common_apply(grads, trainable_variables=trainable_variables)
+        else:
+
+            def _handle_cross_replica(distribution, grads, trainable_variables):
+                finite_per_replica = (
+                    distribution.extended.call_for_each_replica(
+                        self.check_finite, args=(grads,)
+                    )
+                )
+                # Each replica computed the same `finite` value, since
+                # `grads` is all-reduced across replicas. Arbitrarily take
+                # `finite` from the first replica.
+                finite = distribution.experimental_local_results(
+                    finite_per_replica
+                )[0]
+
+                def apply_fn():
+                    distribution.extended.call_for_each_replica(
+                        self._stateful_handle_finite_grads,
+                        args=(grads, trainable_variables),
+                    )
+
+                # Note: We must call this cond() in a cross-replica context.
+                # DistributionStrategy does not support having a cond in a
+                # replica context with a branch that calls `merge_call`, and
+                # self._optimizer.apply_gradients calls `merge_call`.
+                ops.cond(
+                    finite, apply_fn, self._stateful_handle_non_finite_grads
+                )
+
+            tf.distribute.get_replica_context().merge_call(
+                _handle_cross_replica, args=(grads, trainable_variables)
+            )
+
+    def check_finite(self, grads):
+        tensor_grads = [g for g in grads if g is not None]
+        finite_grads = [ops.all(ops.isfinite(g)) for g in tensor_grads]
+        return ops.all(ops.convert_to_tensor(finite_grads))
+
+    @property
+    def learning_rate(self):
+        return self.inner_optimizer.learning_rate
+
+    @learning_rate.setter
+    def learning_rate(self, learning_rate):
+        self.inner_optimizer.learning_rate = learning_rate
+
+    def scale_loss(self, loss):
+        scale = self.dynamic_scale if self.built else self.initial_scale
+        return loss * scale
+
+    def finalize_variable_values(self, var_list):
+        self.inner_optimizer.finalize_variable_values(var_list)
+
+    def get_config(self):
+        config = super().get_config()
+        inner_optimizer_config = serialization_lib.serialize_keras_object(
+            self.inner_optimizer
+        )
+        config.update(
+            {
+                "inner_optimizer": inner_optimizer_config,
+                "initial_scale": self.initial_scale,
+                "dynamic_growth_steps": self.dynamic_growth_steps,
+            }
+        )
+        del config["learning_rate"]
+        return config
+
+    @classmethod
+    def from_config(cls, config, custom_objects=None):
+        inner_optimizer = serialization_lib.deserialize_keras_object(
+            config.pop("inner_optimizer"),
+            custom_objects=custom_objects,
+        )
+        return cls(inner_optimizer, **config)
+
+
+LossScaleOptimizer.__doc__ = LossScaleOptimizer.__doc__.replace(
+    "{{base_optimizer_keyword_args}}", optimizer.base_optimizer_keyword_args
+)

SwarmUI/dlbackend/ComfyUI/venv/lib/python3.10/site-packages/keras/src/optimizers/nadam.py

@@ -0,0 +1,174 @@
+from keras.src import backend
+from keras.src import ops
+from keras.src.api_export import keras_export
+from keras.src.optimizers import optimizer
+
+
+@keras_export(["keras.optimizers.Nadam"])
+class Nadam(optimizer.Optimizer):
+    """Optimizer that implements the Nadam algorithm.
+
+    Much like Adam is essentially RMSprop with momentum, Nadam is Adam with
+    Nesterov momentum.
+
+    Args:
+        learning_rate: A float, a
+            `keras.optimizers.schedules.LearningRateSchedule` instance, or
+            a callable that takes no arguments and returns the actual value to
+            use. The learning rate. Defaults to `0.001`.
+        beta_1: A float value or a constant float tensor, or a callable
+            that takes no arguments and returns the actual value to use. The
+            exponential decay rate for the 1st moment estimates.
+            Defaults to `0.9`.
+        beta_2: A float value or a constant float tensor, or a callable
+            that takes no arguments and returns the actual value to use. The
+            exponential decay rate for the 2nd moment estimates. Defaults to
+            `0.999`.
+        epsilon: A small constant for numerical stability. This epsilon is
+            "epsilon hat" in the Kingma and Ba paper (in the formula just before
+            Section 2.1), not the epsilon in Algorithm 1 of the paper.
+            Defaults to `1e-7`.
+        {{base_optimizer_keyword_args}}
+
+    Reference:
+
+    - [Dozat, 2015](http://cs229.stanford.edu/proj2015/054_report.pdf).
+
+    """
+
+    def __init__(
+        self,
+        learning_rate=0.001,
+        beta_1=0.9,
+        beta_2=0.999,
+        epsilon=1e-7,
+        weight_decay=None,
+        clipnorm=None,
+        clipvalue=None,
+        global_clipnorm=None,
+        use_ema=False,
+        ema_momentum=0.99,
+        ema_overwrite_frequency=None,
+        loss_scale_factor=None,
+        gradient_accumulation_steps=None,
+        name="nadam",
+        **kwargs,
+    ):
+        super().__init__(
+            learning_rate=learning_rate,
+            name=name,
+            weight_decay=weight_decay,
+            clipnorm=clipnorm,
+            clipvalue=clipvalue,
+            global_clipnorm=global_clipnorm,
+            use_ema=use_ema,
+            ema_momentum=ema_momentum,
+            ema_overwrite_frequency=ema_overwrite_frequency,
+            loss_scale_factor=loss_scale_factor,
+            gradient_accumulation_steps=gradient_accumulation_steps,
+            **kwargs,
+        )
+        self.beta_1 = beta_1
+        self.beta_2 = beta_2
+        self.epsilon = epsilon
+
+    def build(self, var_list):
+        """Initialize optimizer variables.
+
+        Nadam optimizer has 2 types of variables: momentums and velocities.
+
+        Args:
+            var_list: list of model variables to build Nadam variables on.
+        """
+        if self.built:
+            return
+        if var_list:
+            dtype = var_list[0].dtype
+        else:
+            dtype = backend.floatx()
+        super().build(var_list)
+        self._momentums = []
+        self._velocities = []
+        self._u_product = backend.Variable(1.0, dtype=dtype)
+
+        for var in var_list:
+            self._momentums.append(
+                self.add_variable_from_reference(
+                    reference_variable=var, name="momentum"
+                )
+            )
+            self._velocities.append(
+                self.add_variable_from_reference(
+                    reference_variable=var, name="velocity"
+                )
+            )
+
+    def _backend_update_step(self, grads, trainable_variables, learning_rate):
+        dtype = self._u_product.dtype
+        self.assign(
+            self._u_product,
+            self._u_product
+            * self.beta_1
+            * (
+                1.0
+                - 0.5 * ops.power(0.96, ops.cast(self.iterations + 1, dtype))
+            ),
+        )
+        super()._backend_update_step(grads, trainable_variables, learning_rate)
+
+    def update_step(self, gradient, variable, learning_rate):
+        """Update step given gradient and the associated model variable."""
+        var_dtype = variable.dtype
+        lr = ops.cast(learning_rate, var_dtype)
+        gradient = ops.cast(gradient, var_dtype)
+
+        local_step = ops.cast(self.iterations + 1, var_dtype)
+        next_step = ops.cast(self.iterations + 2, var_dtype)
+        decay = ops.cast(0.96, var_dtype)
+        beta_1 = ops.cast(self.beta_1, var_dtype)
+        beta_2 = ops.cast(self.beta_2, var_dtype)
+        u_t = beta_1 * (1.0 - 0.5 * (ops.power(decay, local_step)))
+        u_t_1 = beta_1 * (1.0 - 0.5 * (ops.power(decay, next_step)))
+        u_product_t = ops.cast(self._u_product, var_dtype)
+
+        u_product_t_1 = u_product_t * u_t_1
+        beta_2_power = ops.power(beta_2, local_step)
+
+        m = self._momentums[self._get_variable_index(variable)]
+        v = self._velocities[self._get_variable_index(variable)]
+
+        self.assign_add(
+            m, ops.multiply(ops.subtract(gradient, m), (1 - beta_1))
+        )
+        self.assign_add(
+            v, ops.multiply(ops.subtract(ops.square(gradient), v), (1 - beta_2))
+        )
+        m_hat = ops.add(
+            ops.divide(ops.multiply(u_t_1, m), 1 - u_product_t_1),
+            ops.divide(ops.multiply(1 - u_t, gradient), 1 - u_product_t),
+        )
+        v_hat = ops.divide(v, (1 - beta_2_power))
+
+        self.assign_sub(
+            variable,
+            ops.divide(
+                ops.multiply(m_hat, lr), ops.add(ops.sqrt(v_hat), self.epsilon)
+            ),
+        )
+
+    def get_config(self):
+        config = super().get_config()
+
+        config.update(
+            {
+                "beta_1": self.beta_1,
+                "beta_2": self.beta_2,
+                "epsilon": self.epsilon,
+            }
+        )
+        return config
+
+
+Nadam.__doc__ = Nadam.__doc__.replace(
+    "{{base_optimizer_keyword_args}}", optimizer.base_optimizer_keyword_args
+)

SwarmUI/dlbackend/ComfyUI/venv/lib/python3.10/site-packages/keras/src/optimizers/optimizer.py

@@ -0,0 +1,27 @@
+from keras.src import backend
+from keras.src.api_export import keras_export
+from keras.src.optimizers import base_optimizer
+
+if backend.backend() == "tensorflow":
+    from keras.src.backend.tensorflow.optimizer import (
+        TFOptimizer as BackendOptimizer,
+    )
+elif backend.backend() == "torch":
+    from keras.src.backend.torch.optimizers import (
+        TorchOptimizer as BackendOptimizer,
+    )
+elif backend.backend() == "jax":
+    from keras.src.backend.jax.optimizer import JaxOptimizer as BackendOptimizer
+else:
+
+    class BackendOptimizer(base_optimizer.BaseOptimizer):
+        pass
+
+
+@keras_export(["keras.Optimizer", "keras.optimizers.Optimizer"])
+class Optimizer(BackendOptimizer, base_optimizer.BaseOptimizer):
+    pass
+
+
+Optimizer.__doc__ = base_optimizer.BaseOptimizer.__doc__
+base_optimizer_keyword_args = base_optimizer.base_optimizer_keyword_args

SwarmUI/dlbackend/ComfyUI/venv/lib/python3.10/site-packages/keras/src/optimizers/rmsprop.py

@@ -0,0 +1,180 @@
+from keras.src import ops
+from keras.src.api_export import keras_export
+from keras.src.optimizers import optimizer
+
+
+@keras_export(["keras.optimizers.RMSprop"])
+class RMSprop(optimizer.Optimizer):
+    """Optimizer that implements the RMSprop algorithm.
+
+    The gist of RMSprop is to:
+
+    - Maintain a moving (discounted) average of the square of gradients
+    - Divide the gradient by the root of this average
+
+    This implementation of RMSprop uses plain momentum, not Nesterov momentum.
+
+    The centered version additionally maintains a moving average of the
+    gradients, and uses that average to estimate the variance.
+
+    Args:
+        learning_rate: A float, a
+            `keras.optimizers.schedules.LearningRateSchedule` instance, or
+            a callable that takes no arguments and returns the actual value to
+            use. The learning rate. Defaults to `0.001`.
+        rho: float, defaults to 0.9. Discounting factor for the old gradients.
+        momentum: float, defaults to 0.0. If not 0.0., the optimizer tracks the
+            momentum value, with a decay rate equals to `1 - momentum`.
+        epsilon: A small constant for numerical stability. This epsilon is
+            "epsilon hat" in the Kingma and Ba paper (in the formula just before
+            Section 2.1), not the epsilon in Algorithm 1 of the paper. Defaults
+            to 1e-7.
+        centered: Boolean. If `True`, gradients are normalized by the estimated
+            variance of the gradient; if False, by the uncentered second moment.
+            Setting this to `True` may help with training, but is slightly more
+            expensive in terms of computation and memory. Defaults to `False`.
+        {{base_optimizer_keyword_args}}
+
+    Example:
+
+    >>> opt = keras.optimizers.RMSprop(learning_rate=0.1)
+    >>> var1 = keras.backend.Variable(10.0)
+    >>> loss = lambda: (var1 ** 2) / 2.0  # d(loss) / d(var1) = var1
+    >>> opt.minimize(loss, [var1])
+    >>> var1
+    9.683772
+
+    Reference:
+
+    - [Hinton, 2012](
+        http://www.cs.toronto.edu/~tijmen/csc321/slides/lecture_slides_lec6.pdf)
+    """
+
+    def __init__(
+        self,
+        learning_rate=0.001,
+        rho=0.9,
+        momentum=0.0,
+        epsilon=1e-7,
+        centered=False,
+        weight_decay=None,
+        clipnorm=None,
+        clipvalue=None,
+        global_clipnorm=None,
+        use_ema=False,
+        ema_momentum=0.99,
+        ema_overwrite_frequency=None,
+        loss_scale_factor=None,
+        gradient_accumulation_steps=None,
+        name="rmsprop",
+        **kwargs,
+    ):
+        super().__init__(
+            learning_rate=learning_rate,
+            weight_decay=weight_decay,
+            clipnorm=clipnorm,
+            clipvalue=clipvalue,
+            global_clipnorm=global_clipnorm,
+            use_ema=use_ema,
+            ema_momentum=ema_momentum,
+            ema_overwrite_frequency=ema_overwrite_frequency,
+            loss_scale_factor=loss_scale_factor,
+            gradient_accumulation_steps=gradient_accumulation_steps,
+            name=name,
+            **kwargs,
+        )
+        self.rho = rho
+        self.momentum = momentum
+        self.epsilon = epsilon
+        self.centered = centered
+
+    def build(self, var_list):
+        if self.built:
+            return
+
+        super().build(var_list)
+
+        self._velocities = []
+        for var in var_list:
+            self._velocities.append(
+                self.add_variable_from_reference(var, "velocity")
+            )
+
+        self._momentums = []
+        if self.momentum > 0:
+            for var in var_list:
+                self._momentums.append(
+                    self.add_variable_from_reference(var, "momentum")
+                )
+
+        self._average_gradients = []
+        if self.centered:
+            for var in var_list:
+                self._average_gradients.append(
+                    self.add_variable_from_reference(var, "average_gradient")
+                )
+
+    def update_step(self, gradient, variable, learning_rate):
+        """Update step given gradient and the associated model variable."""
+        lr = ops.cast(learning_rate, variable.dtype)
+        gradient = ops.cast(gradient, variable.dtype)
+
+        velocity = self._velocities[self._get_variable_index(variable)]
+        momentum = None
+        if self.momentum > 0:
+            momentum = self._momentums[self._get_variable_index(variable)]
+        average_grad = None
+        if self.centered:
+            average_grad = self._average_gradients[
+                self._get_variable_index(variable)
+            ]
+
+        rho = self.rho
+
+        self.assign(
+            velocity,
+            ops.add(
+                ops.multiply(rho, velocity),
+                ops.multiply(1 - rho, ops.square(gradient)),
+            ),
+        )
+        if self.centered:
+            self.assign(
+                average_grad,
+                ops.add(
+                    ops.multiply(rho, average_grad),
+                    ops.multiply(1 - rho, gradient),
+                ),
+            )
+            denominator = velocity - ops.square(average_grad) + self.epsilon
+        else:
+            denominator = ops.add(velocity, self.epsilon)
+        increment = ops.divide(
+            ops.multiply(lr, gradient), ops.sqrt(denominator)
+        )
+        if self.momentum > 0:
+            self.assign(
+                momentum,
+                ops.add(ops.multiply(self.momentum, momentum), increment),
+            )
+            self.assign_sub(variable, momentum)
+        else:
+            self.assign_sub(variable, increment)
+
+    def get_config(self):
+        config = super().get_config()
+
+        config.update(
+            {
+                "rho": self.rho,
+                "momentum": self.momentum,
+                "epsilon": self.epsilon,
+                "centered": self.centered,
+            }
+        )
+        return config
+
+
+RMSprop.__doc__ = RMSprop.__doc__.replace(
+    "{{base_optimizer_keyword_args}}", optimizer.base_optimizer_keyword_args
+)

SwarmUI/dlbackend/ComfyUI/venv/lib/python3.10/site-packages/keras/src/optimizers/schedules/__init__.py

@@ -0,0 +1,16 @@
+from keras.src.optimizers.schedules.learning_rate_schedule import CosineDecay
+from keras.src.optimizers.schedules.learning_rate_schedule import (
+    CosineDecayRestarts,
+)
+from keras.src.optimizers.schedules.learning_rate_schedule import (
+    ExponentialDecay,
+)
+from keras.src.optimizers.schedules.learning_rate_schedule import (
+    InverseTimeDecay,
+)
+from keras.src.optimizers.schedules.learning_rate_schedule import (
+    PiecewiseConstantDecay,
+)
+from keras.src.optimizers.schedules.learning_rate_schedule import (
+    PolynomialDecay,
+)

SwarmUI/dlbackend/ComfyUI/venv/lib/python3.10/site-packages/keras/src/optimizers/schedules/learning_rate_schedule.py

@@ -0,0 +1,969 @@
+"""Various learning rate schedule functions."""
+
+import math
+
+from keras.src import ops
+from keras.src.api_export import keras_export
+from keras.src.saving import serialization_lib
+
+
+@keras_export("keras.optimizers.schedules.LearningRateSchedule")
+class LearningRateSchedule:
+    """The learning rate schedule base class.
+
+    You can use a learning rate schedule to modulate how the learning rate
+    of your optimizer changes over time.
+
+    Several built-in learning rate schedules are available, such as
+    `keras.optimizers.schedules.ExponentialDecay` or
+    `keras.optimizers.schedules.PiecewiseConstantDecay`:
+
+    ```python
+    lr_schedule = keras.optimizers.schedules.ExponentialDecay(
+        initial_learning_rate=1e-2,
+        decay_steps=10000,
+        decay_rate=0.9)
+    optimizer = keras.optimizers.SGD(learning_rate=lr_schedule)
+    ```
+
+    A `LearningRateSchedule` instance can be passed in as the `learning_rate`
+    argument of any optimizer.
+
+    To implement your own schedule object, you should implement the `__call__`
+    method, which takes a `step` argument (scalar integer tensor, the
+    current training step count).
+    Like for any other Keras object, you can also optionally
+    make your object serializable by implementing the `get_config`
+    and `from_config` methods.
+
+    Example:
+
+    ```python
+    class MyLRSchedule(keras.optimizers.schedules.LearningRateSchedule):
+
+        def __init__(self, initial_learning_rate):
+            self.initial_learning_rate = initial_learning_rate
+
+        def __call__(self, step):
+            return self.initial_learning_rate / (step + 1)
+
+    optimizer = keras.optimizers.SGD(learning_rate=MyLRSchedule(0.1))
+    ```
+    """
+
+    def __call__(self, step):
+        raise NotImplementedError(
+            f"Learning rate schedule '{self.__class__.__name__}' "
+            "must override `__call__(self, step)`."
+        )
+
+    def get_config(self):
+        raise NotImplementedError(
+            f"Learning rate schedule '{self.__class__.__name__}' "
+            "must override `get_config()` in order to be serializable."
+        )
+
+    @classmethod
+    def from_config(cls, config):
+        """Instantiates a `LearningRateSchedule` from its config.
+
+        Args:
+            config: Output of `get_config()`.
+
+        Returns:
+            A `LearningRateSchedule` instance.
+        """
+        return cls(**config)
+
+
+@keras_export("keras.optimizers.schedules.ExponentialDecay")
+class ExponentialDecay(LearningRateSchedule):
+    """A `LearningRateSchedule` that uses an exponential decay schedule.
+
+    When training a model, it is often useful to lower the learning rate as
+    the training progresses. This schedule applies an exponential decay function
+    to an optimizer step, given a provided initial learning rate.
+
+    The schedule is a 1-arg callable that produces a decayed learning
+    rate when passed the current optimizer step. This can be useful for changing
+    the learning rate value across different invocations of optimizer functions.
+    It is computed as:
+
+    ```python
+    def decayed_learning_rate(step):
+        return initial_learning_rate * decay_rate ^ (step / decay_steps)
+    ```
+
+    If the argument `staircase` is `True`, then `step / decay_steps` is
+    an integer division and the decayed learning rate follows a
+    staircase function.
+
+    You can pass this schedule directly into a `keras.optimizers.Optimizer`
+    as the learning rate.
+    Example: When fitting a Keras model, decay every 100000 steps with a base
+    of 0.96:
+
+    ```python
+    initial_learning_rate = 0.1
+    lr_schedule = keras.optimizers.schedules.ExponentialDecay(
+        initial_learning_rate,
+        decay_steps=100000,
+        decay_rate=0.96,
+        staircase=True)
+
+    model.compile(optimizer=keras.optimizers.SGD(learning_rate=lr_schedule),
+                  loss='sparse_categorical_crossentropy',
+                  metrics=['accuracy'])
+
+    model.fit(data, labels, epochs=5)
+    ```
+
+    The learning rate schedule is also serializable and deserializable using
+    `keras.optimizers.schedules.serialize` and
+    `keras.optimizers.schedules.deserialize`.
+
+    Args:
+        initial_learning_rate: A Python float. The initial learning rate.
+        decay_steps: A Python integer. Must be positive. See the decay
+            computation above.
+        decay_rate: A Python float. The decay rate.
+        staircase: Boolean.  If `True` decay the learning rate at discrete
+            intervals.
+        name: String.  Optional name of the operation.  Defaults to
+            `"ExponentialDecay`".
+
+    Returns:
+        A 1-arg callable learning rate schedule that takes the current optimizer
+        step and outputs the decayed learning rate, a scalar tensor of the
+        same type as `initial_learning_rate`.
+    """
+
+    def __init__(
+        self,
+        initial_learning_rate,
+        decay_steps,
+        decay_rate,
+        staircase=False,
+        name="ExponentialDecay",
+    ):
+        super().__init__()
+        self.initial_learning_rate = initial_learning_rate
+        self.decay_steps = decay_steps
+        self.decay_rate = decay_rate
+        self.staircase = staircase
+        self.name = name
+
+        if self.decay_steps <= 0:
+            raise ValueError(
+                "Argument `decay_steps` must be > 0. "
+                f"Received: decay_steps={self.decay_steps}"
+            )
+
+    def __call__(self, step):
+        with ops.name_scope(self.name):
+            initial_learning_rate = ops.convert_to_tensor(
+                self.initial_learning_rate
+            )
+            dtype = initial_learning_rate.dtype
+            decay_steps = ops.cast(self.decay_steps, dtype)
+            decay_rate = ops.cast(self.decay_rate, dtype)
+
+            global_step_recomp = ops.cast(step, dtype)
+            p = global_step_recomp / decay_steps
+            if self.staircase:
+                p = ops.floor(p)
+            return ops.multiply(initial_learning_rate, ops.power(decay_rate, p))
+
+    def get_config(self):
+        return {
+            "initial_learning_rate": self.initial_learning_rate,
+            "decay_steps": self.decay_steps,
+            "decay_rate": self.decay_rate,
+            "staircase": self.staircase,
+            "name": self.name,
+        }
+
+
+@keras_export("keras.optimizers.schedules.PiecewiseConstantDecay")
+class PiecewiseConstantDecay(LearningRateSchedule):
+    """A `LearningRateSchedule` that uses a piecewise constant decay schedule.
+
+    The function returns a 1-arg callable to compute the piecewise constant
+    when passed the current optimizer step. This can be useful for changing the
+    learning rate value across different invocations of optimizer functions.
+
+    Example: use a learning rate that's 1.0 for the first 100001 steps, 0.5
+        for the next 10000 steps, and 0.1 for any additional steps.
+
+    ```python
+    step = ops.array(0)
+    boundaries = [100000, 110000]
+    values = [1.0, 0.5, 0.1]
+    learning_rate_fn = keras.optimizers.schedules.PiecewiseConstantDecay(
+        boundaries, values)
+
+    # Later, whenever we perform an optimization step, we pass in the step.
+    learning_rate = learning_rate_fn(step)
+    ```
+
+    You can pass this schedule directly into a `keras.optimizers.Optimizer`
+    as the learning rate. The learning rate schedule is also serializable and
+    deserializable using `keras.optimizers.schedules.serialize` and
+    `keras.optimizers.schedules.deserialize`.
+
+    Args:
+        boundaries: A list of Python numbers with strictly increasing
+            entries, and with all elements having the same type as the
+            optimizer step.
+        values: A list of Python numbers that specifies the values for the
+            intervals defined by `boundaries`. It should have one more
+            element than `boundaries`, and all elements should have the same
+            type.
+        name: A string. Optional name of the operation. Defaults to
+            `"PiecewiseConstant"`.
+
+    Returns:
+        A 1-arg callable learning rate schedule that takes the current optimizer
+        step and outputs the decayed learning rate, a scalar tensor of the
+        same type as the boundary tensors.
+
+        The output of the 1-arg function that takes the `step`
+        is `values[0]` when `step <= boundaries[0]`,
+        `values[1]` when `step > boundaries[0]` and `step <= boundaries[1]`,
+        ..., and `values[-1]` when `step > boundaries[-1]`.
+
+
+    Raises:
+        ValueError: if the number of elements in the `boundaries` and `values`
+        lists do not match.
+    """
+
+    def __init__(self, boundaries, values, name="PiecewiseConstant"):
+        super().__init__()
+
+        if len(boundaries) != len(values) - 1:
+            raise ValueError(
+                "The length of boundaries should be 1 less than the length of "
+                f"values. Received: boundaries={boundaries} of length "
+                f"{len(boundaries)}, and values={values} "
+                f"of length {len(values)}."
+            )
+
+        self.boundaries = boundaries
+        self.values = values
+        self.name = name
+
+    def __call__(self, step):
+        with ops.name_scope(self.name):
+            boundaries = [ops.convert_to_tensor(x) for x in self.boundaries]
+            values = [ops.convert_to_tensor(x) for x in self.values]
+            step = ops.convert_to_tensor(step)
+
+            for i, b in enumerate(boundaries):
+                if b.dtype != step.dtype:
+                    # We cast the boundaries to have the same type as the step
+                    b = ops.cast(b, step.dtype)
+                    boundaries[i] = b
+
+            result_dtype = values[0].dtype
+            result_value = ops.array(0, dtype=result_dtype)
+
+            # For each range between boundaries, we check whether the step is
+            # within that range, cast the resulting boolean to a number,
+            # and multiply the result by the corresponding value for the range.
+            # Taking the sum of these yields a piecewise constant function.
+            step_less_than_first_boundary = ops.cast(
+                step <= boundaries[0], result_dtype
+            )
+            result_value += step_less_than_first_boundary * values[0]
+
+            step_greater_than_last_boundary = ops.cast(
+                step > boundaries[-1], result_dtype
+            )
+            result_value += step_greater_than_last_boundary * values[-1]
+
+            for low, high, value in zip(
+                boundaries[:-1], boundaries[1:], values[1:-1]
+            ):
+                # Need to bind v here; can do this with lambda v=v: ...
+                step_in_range = ops.cast(
+                    (step > low) & (step <= high), result_dtype
+                )
+                result_value += step_in_range * value
+
+            return result_value
+
+    def get_config(self):
+        return {
+            "boundaries": self.boundaries,
+            "values": self.values,
+            "name": self.name,
+        }
+
+
+@keras_export("keras.optimizers.schedules.PolynomialDecay")
+class PolynomialDecay(LearningRateSchedule):
+    """A `LearningRateSchedule` that uses a polynomial decay schedule.
+
+    It is commonly observed that a monotonically decreasing learning rate, whose
+    degree of change is carefully chosen, results in a better performing model.
+    This schedule applies a polynomial decay function to an optimizer step,
+    given a provided `initial_learning_rate`, to reach an `end_learning_rate`
+    in the given `decay_steps`.
+
+    It requires a `step` value to compute the decayed learning rate. You
+    can just pass a backend variable that you increment at each training
+    step.
+
+    The schedule is a 1-arg callable that produces a decayed learning rate
+    when passed the current optimizer step. This can be useful for changing the
+    learning rate value across different invocations of optimizer functions.
+    It is computed as:
+
+    ```python
+    def decayed_learning_rate(step):
+        step = min(step, decay_steps)
+        return ((initial_learning_rate - end_learning_rate) *
+                (1 - step / decay_steps) ^ (power)
+               ) + end_learning_rate
+    ```
+
+    If `cycle` is True then a multiple of `decay_steps` is used, the first one
+    that is bigger than `step`.
+
+    ```python
+    def decayed_learning_rate(step):
+        decay_steps = decay_steps * ceil(step / decay_steps)
+        return ((initial_learning_rate - end_learning_rate) *
+                (1 - step / decay_steps) ^ (power)
+               ) + end_learning_rate
+    ```
+
+    You can pass this schedule directly into a `keras.optimizers.Optimizer`
+    as the learning rate.
+    Example: Fit a model while decaying from 0.1 to 0.01 in 10000 steps using
+    sqrt (i.e. power=0.5):
+
+    ```python
+    ...
+    starter_learning_rate = 0.1
+    end_learning_rate = 0.01
+    decay_steps = 10000
+    learning_rate_fn = keras.optimizers.schedules.PolynomialDecay(
+        starter_learning_rate,
+        decay_steps,
+        end_learning_rate,
+        power=0.5)
+
+    model.compile(optimizer=keras.optimizers.SGD(
+                      learning_rate=learning_rate_fn),
+                  loss='sparse_categorical_crossentropy',
+                  metrics=['accuracy'])
+
+    model.fit(data, labels, epochs=5)
+    ```
+
+    The learning rate schedule is also serializable and deserializable using
+    `keras.optimizers.schedules.serialize` and
+    `keras.optimizers.schedules.deserialize`.
+
+    Args:
+        initial_learning_rate: A Python float. The initial learning rate.
+        decay_steps: A Python integer. Must be positive. See the decay
+            computation above.
+        end_learning_rate: A Python float. The minimal end learning rate.
+        power: A Python float. The power of the polynomial. Defaults to
+            `1.0`.
+        cycle: A boolean, whether it should cycle beyond decay_steps.
+        name: String.  Optional name of the operation. Defaults to
+            `"PolynomialDecay"`.
+
+    Returns:
+        A 1-arg callable learning rate schedule that takes the current optimizer
+        step and outputs the decayed learning rate, a scalar tensor of the
+        same type as `initial_learning_rate`.
+    """
+
+    def __init__(
+        self,
+        initial_learning_rate,
+        decay_steps,
+        end_learning_rate=0.0001,
+        power=1.0,
+        cycle=False,
+        name="PolynomialDecay",
+    ):
+        super().__init__()
+
+        self.initial_learning_rate = initial_learning_rate
+        self.decay_steps = decay_steps
+        self.end_learning_rate = end_learning_rate
+        self.power = power
+        self.cycle = cycle
+        self.name = name
+
+        if self.decay_steps <= 0:
+            raise ValueError(
+                "Argument `decay_steps` must be > 0. "
+                f"Received: decay_steps={self.decay_steps}"
+            )
+
+    def __call__(self, step):
+        with ops.name_scope(self.name):
+            initial_learning_rate = ops.convert_to_tensor(
+                self.initial_learning_rate
+            )
+            dtype = initial_learning_rate.dtype
+            end_learning_rate = ops.cast(self.end_learning_rate, dtype)
+            power = ops.cast(self.power, dtype)
+
+            global_step_recomp = ops.cast(step, dtype)
+            decay_steps_recomp = ops.cast(self.decay_steps, dtype)
+            if self.cycle:
+                # Find the first multiple of decay_steps that is bigger than
+                # global_step. If global_step is zero set the multiplier to 1
+                multiplier = ops.where(
+                    ops.equal(global_step_recomp, 0),
+                    1.0,
+                    ops.ceil(global_step_recomp / self.decay_steps),
+                )
+                decay_steps_recomp = ops.multiply(
+                    decay_steps_recomp, multiplier
+                )
+            else:
+                # Make sure that the global_step used is not bigger than
+                # decay_steps.
+                global_step_recomp = ops.minimum(
+                    global_step_recomp, decay_steps_recomp
+                )
+
+            p = ops.divide(global_step_recomp, decay_steps_recomp)
+            return ops.add(
+                ops.multiply(
+                    initial_learning_rate - end_learning_rate,
+                    ops.power(1 - p, power),
+                ),
+                end_learning_rate,
+            )
+
+    def get_config(self):
+        return {
+            "initial_learning_rate": self.initial_learning_rate,
+            "decay_steps": self.decay_steps,
+            "end_learning_rate": self.end_learning_rate,
+            "power": self.power,
+            "cycle": self.cycle,
+            "name": self.name,
+        }
+
+
+@keras_export("keras.optimizers.schedules.InverseTimeDecay")
+class InverseTimeDecay(LearningRateSchedule):
+    """A `LearningRateSchedule` that uses an inverse time decay schedule.
+
+    When training a model, it is often useful to lower the learning rate as
+    the training progresses. This schedule applies the inverse decay function
+    to an optimizer step, given a provided initial learning rate.
+    It requires a `step` value to compute the decayed learning rate. You can
+    just pass a backend variable that you increment at each training step.
+
+    The schedule is a 1-arg callable that produces a decayed learning
+    rate when passed the current optimizer step. This can be useful for changing
+    the learning rate value across different invocations of optimizer functions.
+    It is computed as:
+
+    ```python
+    def decayed_learning_rate(step):
+        return initial_learning_rate / (1 + decay_rate * step / decay_step)
+    ```
+
+    or, if `staircase` is `True`, as:
+
+    ```python
+    def decayed_learning_rate(step):
+        return initial_learning_rate /
+               (1 + decay_rate * floor(step / decay_step))
+    ```
+
+    You can pass this schedule directly into a `keras.optimizers.Optimizer`
+    as the learning rate.
+    Example: Fit a Keras model when decaying 1/t with a rate of 0.5:
+
+    ```python
+    ...
+    initial_learning_rate = 0.1
+    decay_steps = 1.0
+    decay_rate = 0.5
+    learning_rate_fn = keras.optimizers.schedules.InverseTimeDecay(
+        initial_learning_rate, decay_steps, decay_rate)
+
+    model.compile(optimizer=keras.optimizers.SGD(
+                      learning_rate=learning_rate_fn),
+                  loss='sparse_categorical_crossentropy',
+                  metrics=['accuracy'])
+
+    model.fit(data, labels, epochs=5)
+    ```
+
+    Args:
+        initial_learning_rate: A Python float. The initial learning rate.
+        decay_steps: How often to apply decay.
+        decay_rate: A Python number.  The decay rate.
+        staircase: Whether to apply decay in a discrete staircase, as o
+        pposed to continuous, fashion.
+        name: String.  Optional name of the operation.  Defaults to
+            `"InverseTimeDecay"`.
+
+    Returns:
+        A 1-arg callable learning rate schedule that takes the current optimizer
+        step and outputs the decayed learning rate, a scalar tensor of the
+        same type as `initial_learning_rate`.
+    """
+
+    def __init__(
+        self,
+        initial_learning_rate,
+        decay_steps,
+        decay_rate,
+        staircase=False,
+        name="InverseTimeDecay",
+    ):
+        super().__init__()
+
+        self.initial_learning_rate = initial_learning_rate
+        self.decay_steps = decay_steps
+        self.decay_rate = decay_rate
+        self.staircase = staircase
+        self.name = name
+
+        if self.decay_steps <= 0:
+            raise ValueError(
+                "Argument `decay_steps` must be > 0. "
+                f"Received: decay_steps={self.decay_steps}"
+            )
+
+    def __call__(self, step):
+        with ops.name_scope(self.name):
+            initial_learning_rate = ops.convert_to_tensor(
+                self.initial_learning_rate
+            )
+            dtype = initial_learning_rate.dtype
+            decay_steps = ops.cast(self.decay_steps, dtype)
+            decay_rate = ops.cast(self.decay_rate, dtype)
+
+            global_step_recomp = ops.cast(step, dtype)
+            p = global_step_recomp / decay_steps
+            if self.staircase:
+                p = ops.floor(p)
+            const = ops.cast(ops.array(1), dtype)
+            denom = ops.add(const, ops.multiply(decay_rate, p))
+            return ops.divide(initial_learning_rate, denom)
+
+    def get_config(self):
+        return {
+            "initial_learning_rate": self.initial_learning_rate,
+            "decay_steps": self.decay_steps,
+            "decay_rate": self.decay_rate,
+            "staircase": self.staircase,
+            "name": self.name,
+        }
+
+
+@keras_export("keras.optimizers.schedules.CosineDecay")
+class CosineDecay(LearningRateSchedule):
+    """A `LearningRateSchedule` that uses a cosine decay with optional warmup.
+
+    See [Loshchilov & Hutter, ICLR2016](https://arxiv.org/abs/1608.03983),
+    SGDR: Stochastic Gradient Descent with Warm Restarts.
+
+    For the idea of a linear warmup of our learning rate,
+    see [Goyal et al.](https://arxiv.org/pdf/1706.02677.pdf).
+
+    When we begin training a model, we often want an initial increase in our
+    learning rate followed by a decay. If `warmup_target` is an int, this
+    schedule applies a linear increase per optimizer step to our learning rate
+    from `initial_learning_rate` to `warmup_target` for a duration of
+    `warmup_steps`. Afterwards, it applies a cosine decay function taking our
+    learning rate from `warmup_target` to `alpha` for a duration of
+    `decay_steps`. If `warmup_target` is None we skip warmup and our decay
+    will take our learning rate from `initial_learning_rate` to `alpha`.
+    It requires a `step` value to  compute the learning rate. You can
+    just pass a backend variable that you increment at each training step.
+
+    The schedule is a 1-arg callable that produces a warmup followed by a
+    decayed learning rate when passed the current optimizer step. This can be
+    useful for changing the learning rate value across different invocations of
+    optimizer functions.
+
+    Our warmup is computed as:
+
+    ```python
+    def warmup_learning_rate(step):
+        completed_fraction = step / warmup_steps
+        total_delta = target_warmup - initial_learning_rate
+        return completed_fraction * total_delta
+    ```
+
+    And our decay is computed as:
+
+    ```python
+    if warmup_target is None:
+        initial_decay_lr = initial_learning_rate
+    else:
+        initial_decay_lr = warmup_target
+
+    def decayed_learning_rate(step):
+        step = min(step, decay_steps)
+        cosine_decay = 0.5 * (1 + cos(pi * step / decay_steps))
+        decayed = (1 - alpha) * cosine_decay + alpha
+        return initial_decay_lr * decayed
+    ```
+
+    Example usage without warmup:
+
+    ```python
+    decay_steps = 1000
+    initial_learning_rate = 0.1
+    lr_decayed_fn = keras.optimizers.schedules.CosineDecay(
+        initial_learning_rate, decay_steps)
+    ```
+
+    Example usage with warmup:
+
+    ```python
+    decay_steps = 1000
+    initial_learning_rate = 0
+    warmup_steps = 1000
+    target_learning_rate = 0.1
+    lr_warmup_decayed_fn = keras.optimizers.schedules.CosineDecay(
+        initial_learning_rate, decay_steps, warmup_target=target_learning_rate,
+        warmup_steps=warmup_steps
+    )
+    ```
+
+    You can pass this schedule directly into a `keras.optimizers.Optimizer`
+    as the learning rate. The learning rate schedule is also serializable and
+    deserializable using `keras.optimizers.schedules.serialize` and
+    `keras.optimizers.schedules.deserialize`.
+
+    Args:
+        initial_learning_rate: A Python float. The initial learning rate.
+        decay_steps: A Python int. Number of steps to decay over.
+        alpha: A Python float. Minimum learning rate value for decay as a
+            fraction of `initial_learning_rate`.
+        name: String. Optional name of the operation.  Defaults to
+            `"CosineDecay"`.
+        warmup_target: A Python float. The target learning rate for our
+            warmup phase. Will cast to the `initial_learning_rate` datatype.
+            Setting to `None` will skip warmup and begins decay phase from
+            `initial_learning_rate`. Otherwise scheduler will warmup from
+            `initial_learning_rate` to `warmup_target`.
+        warmup_steps: A Python int. Number of steps to warmup over.
+
+    Returns:
+        A 1-arg callable learning rate schedule that takes the current optimizer
+        step and outputs the decayed learning rate, a scalar tensor of the
+        same type as `initial_learning_rate`.
+    """
+
+    def __init__(
+        self,
+        initial_learning_rate,
+        decay_steps,
+        alpha=0.0,
+        name="CosineDecay",
+        warmup_target=None,
+        warmup_steps=0,
+    ):
+        super().__init__()
+
+        self.initial_learning_rate = initial_learning_rate
+        self.decay_steps = decay_steps
+        self.alpha = alpha
+        self.name = name
+        self.warmup_steps = warmup_steps
+        self.warmup_target = warmup_target
+
+        if self.decay_steps <= 0:
+            raise ValueError(
+                "Argument `decay_steps` must be > 0. "
+                f"Received: decay_steps={self.decay_steps}"
+            )
+
+    def _decay_function(self, step, decay_steps, decay_from_lr, dtype):
+        with ops.name_scope(self.name):
+            completed_fraction = step / decay_steps
+            pi = ops.array(math.pi, dtype=dtype)
+            cosine_decayed = 0.5 * (1.0 + ops.cos(pi * completed_fraction))
+            decayed = (1 - self.alpha) * cosine_decayed + self.alpha
+            return ops.multiply(decay_from_lr, decayed)
+
+    def _warmup_function(
+        self, step, warmup_steps, warmup_target, initial_learning_rate
+    ):
+        with ops.name_scope(self.name):
+            completed_fraction = step / warmup_steps
+            total_step_delta = warmup_target - initial_learning_rate
+            return total_step_delta * completed_fraction + initial_learning_rate
+
+    def __call__(self, step):
+        with ops.name_scope(self.name):
+            initial_learning_rate = ops.convert_to_tensor(
+                self.initial_learning_rate
+            )
+            dtype = initial_learning_rate.dtype
+            decay_steps = ops.cast(self.decay_steps, dtype)
+            global_step_recomp = ops.cast(step, dtype)
+
+            if self.warmup_target is None:
+                global_step_recomp = ops.minimum(
+                    global_step_recomp, decay_steps
+                )
+                return self._decay_function(
+                    global_step_recomp,
+                    decay_steps,
+                    initial_learning_rate,
+                    dtype,
+                )
+
+            warmup_target = ops.cast(self.warmup_target, dtype)
+            warmup_steps = ops.cast(self.warmup_steps, dtype)
+
+            global_step_recomp = ops.minimum(
+                global_step_recomp, decay_steps + warmup_steps
+            )
+
+            return ops.cond(
+                global_step_recomp < warmup_steps,
+                lambda: self._warmup_function(
+                    global_step_recomp,
+                    warmup_steps,
+                    warmup_target,
+                    initial_learning_rate,
+                ),
+                lambda: self._decay_function(
+                    global_step_recomp - warmup_steps,
+                    decay_steps,
+                    warmup_target,
+                    dtype,
+                ),
+            )
+
+    def get_config(self):
+        return {
+            "initial_learning_rate": self.initial_learning_rate,
+            "decay_steps": self.decay_steps,
+            "alpha": self.alpha,
+            "name": self.name,
+            "warmup_target": self.warmup_target,
+            "warmup_steps": self.warmup_steps,
+        }
+
+
+@keras_export("keras.optimizers.schedules.CosineDecayRestarts")
+class CosineDecayRestarts(LearningRateSchedule):
+    """A `LearningRateSchedule` that uses a cosine decay schedule with restarts.
+
+    See [Loshchilov & Hutter, ICLR2016](https://arxiv.org/abs/1608.03983),
+    SGDR: Stochastic Gradient Descent with Warm Restarts.
+
+    When training a model, it is often useful to lower the learning rate as
+    the training progresses. This schedule applies a cosine decay function with
+    restarts to an optimizer step, given a provided initial learning rate.
+    It requires a `step` value to compute the decayed learning rate. You can
+    just pass a backend variable that you increment at each training step.
+
+    The schedule is a 1-arg callable that produces a decayed learning
+    rate when passed the current optimizer step. This can be useful for changing
+    the learning rate value across different invocations of optimizer functions.
+
+    The learning rate multiplier first decays
+    from 1 to `alpha` for `first_decay_steps` steps. Then, a warm
+    restart is performed. Each new warm restart runs for `t_mul` times more
+    steps and with `m_mul` times initial learning rate as the new learning rate.
+
+    Example:
+    ```python
+    first_decay_steps = 1000
+    lr_decayed_fn = (
+        keras.optimizers.schedules.CosineDecayRestarts(
+            initial_learning_rate,
+            first_decay_steps))
+    ```
+
+    You can pass this schedule directly into a `keras.optimizers.Optimizer`
+    as the learning rate. The learning rate schedule is also serializable and
+    deserializable using `keras.optimizers.schedules.serialize` and
+    `keras.optimizers.schedules.deserialize`.
+
+    Args:
+        initial_learning_rate: A Python float. The initial learning rate.
+        first_decay_steps: A Python integer. Number of steps to decay over.
+        t_mul: A Python float. Used to derive the number of iterations in
+            the i-th period.
+        m_mul: A Python float. Used to derive the initial learning rate of
+            the i-th period.
+        alpha: A Python float. Minimum learning rate value as a fraction of
+            the `initial_learning_rate`.
+        name: String. Optional name of the operation. Defaults to
+            `"SGDRDecay"`.
+
+    Returns:
+        A 1-arg callable learning rate schedule that takes the current optimizer
+        step and outputs the decayed learning rate, a scalar tensor of the
+        same type as `initial_learning_rate`.
+    """
+
+    def __init__(
+        self,
+        initial_learning_rate,
+        first_decay_steps,
+        t_mul=2.0,
+        m_mul=1.0,
+        alpha=0.0,
+        name="SGDRDecay",
+    ):
+        super().__init__()
+
+        self.initial_learning_rate = initial_learning_rate
+        self.first_decay_steps = first_decay_steps
+        self._t_mul = t_mul
+        self._m_mul = m_mul
+        self.alpha = alpha
+        self.name = name
+
+        if self.first_decay_steps <= 0:
+            raise ValueError(
+                "Argument `first_decay_steps` must be > 0. "
+                f"Received: first_decay_steps={self.first_decay_steps}"
+            )
+
+    def __call__(self, step):
+        with ops.name_scope(self.name):
+            initial_learning_rate = ops.convert_to_tensor(
+                self.initial_learning_rate
+            )
+            dtype = initial_learning_rate.dtype
+            first_decay_steps = ops.cast(self.first_decay_steps, dtype)
+            alpha = ops.cast(self.alpha, dtype)
+            t_mul = ops.cast(self._t_mul, dtype)
+            m_mul = ops.cast(self._m_mul, dtype)
+
+            global_step_recomp = ops.cast(step, dtype)
+            completed_fraction = global_step_recomp / first_decay_steps
+
+            def compute_step(completed_fraction, geometric=False):
+                """Helper for `cond` operation."""
+                if geometric:
+                    # ops.log is sensitive to the precision of dtype, so we need
+                    # the additional casting
+                    i_restart = ops.floor(
+                        ops.log(
+                            ops.cast(
+                                1.0 - completed_fraction * (1.0 - t_mul), dtype
+                            )
+                        )
+                        / ops.log(t_mul)
+                    )
+
+                    sum_r = (1.0 - t_mul**i_restart) / (1.0 - t_mul)
+                    completed_fraction = (
+                        completed_fraction - sum_r
+                    ) / t_mul**i_restart
+
+                else:
+                    i_restart = ops.floor(completed_fraction)
+                    completed_fraction -= i_restart
+
+                return i_restart, completed_fraction
+
+            i_restart, completed_fraction = ops.cond(
+                ops.equal(t_mul, 1.0),
+                lambda: compute_step(completed_fraction, geometric=False),
+                lambda: compute_step(completed_fraction, geometric=True),
+            )
+
+            m_fac = m_mul**i_restart
+            cosine_decayed = (
+                0.5
+                * m_fac
+                * (
+                    1.0
+                    + ops.cos(
+                        ops.array(math.pi, dtype=dtype) * completed_fraction
+                    )
+                )
+            )
+            decayed = (1 - alpha) * cosine_decayed + alpha
+
+            return ops.multiply(initial_learning_rate, decayed)
+
+    def get_config(self):
+        return {
+            "initial_learning_rate": self.initial_learning_rate,
+            "first_decay_steps": self.first_decay_steps,
+            "t_mul": self._t_mul,
+            "m_mul": self._m_mul,
+            "alpha": self.alpha,
+            "name": self.name,
+        }
+
+
+@keras_export("keras.optimizers.schedules.serialize")
+def serialize(learning_rate_schedule):
+    """Serializes a `LearningRateSchedule` into a JSON-compatible dict.
+
+    Args:
+        learning_rate_schedule: The `LearningRateSchedule` object to serialize.
+
+    Returns:
+        A JSON-serializable dict representing the object's config.
+
+    Example:
+
+    >>> lr_schedule = keras.optimizers.schedules.ExponentialDecay(
+    ...     0.1, decay_steps=100000, decay_rate=0.96, staircase=True)
+    >>> keras.optimizers.schedules.serialize(lr_schedule)
+    {'module': 'keras.optimizers.schedules',
+    'class_name': 'ExponentialDecay', 'config': {...},
+    'registered_name': None}
+    """
+    return serialization_lib.serialize_keras_object(learning_rate_schedule)
+
+
+@keras_export("keras.optimizers.schedules.deserialize")
+def deserialize(config, custom_objects=None):
+    """Instantiates a `LearningRateSchedule` object from a serialized form.
+
+    Args:
+        config: The serialized form of the `LearningRateSchedule`. Dictionary of
+            the form {'class_name': str, 'config': dict}.
+        custom_objects: A dictionary mapping class names (or function names) of
+            custom (non-Keras) objects to class/functions.
+
+    Returns:
+        A `LearningRateSchedule` object.
+
+    Example:
+
+    ```python
+    # Configuration for PolynomialDecay
+    config = {
+        'class_name': 'PolynomialDecay',
+        'config': {'cycle': False,
+            'decay_steps': 10000,
+            'end_learning_rate': 0.01,
+            'initial_learning_rate': 0.1,
+            'name': None,
+            'power': 0.5
+        }
+    }
+    lr_schedule = keras.optimizers.schedules.deserialize(config)
+    ```
+    """
+    return serialization_lib.deserialize_keras_object(
+        config,
+        module_objects=globals(),
+        custom_objects=custom_objects,
+        printable_module_name="decay",
+    )

SwarmUI/dlbackend/ComfyUI/venv/lib/python3.10/site-packages/keras/src/optimizers/sgd.py

@@ -0,0 +1,143 @@
+from keras.src import ops
+from keras.src.api_export import keras_export
+from keras.src.optimizers import optimizer
+
+
+@keras_export("keras.optimizers.SGD")
+class SGD(optimizer.Optimizer):
+    """Gradient descent (with momentum) optimizer.
+
+    Update rule for parameter `w` with gradient `g` when `momentum` is 0:
+
+    ```python
+    w = w - learning_rate * g
+    ```
+
+    Update rule when `momentum` is larger than 0:
+
+    ```python
+    velocity = momentum * velocity - learning_rate * g
+    w = w + velocity
+    ```
+
+    When `nesterov=True`, this rule becomes:
+
+    ```python
+    velocity = momentum * velocity - learning_rate * g
+    w = w + momentum * velocity - learning_rate * g
+    ```
+
+    Args:
+        learning_rate: A float, a
+            `keras.optimizers.schedules.LearningRateSchedule` instance, or
+            a callable that takes no arguments and returns the actual value to
+            use. The learning rate. Defaults to `0.01`.
+        momentum: float hyperparameter >= 0 that accelerates gradient descent in
+            the relevant direction and dampens oscillations. 0 is vanilla
+            gradient descent. Defaults to `0.0`.
+        nesterov: boolean. Whether to apply Nesterov momentum.
+            Defaults to `False`.
+        {{base_optimizer_keyword_args}}
+    """
+
+    def __init__(
+        self,
+        learning_rate=0.01,
+        momentum=0.0,
+        nesterov=False,
+        weight_decay=None,
+        clipnorm=None,
+        clipvalue=None,
+        global_clipnorm=None,
+        use_ema=False,
+        ema_momentum=0.99,
+        ema_overwrite_frequency=None,
+        loss_scale_factor=None,
+        gradient_accumulation_steps=None,
+        name="SGD",
+        **kwargs,
+    ):
+        super().__init__(
+            learning_rate=learning_rate,
+            name=name,
+            weight_decay=weight_decay,
+            clipnorm=clipnorm,
+            clipvalue=clipvalue,
+            global_clipnorm=global_clipnorm,
+            use_ema=use_ema,
+            ema_momentum=ema_momentum,
+            ema_overwrite_frequency=ema_overwrite_frequency,
+            loss_scale_factor=loss_scale_factor,
+            gradient_accumulation_steps=gradient_accumulation_steps,
+            **kwargs,
+        )
+        if not isinstance(momentum, float) or momentum < 0 or momentum > 1:
+            raise ValueError("`momentum` must be a float between [0, 1].")
+        self.momentum = momentum
+        self.nesterov = nesterov
+
+    def build(self, variables):
+        """Initialize optimizer variables.
+
+        SGD optimizer has one variable `momentums`, only set if `self.momentum`
+        is not 0.
+
+        Args:
+          var_list: list of model variables to build SGD variables on.
+        """
+        if self.built:
+            return
+        super().build(variables)
+        self.momentums = []
+        if self.momentum != 0:
+            for variable in variables:
+                self.momentums.append(
+                    self.add_variable_from_reference(
+                        reference_variable=variable, name="momentum"
+                    )
+                )
+
+    def update_step(self, gradient, variable, learning_rate):
+        """Update step given gradient and the associated model variable."""
+        learning_rate = ops.cast(learning_rate, variable.dtype)
+        gradient = ops.cast(gradient, variable.dtype)
+        m = None
+        if self.momentum != 0:
+            m = self.momentums[self._get_variable_index(variable)]
+
+        if m is not None:
+            momentum = ops.cast(self.momentum, variable.dtype)
+            self.assign(
+                m,
+                ops.subtract(
+                    ops.multiply(m, momentum),
+                    ops.multiply(gradient, learning_rate),
+                ),
+            )
+            if self.nesterov:
+                self.assign_add(
+                    variable,
+                    ops.subtract(
+                        ops.multiply(m, momentum),
+                        ops.multiply(gradient, learning_rate),
+                    ),
+                )
+            else:
+                self.assign_add(variable, m)
+        else:
+            self.assign_sub(variable, ops.multiply(gradient, learning_rate))
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "momentum": self.momentum,
+                "nesterov": self.nesterov,
+            }
+        )
+        return config
+
+
+SGD.__doc__ = SGD.__doc__.replace(
+    "{{base_optimizer_keyword_args}}", optimizer.base_optimizer_keyword_args
+)