backend.py

# Copyright 2015 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================


"""Keras backend API."""

import collections
import itertools
import json
import os
import random
import sys
import threading
import warnings
import weakref

import numpy as np
import tensorflow.compat.v2 as tf

from tf_keras import backend_config
from tf_keras.distribute import distribute_coordinator_utils as dc
from tf_keras.dtensor import dtensor_api as dtensor
from tf_keras.engine import keras_tensor
from tf_keras.utils import control_flow_util
from tf_keras.utils import object_identity
from tf_keras.utils import tf_contextlib
from tf_keras.utils import tf_inspect
from tf_keras.utils import tf_utils

# isort: off
from tensorflow.core.protobuf import config_pb2
from tensorflow.python.eager import context
from tensorflow.python.eager.context import get_config
from tensorflow.python.platform import tf_logging as logging
from tensorflow.python.util.tf_export import keras_export
from tensorflow.tools.docs import doc_controls

py_all = all
py_sum = sum
py_any = any

# INTERNAL UTILS

# The internal graph maintained by TF-Keras and used by the symbolic TF-Keras
# APIs while executing eagerly (such as the functional API for model-building).
# This is thread-local to allow building separate models in different threads
# concurrently, but comes at the cost of not being able to build one model
# across threads.
_GRAPH = threading.local()

# A graph which is used for constructing functions in eager mode.
_CURRENT_SCRATCH_GRAPH = threading.local()


# This is a thread local object that will hold the default internal TF session
# used by TF-Keras. It can be set manually via `set_session(sess)`.
class SessionLocal(threading.local):
    def __init__(self):
        super().__init__()
        self.session = None


_SESSION = SessionLocal()


# A global dictionary mapping graph objects to an index of counters used
# for various layer/optimizer names in each graph.
# Allows to give unique autogenerated names to layers, in a graph-specific way.
PER_GRAPH_OBJECT_NAME_UIDS = weakref.WeakKeyDictionary()


# A global set tracking what object names have been seen so far.
# Optionally used as an avoid-list when generating names
OBSERVED_NAMES = set()


# _DUMMY_EAGER_GRAPH.key is used as a key in _GRAPH_LEARNING_PHASES.
# We keep a separate reference to it to make sure it does not get removed from
# _GRAPH_LEARNING_PHASES.
# _DummyEagerGraph inherits from threading.local to make its `key` attribute
# thread local. This is needed to make set_learning_phase affect only the
# current thread during eager execution (see b/123096885 for more details).
class _DummyEagerGraph(threading.local):
    """_DummyEagerGraph provides a thread local `key` attribute.

    We can't use threading.local directly, i.e. without subclassing, because
    gevent monkey patches threading.local and its version does not support
    weak references.
    """

    class _WeakReferencableClass:
        """This dummy class is needed for two reasons.

        - We need something that supports weak references. Basic types like
        string and ints don't.
        - We need something whose hash and equality are based on object identity
        to make sure they are treated as different keys to
        _GRAPH_LEARNING_PHASES.

        An empty Python class satisfies both of these requirements.
        """

        pass

    def __init__(self):
        # Constructors for classes subclassing threading.local run once
        # per thread accessing something in the class. Thus, each thread will
        # get a different key.
        super().__init__()
        self.key = _DummyEagerGraph._WeakReferencableClass()
        self.learning_phase_is_set = False


_DUMMY_EAGER_GRAPH = _DummyEagerGraph()

# This boolean flag can be set to True to leave variable initialization
# up to the user.
# Change its value via `manual_variable_initialization(value)`.
_MANUAL_VAR_INIT = False

# This list holds the available devices.
# It is populated when `_get_available_gpus()` is called for the first time.
# We assume our devices don't change henceforth.
_LOCAL_DEVICES = None

# The below functions are kept accessible from backend for compatibility.
epsilon = backend_config.epsilon
floatx = backend_config.floatx
image_data_format = backend_config.image_data_format
set_epsilon = backend_config.set_epsilon
set_floatx = backend_config.set_floatx
set_image_data_format = backend_config.set_image_data_format


@keras_export("keras.backend.backend")
@doc_controls.do_not_generate_docs
def backend():
    """Publicly accessible method for determining the current backend.

    Only exists for API compatibility with multi-backend TF-Keras.

    Returns:
        The string "tensorflow".
    """
    return "tensorflow"


@keras_export("keras.backend.cast_to_floatx")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def cast_to_floatx(x):
    """Cast a Numpy array to the default TF-Keras float type.

    Args:
        x: Numpy array or TensorFlow tensor.

    Returns:
        The same array (Numpy array if `x` was a Numpy array, or TensorFlow
        tensor if `x` was a tensor), cast to its new type.

    Example:

    >>> tf.keras.backend.floatx()
    'float32'
    >>> arr = np.array([1.0, 2.0], dtype='float64')
    >>> arr.dtype
    dtype('float64')
    >>> new_arr = cast_to_floatx(arr)
    >>> new_arr
    array([1.,  2.], dtype=float32)
    >>> new_arr.dtype
    dtype('float32')

    """
    if isinstance(x, (tf.Tensor, tf.Variable, tf.SparseTensor)):
        return tf.cast(x, dtype=floatx())
    return np.asarray(x, dtype=floatx())


@keras_export("keras.backend.get_uid")
def get_uid(prefix=""):
    """Associates a string prefix with an integer counter in a TensorFlow graph.

    Args:
      prefix: String prefix to index.

    Returns:
      Unique integer ID.

    Example:

    >>> get_uid('dense')
    1
    >>> get_uid('dense')
    2

    """
    graph = get_graph()
    if graph not in PER_GRAPH_OBJECT_NAME_UIDS:
        PER_GRAPH_OBJECT_NAME_UIDS[graph] = collections.defaultdict(int)
    layer_name_uids = PER_GRAPH_OBJECT_NAME_UIDS[graph]
    layer_name_uids[prefix] += 1
    return layer_name_uids[prefix]


@keras_export("keras.backend.reset_uids")
def reset_uids():
    """Resets graph identifiers."""

    PER_GRAPH_OBJECT_NAME_UIDS.clear()
    OBSERVED_NAMES.clear()


@keras_export("keras.backend.clear_session")
def clear_session():
    """Resets all state generated by TF-Keras.

    TF-Keras manages a global state, which it uses to implement the Functional
    model-building API and to uniquify autogenerated layer names.

    If you are creating many models in a loop, this global state will consume
    an increasing amount of memory over time, and you may want to clear it.
    Calling `clear_session()` releases the global state: this helps avoid
    clutter from old models and layers, especially when memory is limited.

    Example 1: calling `clear_session()` when creating models in a loop

    ```python
    for _ in range(100):
      # Without `clear_session()`, each iteration of this loop will
      # slightly increase the size of the global state managed by Keras
      model = tf.keras.Sequential([
          tf.keras.layers.Dense(10) for _ in range(10)])

    for _ in range(100):
      # With `clear_session()` called at the beginning,
      # TF-Keras starts with a blank state at each iteration
      # and memory consumption is constant over time.
      tf.keras.backend.clear_session()
      model = tf.keras.Sequential([
          tf.keras.layers.Dense(10) for _ in range(10)])
    ```

    Example 2: resetting the layer name generation counter

    >>> import tensorflow as tf
    >>> layers = [tf.keras.layers.Dense(10) for _ in range(10)]
    >>> new_layer = tf.keras.layers.Dense(10)
    >>> print(new_layer.name)
    dense_10
    >>> tf.keras.backend.set_learning_phase(1)
    >>> print(tf.keras.backend.learning_phase())
    1
    >>> tf.keras.backend.clear_session()
    >>> new_layer = tf.keras.layers.Dense(10)
    >>> print(new_layer.name)
    dense
    """
    global _SESSION
    global _GRAPH_LEARNING_PHASES
    global _GRAPH_VARIABLES
    global _GRAPH_TF_OPTIMIZERS
    global _GRAPH
    _GRAPH.graph = None
    tf.compat.v1.reset_default_graph()
    reset_uids()
    if _SESSION.session is not None:
        _SESSION.session.close()
        _SESSION.session = None
    graph = get_graph()
    with graph.as_default():
        _DUMMY_EAGER_GRAPH.learning_phase_is_set = False

        _GRAPH_LEARNING_PHASES = {}
        # Create the learning phase placeholder in graph using the default
        # factory
        phase = _default_learning_phase()
        _internal_set_learning_phase(graph, phase)

        _GRAPH_VARIABLES.pop(graph, None)
        _GRAPH_TF_OPTIMIZERS.pop(graph, None)
    if tf.executing_eagerly():
        # Clear pending nodes in eager executors, kernel caches and
        # step_containers.
        context.context().clear_kernel_cache()


# Inject the clear_session function to keras_deps to remove the dependency
# from TFLite to TF-Keras.
tf.__internal__.register_clear_session_function(clear_session)


@keras_export("keras.backend.manual_variable_initialization")
@doc_controls.do_not_generate_docs
def manual_variable_initialization(value):
    """Sets the manual variable initialization flag.

    This boolean flag determines whether
    variables should be initialized
    as they are instantiated (default), or if
    the user should handle the initialization
    (e.g. via `tf.compat.v1.initialize_all_variables()`).

    Args:
        value: Python boolean.
    """
    global _MANUAL_VAR_INIT
    _MANUAL_VAR_INIT = value


@keras_export("keras.backend.learning_phase")
@doc_controls.do_not_generate_docs
def learning_phase():
    """Returns the learning phase flag.

    The learning phase flag is a bool tensor (0 = test, 1 = train)
    to be passed as input to any TF-Keras function
    that uses a different behavior at train time and test time.

    Returns:
        Learning phase (scalar integer tensor or Python integer).
    """
    graph = tf.compat.v1.get_default_graph()
    if graph is getattr(_GRAPH, "graph", None):
        # Don't enter an init_scope for the learning phase if eager execution
        # is enabled but we're inside the TF-Keras workspace graph.
        learning_phase = symbolic_learning_phase()
    else:
        with tf.init_scope():
            # We always check & set the learning phase inside the init_scope,
            # otherwise the wrong default_graph will be used to look up the
            # learning phase inside of functions & defuns.
            #
            # This is because functions & defuns (both in graph & in eager mode)
            # will always execute non-eagerly using a function-specific default
            # subgraph.
            if context.executing_eagerly():
                if _DUMMY_EAGER_GRAPH.key not in _GRAPH_LEARNING_PHASES:
                    return _default_learning_phase()
                else:
                    return _internal_get_learning_phase(_DUMMY_EAGER_GRAPH.key)
            else:
                learning_phase = symbolic_learning_phase()
    _mark_func_graph_as_unsaveable(graph, learning_phase)
    return learning_phase


def global_learning_phase_is_set():
    return _DUMMY_EAGER_GRAPH.learning_phase_is_set


def _mark_func_graph_as_unsaveable(graph, learning_phase):
    """Mark graph as unsaveable due to use of symbolic keras learning phase.

    Functions that capture the symbolic learning phase cannot be exported to
    SavedModel. Mark the funcgraph as unsaveable, so that an error will be
    raised if it is exported.

    Args:
      graph: Graph or FuncGraph object.
      learning_phase: Learning phase placeholder or int defined in the graph.
    """
    if graph.building_function and is_placeholder(learning_phase):
        graph.mark_as_unsaveable(
            "The keras learning phase placeholder was used inside a function. "
            "Exporting placeholders is not supported when saving out a "
            "SavedModel. Please call `tf.keras.backend.set_learning_phase(0)` "
            "in the function to set the learning phase to a constant value."
        )


def symbolic_learning_phase():
    graph = get_graph()
    with graph.as_default():
        if graph not in _GRAPH_LEARNING_PHASES:
            phase = _default_learning_phase()
            _internal_set_learning_phase(graph, phase)

        return _internal_get_learning_phase(graph)


def _internal_set_learning_phase(graph, value):
    global _GRAPH_LEARNING_PHASES

    if isinstance(value, tf.Tensor):
        # The 'value' here is a tf.Tensor with attribute 'graph'.
        # There is a circular reference between key 'graph' and attribute
        # 'graph'.  So we need use a weakref.ref to refer to the 'value' tensor
        # here.  Otherwise, it would lead to memory leak.
        value_ref = weakref.ref(value)
        _GRAPH_LEARNING_PHASES[graph] = value_ref
    else:
        _GRAPH_LEARNING_PHASES[graph] = value


def _internal_get_learning_phase(graph):
    phase = _GRAPH_LEARNING_PHASES.get(graph, None)
    if isinstance(phase, weakref.ref):
        return phase()
    else:
        return phase


def _default_learning_phase():
    if context.executing_eagerly():
        return 0
    else:
        with name_scope(""):
            return tf.compat.v1.placeholder_with_default(
                False, shape=(), name="keras_learning_phase"
            )


@keras_export("keras.backend.set_learning_phase")
@doc_controls.do_not_generate_docs
def set_learning_phase(value):
    """Sets the learning phase to a fixed value.

    The backend learning phase affects any code that calls
    `backend.learning_phase()`
    In particular, all TF-Keras built-in layers use the learning phase as the
    default for the `training` arg to `Layer.__call__`.

    User-written layers and models can achieve the same behavior with code that
    looks like:

    ```python
      def call(self, inputs, training=None):
        if training is None:
          training = backend.learning_phase()
    ```

    Args:
        value: Learning phase value, either 0 or 1 (integers).
               0 = test, 1 = train

    Raises:
        ValueError: if `value` is neither `0` nor `1`.
    """
    warnings.warn(
        "`tf.keras.backend.set_learning_phase` is deprecated and "
        "will be removed after 2020-10-11. To update it, simply "
        "pass a True/False value to the `training` argument of the "
        "`__call__` method of your layer or model."
    )
    deprecated_internal_set_learning_phase(value)


def deprecated_internal_set_learning_phase(value):
    """A deprecated internal implementation of set_learning_phase.

    This method is an internal-only version of `set_learning_phase` that
    does not raise a deprecation error. It is required because
    saved_model needs to keep working with user code that uses the deprecated
    learning phase methods until those APIs are fully removed from the public
    API.

    Specifically SavedModel saving needs to make sure the learning phase is 0
    during tracing even if users overwrote it to a different value.

    But, we don't want to raise deprecation warnings for users when savedmodel
    sets learning phase just for compatibility with code that relied on
    explicitly setting the learning phase for other values.

    Args:
        value: Learning phase value, either 0 or 1 (integers).
            0 = test, 1 = train

    Raises:
        ValueError: if `value` is neither `0` nor `1`.
    """
    if value not in {0, 1}:
        raise ValueError("Expected learning phase to be 0 or 1.")
    with tf.init_scope():
        if tf.executing_eagerly():
            # In an eager context, the learning phase values applies to both the
            # eager context and the internal TF-Keras graph.
            _DUMMY_EAGER_GRAPH.learning_phase_is_set = True
            _internal_set_learning_phase(_DUMMY_EAGER_GRAPH.key, value)

        _internal_set_learning_phase(get_graph(), value)


@keras_export("keras.backend.learning_phase_scope")
@tf_contextlib.contextmanager
@doc_controls.do_not_generate_docs
def learning_phase_scope(value):
    """Provides a scope within which the learning phase is equal to `value`.

    The learning phase gets restored to its original value upon exiting the
    scope.

    Args:
       value: Learning phase value, either 0 or 1 (integers).
              0 = test, 1 = train

    Yields:
      None.

    Raises:
       ValueError: if `value` is neither `0` nor `1`.
    """
    warnings.warn(
        "`tf.keras.backend.learning_phase_scope` is deprecated and "
        "will be removed after 2020-10-11. To update it, simply "
        "pass a True/False value to the `training` argument of the "
        "`__call__` method of your layer or model.",
        stacklevel=2,
    )
    with deprecated_internal_learning_phase_scope(value):
        try:
            yield
        finally:
            pass


@tf_contextlib.contextmanager
def deprecated_internal_learning_phase_scope(value):
    """An internal-only version of `learning_phase_scope`.

    Unlike the public method, this method does not raise a deprecation warning.
    This is needed because saved model saving needs to set learning phase
    to maintain compatibility
    with code that sets/gets the learning phase, but saved model
    saving itself shouldn't raise a deprecation warning.

    We can get rid of this method and its usages when the public API is
    removed.

    Args:
        value: Learning phase value, either 0 or 1 (integers).
            0 = test, 1 = train

    Yields:
        None.

    Raises:
        ValueError: if `value` is neither `0` nor `1`.
    """
    global _GRAPH_LEARNING_PHASES
    if value not in {0, 1}:
        raise ValueError("Expected learning phase to be 0 or 1.")

    with tf.init_scope():
        if tf.executing_eagerly():
            previous_eager_value = _internal_get_learning_phase(
                _DUMMY_EAGER_GRAPH.key
            )
        previous_graph_value = _internal_get_learning_phase(get_graph())

    learning_phase_previously_set = _DUMMY_EAGER_GRAPH.learning_phase_is_set
    try:
        deprecated_internal_set_learning_phase(value)
        yield
    finally:
        # Restore learning phase to initial value.
        if not learning_phase_previously_set:
            _DUMMY_EAGER_GRAPH.learning_phase_is_set = False
        with tf.init_scope():
            if tf.executing_eagerly():
                if previous_eager_value is not None:
                    _internal_set_learning_phase(
                        _DUMMY_EAGER_GRAPH.key, previous_eager_value
                    )
                elif _DUMMY_EAGER_GRAPH.key in _GRAPH_LEARNING_PHASES:
                    del _GRAPH_LEARNING_PHASES[_DUMMY_EAGER_GRAPH.key]

            graph = get_graph()
            if previous_graph_value is not None:
                _internal_set_learning_phase(graph, previous_graph_value)
            elif graph in _GRAPH_LEARNING_PHASES:
                del _GRAPH_LEARNING_PHASES[graph]


@tf_contextlib.contextmanager
def eager_learning_phase_scope(value):
    """Internal scope that sets the learning phase in eager / tf.function only.

    Args:
        value: Learning phase value, either 0 or 1 (integers).
               0 = test, 1 = train

    Yields:
      None.

    Raises:
       ValueError: if `value` is neither `0` nor `1`.
    """
    global _GRAPH_LEARNING_PHASES
    assert value in {0, 1}
    assert tf.compat.v1.executing_eagerly_outside_functions()
    global_learning_phase_was_set = global_learning_phase_is_set()
    if global_learning_phase_was_set:
        previous_value = learning_phase()
    try:
        _internal_set_learning_phase(_DUMMY_EAGER_GRAPH.key, value)
        yield
    finally:
        # Restore learning phase to initial value or unset.
        if global_learning_phase_was_set:
            _internal_set_learning_phase(_DUMMY_EAGER_GRAPH.key, previous_value)
        else:
            del _GRAPH_LEARNING_PHASES[_DUMMY_EAGER_GRAPH.key]


def _as_graph_element(obj):
    """Convert `obj` to a graph element if possible, otherwise return `None`.

    Args:
      obj: Object to convert.

    Returns:
      The result of `obj._as_graph_element()` if that method is available;
          otherwise `None`.
    """
    conv_fn = getattr(obj, "_as_graph_element", None)
    if conv_fn and callable(conv_fn):
        return conv_fn()
    return None


def _assert_same_graph(original_item, item):
    """Fail if the 2 items are from different graphs.

    Args:
      original_item: Original item to check against.
      item: Item to check.

    Raises:
      ValueError: if graphs do not match.
    """
    original_graph = getattr(original_item, "graph", None)
    graph = getattr(item, "graph", None)
    if original_graph and graph and original_graph is not graph:
        raise ValueError(
            "%s must be from the same graph as %s (graphs are %s and %s)."
            % (item, original_item, graph, original_graph)
        )


def _current_graph(op_input_list, graph=None):
    """Returns the appropriate graph to use for the given inputs.

    This library method provides a consistent algorithm for choosing the graph
    in which an Operation should be constructed:

    1. If the default graph is being used to construct a function, we
       use the default graph.
    2. If the "graph" is specified explicitly, we validate that all of the
       inputs in "op_input_list" are compatible with that graph.
    3. Otherwise, we attempt to select a graph from the first Operation-
       or Tensor-valued input in "op_input_list", and validate that all other
       such inputs are in the same graph.
    4. If the graph was not specified and it could not be inferred from
       "op_input_list", we attempt to use the default graph.

    Args:
      op_input_list: A list of inputs to an operation, which may include
        `Tensor`, `Operation`, and other objects that may be converted to a
        graph element.
      graph: (Optional) The explicit graph to use.

    Raises:
      TypeError: If op_input_list is not a list or tuple, or if graph is not a
        Graph.
      ValueError: If a graph is explicitly passed and not all inputs are from
        it, or if the inputs are from multiple graphs, or we could not find a
        graph and there was no default graph.

    Returns:
      The appropriate graph to use for the given inputs.

    """
    current_default_graph = tf.compat.v1.get_default_graph()
    if current_default_graph.building_function:
        return current_default_graph

    op_input_list = tuple(op_input_list)  # Handle generators correctly
    if graph and not isinstance(graph, tf.Graph):
        raise TypeError(f"Input graph needs to be a Graph: {graph}")

    def _is_symbolic_tensor(tensor):
        if hasattr(tf, "is_symbolic_tensor"):
            return tf.is_symbolic_tensor(tensor)
        return type(tensor) == tf.Tensor  # noqa: E721

    # 1. We validate that all of the inputs are from the same graph. This is
    #    either the supplied graph parameter, or the first one selected from one
    #    the graph-element-valued inputs. In the latter case, we hold onto
    #    that input in original_graph_element so we can provide a more
    #    informative error if a mismatch is found.
    original_graph_element = None
    for op_input in op_input_list:
        if isinstance(
            op_input, (tf.Operation, tf.__internal__.CompositeTensor)
        ) or _is_symbolic_tensor(op_input):
            graph_element = op_input
        else:
            graph_element = _as_graph_element(op_input)

        if graph_element is not None:
            if not graph:
                original_graph_element = graph_element
                graph = getattr(graph_element, "graph", None)
            elif original_graph_element is not None:
                _assert_same_graph(original_graph_element, graph_element)
            elif graph_element.graph is not graph:
                raise ValueError(
                    f"{graph_element} is not from the passed-in graph."
                )

    # 2. If all else fails, we use the default graph, which is always there.
    return graph or current_default_graph


def _get_session(op_input_list=()):
    """Returns the session object for the current thread."""
    global _SESSION
    default_session = tf.compat.v1.get_default_session()
    if default_session is not None:
        session = default_session
    else:
        if tf.inside_function():
            raise RuntimeError(
                "Cannot get session inside Tensorflow graph function."
            )
        # If we don't have a session, or that session does not match the current
        # graph, create and cache a new session.
        if getattr(
            _SESSION, "session", None
        ) is None or _SESSION.session.graph is not _current_graph(
            op_input_list
        ):
            # If we are creating the Session inside a tf.distribute.Strategy
            # scope, we ask the strategy for the right session options to use.
            if tf.distribute.has_strategy():
                configure_and_create_distributed_session(
                    tf.distribute.get_strategy()
                )
            else:
                _SESSION.session = tf.compat.v1.Session(
                    config=get_default_session_config()
                )
        session = _SESSION.session
    return session


@keras_export(v1=["keras.backend.get_session"])
def get_session(op_input_list=()):
    """Returns the TF session to be used by the backend.

    If a default TensorFlow session is available, we will return it.

    Else, we will return the global TF-Keras session assuming it matches
    the current graph.

    If no global TF-Keras session exists at this point:
    we will create a new global session.

    Note that you can manually set the global session
    via `K.set_session(sess)`.

    Args:
        op_input_list: An option sequence of tensors or ops, which will be used
          to determine the current graph. Otherwise the default graph will be
          used.

    Returns:
        A TensorFlow session.
    """
    session = _get_session(op_input_list)
    if not _MANUAL_VAR_INIT:
        with session.graph.as_default():
            _initialize_variables(session)
    return session


# Inject the get_session function to keras_deps to remove the dependency
# from TFLite to TF-Keras.
tf.__internal__.register_get_session_function(get_session)

# Inject the get_session function to tracking_util to avoid the backward
# dependency from TF to TF-Keras.
tf.__internal__.tracking.register_session_provider(get_session)


def get_graph():
    if tf.executing_eagerly():
        global _GRAPH
        if not getattr(_GRAPH, "graph", None):
            _GRAPH.graph = tf.__internal__.FuncGraph("keras_graph")
        return _GRAPH.graph
    else:
        return tf.compat.v1.get_default_graph()


@tf_contextlib.contextmanager
def _scratch_graph(graph=None):
    """Retrieve a shared and temporary func graph.

    The eager execution path lifts a subgraph from the keras global graph into
    a scratch graph in order to create a function. DistributionStrategies, in
    turn, constructs multiple functions as well as a final combined function. In
    order for that logic to work correctly, all of the functions need to be
    created on the same scratch FuncGraph.

    Args:
      graph: A graph to be used as the current scratch graph. If not set then
        a scratch graph will either be retrieved or created:

    Yields:
      The current scratch graph.
    """
    global _CURRENT_SCRATCH_GRAPH
    scratch_graph = getattr(_CURRENT_SCRATCH_GRAPH, "graph", None)
    # If scratch graph and `graph` are both configured, they must match.
    if (
        scratch_graph is not None
        and graph is not None
        and scratch_graph is not graph
    ):
        raise ValueError("Multiple scratch graphs specified.")

    if scratch_graph:
        yield scratch_graph
        return

    graph = graph or tf.__internal__.FuncGraph("keras_scratch_graph")
    try:
        _CURRENT_SCRATCH_GRAPH.graph = graph
        yield graph
    finally:
        _CURRENT_SCRATCH_GRAPH.graph = None


@keras_export(v1=["keras.backend.set_session"])
def set_session(session):
    """Sets the global TensorFlow session.

    Args:
        session: A TF Session.
    """
    global _SESSION
    _SESSION.session = session


def get_default_session_config():
    if os.environ.get("OMP_NUM_THREADS"):
        logging.warning(
            "OMP_NUM_THREADS is no longer used by the default TF-Keras config. "
            "To configure the number of threads, use tf.config.threading APIs."
        )

    config = get_config()
    config.allow_soft_placement = True

    return config


def get_default_graph_uid_map():
    graph = tf.compat.v1.get_default_graph()
    name_uid_map = PER_GRAPH_OBJECT_NAME_UIDS.get(graph, None)
    if name_uid_map is None:
        name_uid_map = collections.defaultdict(int)
        PER_GRAPH_OBJECT_NAME_UIDS[graph] = name_uid_map
    return name_uid_map


# DEVICE MANIPULATION


class _TfDeviceCaptureOp:
    """Class for capturing the TF device scope."""

    def __init__(self):
        self.device = None

    def _set_device(self, device):
        """This method captures TF's explicit device scope setting."""
        if isinstance(device, tf.DeviceSpec):
            device = device.to_string()
        self.device = device

    def _set_device_from_string(self, device_str):
        self.device = device_str


def _get_current_tf_device():
    """Return explicit device of current context, otherwise returns `None`.

    Returns:
        If the current device scope is explicitly set, it returns a string with
        the device (`CPU` or `GPU`). If the scope is not explicitly set, it will
        return `None`.
    """
    graph = get_graph()
    op = _TfDeviceCaptureOp()
    graph._apply_device_functions(op)
    if tf.__internal__.tf2.enabled():
        return tf.DeviceSpec.from_string(op.device)
    else:
        return tf.compat.v1.DeviceSpec.from_string(op.device)


def _is_current_explicit_device(device_type):
    """Check if the current device is explicitly set to `device_type`.

    Args:
        device_type: A string containing `GPU` or `CPU` (case-insensitive).

    Returns:
        A boolean indicating if the current device scope is explicitly set on
        the device type.

    Raises:
        ValueError: If the `device_type` string indicates an unsupported device.
    """
    device_type = device_type.upper()
    if device_type not in ["CPU", "GPU"]:
        raise ValueError('`device_type` should be either "CPU" or "GPU".')
    device = _get_current_tf_device()
    return device is not None and device.device_type == device_type.upper()


def _get_available_gpus():
    """Get a list of available GPU devices (formatted as strings).

    Returns:
        A list of available GPU devices.
    """
    if tf.compat.v1.executing_eagerly_outside_functions():
        # Returns names of devices directly.
        return [d.name for d in tf.config.list_logical_devices("GPU")]

    global _LOCAL_DEVICES
    if _LOCAL_DEVICES is None:
        _LOCAL_DEVICES = get_session().list_devices()
    return [x.name for x in _LOCAL_DEVICES if x.device_type == "GPU"]


def _has_nchw_support():
    """Check whether the current scope supports NCHW ops.

    TensorFlow does not support NCHW on CPU. Therefore we check if we are not
    explicitly put on
    CPU, and have GPUs available. In this case there will be soft-placing on the
    GPU device.

    Returns:
        bool: if the current scope device placement would support nchw
    """
    explicitly_on_cpu = _is_current_explicit_device("CPU")
    gpus_available = bool(_get_available_gpus())
    return not explicitly_on_cpu and gpus_available


# VARIABLE MANIPULATION


def _constant_to_tensor(x, dtype):
    """Convert the input `x` to a tensor of type `dtype`.

    This is slightly faster than the _to_tensor function, at the cost of
    handling fewer cases.

    Args:
        x: An object to be converted (numpy arrays, floats, ints and lists of
          them).
        dtype: The destination type.

    Returns:
        A tensor.
    """
    return tf.constant(x, dtype=dtype)


def _to_tensor(x, dtype):
    """Convert the input `x` to a tensor of type `dtype`.

    Args:
        x: An object to be converted (numpy array, list, tensors).
        dtype: The destination type.

    Returns:
        A tensor.
    """
    return tf.convert_to_tensor(x, dtype=dtype)


@keras_export("keras.backend.is_sparse")
@doc_controls.do_not_generate_docs
def is_sparse(tensor):
    """Returns whether a tensor is a sparse tensor.

    Args:
        tensor: A tensor instance.

    Returns:
        A boolean.

    Example:


    >>> a = tf.keras.backend.placeholder((2, 2), sparse=False)
    >>> print(tf.keras.backend.is_sparse(a))
    False
    >>> b = tf.keras.backend.placeholder((2, 2), sparse=True)
    >>> print(tf.keras.backend.is_sparse(b))
    True

    """
    spec = getattr(tensor, "_type_spec", None)
    if spec is not None:
        return isinstance(spec, tf.SparseTensorSpec)
    return isinstance(tensor, tf.SparseTensor)


@keras_export("keras.backend.to_dense")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def to_dense(tensor):
    """Converts a sparse tensor into a dense tensor and returns it.

    Args:
        tensor: A tensor instance (potentially sparse).

    Returns:
        A dense tensor.

    Examples:


    >>> b = tf.keras.backend.placeholder((2, 2), sparse=True)
    >>> print(tf.keras.backend.is_sparse(b))
    True
    >>> c = tf.keras.backend.to_dense(b)
    >>> print(tf.keras.backend.is_sparse(c))
    False

    """
    if is_sparse(tensor):
        return tf.sparse.to_dense(tensor)
    else:
        return tensor


@keras_export("keras.backend.name_scope", v1=[])
@doc_controls.do_not_generate_docs
def name_scope(name):
    """A context manager for use when defining a Python op.

    This context manager pushes a name scope, which will make the name of all
    operations added within it have a prefix.

    For example, to define a new Python op called `my_op`:


    def my_op(a):
      with tf.name_scope("MyOp") as scope:
        a = tf.convert_to_tensor(a, name="a")
        # Define some computation that uses `a`.
        return foo_op(..., name=scope)


    When executed, the Tensor `a` will have the name `MyOp/a`.

    Args:
      name: The prefix to use on all names created within the name scope.

    Returns:
      Name scope context manager.
    """
    return tf.name_scope(name)


# Export V1 version.
_v1_name_scope = tf.compat.v1.name_scope
keras_export(v1=["keras.backend.name_scope"], allow_multiple_exports=True)(
    _v1_name_scope
)


@keras_export("keras.backend.variable")
@doc_controls.do_not_generate_docs
def variable(value, dtype=None, name=None, constraint=None):
    """Instantiates a variable and returns it.

    Args:
        value: Numpy array, initial value of the tensor.
        dtype: Tensor type.
        name: Optional name string for the tensor.
        constraint: Optional projection function to be
            applied to the variable after an optimizer update.

    Returns:
        A variable instance (with TF-Keras metadata included).

    Examples:

    >>> val = np.array([[1, 2], [3, 4]])
    >>> kvar = tf.keras.backend.variable(value=val, dtype='float64',
    ...                                  name='example_var')
    >>> tf.keras.backend.dtype(kvar)
    'float64'
    >>> print(kvar)
    <tf.Variable 'example_var:...' shape=(2, 2) dtype=float64, numpy=
      array([[1., 2.],
             [3., 4.]])>

    """
    if dtype is None:
        dtype = floatx()
    if hasattr(value, "tocoo"):
        sparse_coo = value.tocoo()
        indices = np.concatenate(
            (
                np.expand_dims(sparse_coo.row, 1),
                np.expand_dims(sparse_coo.col, 1),
            ),
            1,
        )
        v = tf.SparseTensor(
            indices=indices,
            values=sparse_coo.data,
            dense_shape=sparse_coo.shape,
        )
        v._keras_shape = sparse_coo.shape
        return v
    v = tf.Variable(
        value, dtype=tf.as_dtype(dtype), name=name, constraint=constraint
    )
    if isinstance(value, np.ndarray):
        v._keras_shape = value.shape
    elif hasattr(value, "shape"):
        v._keras_shape = int_shape(value)
    track_variable(v)
    return v


def track_tf_optimizer(tf_optimizer):
    """Tracks the given TF optimizer for initialization of its variables."""
    if tf.executing_eagerly():
        return
    optimizers = _GRAPH_TF_OPTIMIZERS[None]
    optimizers.add(tf_optimizer)


@keras_export("keras.__internal__.backend.track_variable", v1=[])
def track_variable(v):
    """Tracks the given variable for initialization."""
    if tf.executing_eagerly():
        return
    graph = v.graph if hasattr(v, "graph") else get_graph()
    _GRAPH_VARIABLES[graph].add(v)


def observe_object_name(name):
    """Observe a name and make sure it won't be used by `unique_object_name`."""
    OBSERVED_NAMES.add(name)


def unique_object_name(
    name,
    name_uid_map=None,
    avoid_names=None,
    namespace="",
    zero_based=False,
    avoid_observed_names=False,
):
    """Makes a object name (or any string) unique within a TF-Keras session.

    Args:
      name: String name to make unique.
      name_uid_map: An optional defaultdict(int) to use when creating unique
        names. If None (default), uses a per-Graph dictionary.
      avoid_names: An optional set or dict with names which should not be used.
        If None (default), don't avoid any names unless `avoid_observed_names`
        is True.
      namespace: Gets a name which is unique within the (graph, namespace).
        Layers which are not Networks use a blank namespace and so get
        graph-global names.
      zero_based: If True, name sequences start with no suffix (e.g. "dense",
        "dense_1"). If False, naming is one-based ("dense_1", "dense_2").
      avoid_observed_names: If True, avoid any names that have been observed by
        `backend.observe_object_name`.

    Returns:
      Unique string name.

    Example:


    unique_object_name('dense')  # dense_1
    unique_object_name('dense')  # dense_2

    """
    if name_uid_map is None:
        name_uid_map = get_default_graph_uid_map()
    if avoid_names is None:
        if avoid_observed_names:
            avoid_names = OBSERVED_NAMES
        else:
            avoid_names = set()
    proposed_name = None
    while proposed_name is None or proposed_name in avoid_names:
        name_key = (namespace, name)
        if zero_based:
            number = name_uid_map[name_key]
            if number:
                proposed_name = name + "_" + str(number)
            else:
                proposed_name = name
            name_uid_map[name_key] += 1
        else:
            name_uid_map[name_key] += 1
            proposed_name = name + "_" + str(name_uid_map[name_key])
    return proposed_name


def _get_variables(graph=None):
    """Returns variables corresponding to the given graph for initialization."""
    assert not tf.executing_eagerly()
    variables = _GRAPH_VARIABLES[graph]
    for opt in _GRAPH_TF_OPTIMIZERS[graph]:
        variables.update(opt.optimizer.variables())
    return variables


@keras_export("keras.__internal__.backend.initialize_variables", v1=[])
def _initialize_variables(session):
    """Utility to initialize uninitialized variables on the fly."""
    variables = _get_variables(get_graph())
    candidate_vars = []
    for v in variables:
        if not getattr(v, "_keras_initialized", False):
            candidate_vars.append(v)
    if candidate_vars:
        # This step is expensive, so we only run it on variables not already
        # marked as initialized.
        is_initialized = session.run(
            [tf.compat.v1.is_variable_initialized(v) for v in candidate_vars]
        )
        # TODO(kathywu): Some metric variables loaded from SavedModel are never
        # actually used, and do not have an initializer.
        should_be_initialized = [
            (not is_initialized[n]) and v.initializer is not None
            for n, v in enumerate(candidate_vars)
        ]
        uninitialized_vars = []
        for flag, v in zip(should_be_initialized, candidate_vars):
            if flag:
                uninitialized_vars.append(v)
            v._keras_initialized = True
        if uninitialized_vars:
            session.run(tf.compat.v1.variables_initializer(uninitialized_vars))


@keras_export("keras.backend.constant")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def constant(value, dtype=None, shape=None, name=None):
    """Creates a constant tensor.

    Args:
        value: A constant value (or list)
        dtype: The type of the elements of the resulting tensor.
        shape: Optional dimensions of resulting tensor.
        name: Optional name for the tensor.

    Returns:
        A Constant Tensor.
    """
    if dtype is None:
        dtype = floatx()

    return tf.constant(value, dtype=dtype, shape=shape, name=name)


@keras_export("keras.backend.is_keras_tensor")
def is_keras_tensor(x):
    """Returns whether `x` is a TF-Keras tensor.

    A "Keras tensor" is a tensor that was returned by a TF-Keras layer,
    (`Layer` class) or by `Input`.

    Args:
        x: A candidate tensor.

    Returns:
        A boolean: Whether the argument is a TF-Keras tensor.

    Raises:
        ValueError: In case `x` is not a symbolic tensor.

    Examples:

    >>> np_var = np.array([1, 2])
    >>> # A numpy array is not a symbolic tensor.
    >>> tf.keras.backend.is_keras_tensor(np_var)
    Traceback (most recent call last):
    ...
    ValueError: Unexpectedly found an instance of type
    `<class 'numpy.ndarray'>`.
    Expected a symbolic tensor instance.
    >>> keras_var = tf.keras.backend.variable(np_var)
    >>> # A variable created with the keras backend is not a TF-Keras tensor.
    >>> tf.keras.backend.is_keras_tensor(keras_var)
    False
    >>> keras_placeholder = tf.keras.backend.placeholder(shape=(2, 4, 5))
    >>> # A placeholder is a TF-Keras tensor.
    >>> tf.keras.backend.is_keras_tensor(keras_placeholder)
    True
    >>> keras_input = tf.keras.layers.Input([10])
    >>> # An Input is a TF-Keras tensor.
    >>> tf.keras.backend.is_keras_tensor(keras_input)
    True
    >>> keras_layer_output = tf.keras.layers.Dense(10)(keras_input)
    >>> # Any TF-Keras layer output is a TF-Keras tensor.
    >>> tf.keras.backend.is_keras_tensor(keras_layer_output)
    True

    """
    if not isinstance(
        x,
        (
            tf.Tensor,
            tf.Variable,
            tf.SparseTensor,
            tf.RaggedTensor,
            keras_tensor.KerasTensor,
        ),
    ):
        raise ValueError(
            "Unexpectedly found an instance of type `"
            + str(type(x))
            + "`. Expected a symbolic tensor instance."
        )
    if tf.compat.v1.executing_eagerly_outside_functions():
        return isinstance(x, keras_tensor.KerasTensor)
    return hasattr(x, "_keras_history")


@keras_export("keras.backend.placeholder")
@doc_controls.do_not_generate_docs
def placeholder(
    shape=None, ndim=None, dtype=None, sparse=False, name=None, ragged=False
):
    """Instantiates a placeholder tensor and returns it.

    Args:
        shape: Shape of the placeholder
            (integer tuple, may include `None` entries).
        ndim: Number of axes of the tensor.
            At least one of {`shape`, `ndim`} must be specified.
            If both are specified, `shape` is used.
        dtype: Placeholder type.
        sparse: Boolean, whether the placeholder should have a sparse type.
        name: Optional name string for the placeholder.
        ragged: Boolean, whether the placeholder should have a ragged type.
            In this case, values of 'None' in the 'shape' argument represent
            ragged dimensions. For more information about RaggedTensors, see
            this [guide](https://www.tensorflow.org/guide/ragged_tensor).

    Raises:
        ValueError: If called with sparse = True and ragged = True.

    Returns:
        Tensor instance (with TF-Keras metadata included).

    Examples:


    >>> input_ph = tf.keras.backend.placeholder(shape=(2, 4, 5))
    >>> input_ph
    <KerasTensor: shape=(2, 4, 5) dtype=float32 (created by layer ...)>

    """
    if sparse and ragged:
        raise ValueError(
            "Cannot set both sparse and ragged to "
            "True when creating a placeholder."
        )
    if dtype is None:
        dtype = floatx()
    if not shape:
        if ndim:
            shape = (None,) * ndim
    if tf.compat.v1.executing_eagerly_outside_functions():
        if sparse:
            spec = tf.SparseTensorSpec(shape=shape, dtype=dtype)
        elif ragged:
            ragged_rank = 0
            for i in range(1, len(shape)):
                # Hacky because could be tensorshape or tuple maybe?
                # Or just tensorshape?
                if shape[i] is None or (
                    hasattr(shape[i], "value") and shape[i].value is None
                ):
                    ragged_rank = i
            spec = tf.RaggedTensorSpec(
                shape=shape, dtype=dtype, ragged_rank=ragged_rank
            )
        else:
            spec = tf.TensorSpec(shape=shape, dtype=dtype, name=name)
        x = keras_tensor.keras_tensor_from_type_spec(spec, name=name)
    else:
        with get_graph().as_default():
            if sparse:
                x = tf.compat.v1.sparse_placeholder(
                    dtype, shape=shape, name=name
                )
            elif ragged:
                ragged_rank = 0
                for i in range(1, len(shape)):
                    if shape[i] is None:
                        ragged_rank = i
                type_spec = tf.RaggedTensorSpec(
                    shape=shape, dtype=dtype, ragged_rank=ragged_rank
                )

                def tensor_spec_to_placeholder(tensorspec):
                    return tf.compat.v1.placeholder(
                        tensorspec.dtype, tensorspec.shape
                    )

                x = tf.nest.map_structure(
                    tensor_spec_to_placeholder,
                    type_spec,
                    expand_composites=True,
                )
            else:
                x = tf.compat.v1.placeholder(dtype, shape=shape, name=name)

    if tf.executing_eagerly():
        # Add keras_history connectivity information to the placeholder
        # when the placeholder is built in a top-level eager context
        # (intended to be used with keras.backend.function)
        from tf_keras.engine import (
            input_layer,
        )

        x = input_layer.Input(tensor=x)
        x._is_backend_placeholder = True

    return x


def is_placeholder(x):
    """Returns whether `x` is a placeholder.

    Args:
        x: A candidate placeholder.

    Returns:
        Boolean.
    """
    try:
        if tf.compat.v1.executing_eagerly_outside_functions():
            return hasattr(x, "_is_backend_placeholder")

        # TODO(b/246438937): Remove the special case for tf.Variable once
        # tf.Variable becomes CompositeTensor and will be expanded into
        # dt_resource tensors.
        if tf_utils.is_extension_type(x) and not isinstance(x, tf.Variable):
            flat_components = tf.nest.flatten(x, expand_composites=True)
            return py_any(is_placeholder(c) for c in flat_components)
        else:
            return x.op.type == "Placeholder"
    except AttributeError:
        return False


@keras_export("keras.backend.shape")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def shape(x):
    """Returns the symbolic shape of a tensor or variable.

    Args:
        x: A tensor or variable.

    Returns:
        A symbolic shape (which is itself a tensor).

    Examples:

    >>> val = np.array([[1, 2], [3, 4]])
    >>> kvar = tf.keras.backend.variable(value=val)
    >>> tf.keras.backend.shape(kvar)
    <tf.Tensor: shape=(2,), dtype=int32, numpy=array([2, 2], dtype=int32)>
    >>> input = tf.keras.backend.placeholder(shape=(2, 4, 5))
    >>> tf.keras.backend.shape(input)
    <KerasTensor: shape=(3,) dtype=int32 inferred_value=[2, 4, 5] ...>

    """
    return tf.shape(x)


@keras_export("keras.backend.int_shape")
@doc_controls.do_not_generate_docs
def int_shape(x):
    """Returns shape of tensor/variable as a tuple of int/None entries.

    Args:
        x: Tensor or variable.

    Returns:
        A tuple of integers (or None entries).

    Examples:

    >>> input = tf.keras.backend.placeholder(shape=(2, 4, 5))
    >>> tf.keras.backend.int_shape(input)
    (2, 4, 5)
    >>> val = np.array([[1, 2], [3, 4]])
    >>> kvar = tf.keras.backend.variable(value=val)
    >>> tf.keras.backend.int_shape(kvar)
    (2, 2)

    """
    try:
        shape = x.shape
        if not isinstance(shape, tuple):
            shape = tuple(shape.as_list())
        return shape
    except ValueError:
        return None


@keras_export("keras.backend.ndim")
@doc_controls.do_not_generate_docs
def ndim(x):
    """Returns the number of axes in a tensor, as an integer.

    Args:
        x: Tensor or variable.

    Returns:
        Integer (scalar), number of axes.

    Examples:


    >>> input = tf.keras.backend.placeholder(shape=(2, 4, 5))
    >>> val = np.array([[1, 2], [3, 4]])
    >>> kvar = tf.keras.backend.variable(value=val)
    >>> tf.keras.backend.ndim(input)
    3
    >>> tf.keras.backend.ndim(kvar)
    2

    """
    return x.shape.rank


@keras_export("keras.backend.dtype")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def dtype(x):
    """Returns the dtype of a TF-Keras tensor or variable, as a string.

    Args:
        x: Tensor or variable.

    Returns:
        String, dtype of `x`.

    Examples:

    >>> tf.keras.backend.dtype(tf.keras.backend.placeholder(shape=(2,4,5)))
    'float32'
    >>> tf.keras.backend.dtype(tf.keras.backend.placeholder(shape=(2,4,5),
    ...                                                     dtype='float32'))
    'float32'
    >>> tf.keras.backend.dtype(tf.keras.backend.placeholder(shape=(2,4,5),
    ...                                                     dtype='float64'))
    'float64'
    >>> kvar = tf.keras.backend.variable(np.array([[1, 2], [3, 4]]))
    >>> tf.keras.backend.dtype(kvar)
    'float32'
    >>> kvar = tf.keras.backend.variable(np.array([[1, 2], [3, 4]]),
    ...                                  dtype='float32')
    >>> tf.keras.backend.dtype(kvar)
    'float32'

    """
    return x.dtype.base_dtype.name


@doc_controls.do_not_generate_docs
def dtype_numpy(x):
    """Returns the numpy dtype of a TF-Keras tensor or variable.

    Args:
        x: Tensor or variable.

    Returns:
        numpy.dtype, dtype of `x`.
    """
    return tf.as_dtype(x.dtype).as_numpy_dtype


@keras_export("keras.backend.eval")
@doc_controls.do_not_generate_docs
def eval(x):
    """Evaluates the value of a variable.

    Args:
        x: A variable.

    Returns:
        A Numpy array.

    Examples:

    >>> kvar = tf.keras.backend.variable(np.array([[1, 2], [3, 4]]),
    ...                                  dtype='float32')
    >>> tf.keras.backend.eval(kvar)
    array([[1.,  2.],
           [3.,  4.]], dtype=float32)

    """
    return get_value(to_dense(x))


@keras_export("keras.backend.zeros")
@doc_controls.do_not_generate_docs
def zeros(shape, dtype=None, name=None):
    """Instantiates an all-zeros variable and returns it.

    Args:
        shape: Tuple or list of integers, shape of returned TF-Keras variable
        dtype: data type of returned TF-Keras variable
        name: name of returned TF-Keras variable

    Returns:
        A variable (including TF-Keras metadata), filled with `0.0`.
        Note that if `shape` was symbolic, we cannot return a variable,
        and will return a dynamically-shaped tensor instead.

    Example:

    >>> kvar = tf.keras.backend.zeros((3,4))
    >>> tf.keras.backend.eval(kvar)
    array([[0.,  0.,  0.,  0.],
           [0.,  0.,  0.,  0.],
           [0.,  0.,  0.,  0.]], dtype=float32)
    >>> A = tf.constant([1,2,3])
    >>> kvar2 = tf.keras.backend.zeros(A.shape) # [0., 0., 0.]
    >>> tf.keras.backend.eval(kvar2)
    array([0., 0., 0.], dtype=float32)
    >>> kvar3 = tf.keras.backend.zeros(A.shape,dtype=tf.int32)
    >>> tf.keras.backend.eval(kvar3)
    array([0, 0, 0], dtype=int32)
    >>> kvar4 = tf.keras.backend.zeros([2,3])
    >>> tf.keras.backend.eval(kvar4)
    array([[0., 0., 0.],
           [0., 0., 0.]], dtype=float32)

    """
    with tf.init_scope():
        if dtype is None:
            dtype = floatx()
        tf_dtype = tf.as_dtype(dtype)
        v = tf.zeros(shape=shape, dtype=tf_dtype, name=name)
        if py_all(v.shape.as_list()):
            return variable(v, dtype=dtype, name=name)
        return v


@keras_export("keras.backend.ones")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def ones(shape, dtype=None, name=None):
    """Instantiates an all-ones variable and returns it.

    Args:
        shape: Tuple of integers, shape of returned TF-Keras variable.
        dtype: String, data type of returned TF-Keras variable.
        name: String, name of returned TF-Keras variable.

    Returns:
        A TF-Keras variable, filled with `1.0`.
        Note that if `shape` was symbolic, we cannot return a variable,
        and will return a dynamically-shaped tensor instead.

    Example:


    >>> kvar = tf.keras.backend.ones((3,4))
    >>> tf.keras.backend.eval(kvar)
    array([[1.,  1.,  1.,  1.],
           [1.,  1.,  1.,  1.],
           [1.,  1.,  1.,  1.]], dtype=float32)

    """
    with tf.init_scope():
        if dtype is None:
            dtype = floatx()
        tf_dtype = tf.as_dtype(dtype)
        v = tf.ones(shape=shape, dtype=tf_dtype, name=name)
        if py_all(v.shape.as_list()):
            return variable(v, dtype=dtype, name=name)
        return v


@keras_export("keras.backend.eye")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def eye(size, dtype=None, name=None):
    """Instantiate an identity matrix and returns it.

    Args:
        size: Integer, number of rows/columns.
        dtype: String, data type of returned TF-Keras variable.
        name: String, name of returned TF-Keras variable.

    Returns:
        A TF-Keras variable, an identity matrix.

    Example:


    >>> kvar = tf.keras.backend.eye(3)
    >>> tf.keras.backend.eval(kvar)
    array([[1.,  0.,  0.],
           [0.,  1.,  0.],
           [0.,  0.,  1.]], dtype=float32)


    """
    if dtype is None:
        dtype = floatx()
    tf_dtype = tf.as_dtype(dtype)
    return variable(tf.eye(size, dtype=tf_dtype), dtype, name)


@keras_export("keras.backend.zeros_like")
@doc_controls.do_not_generate_docs
def zeros_like(x, dtype=None, name=None):
    """Instantiates an all-zeros variable of the same shape as another tensor.

    Args:
        x: TF-Keras variable or TF-Keras tensor.
        dtype: dtype of returned TF-Keras variable.
               `None` uses the dtype of `x`.
        name: name for the variable to create.

    Returns:
        A TF-Keras variable with the shape of `x` filled with zeros.

    Example:

    ```python
    kvar = tf.keras.backend.variable(np.random.random((2,3)))
    kvar_zeros = tf.keras.backend.zeros_like(kvar)
    K.eval(kvar_zeros)
    # array([[ 0.,  0.,  0.], [ 0.,  0.,  0.]], dtype=float32)
    ```
    """
    return tf.zeros_like(x, dtype=dtype, name=name)


@keras_export("keras.backend.ones_like")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def ones_like(x, dtype=None, name=None):
    """Instantiates an all-ones variable of the same shape as another tensor.

    Args:
        x: TF-Keras variable or tensor.
        dtype: String, dtype of returned TF-Keras variable.
             None uses the dtype of x.
        name: String, name for the variable to create.

    Returns:
        A TF-Keras variable with the shape of x filled with ones.

    Example:

    >>> kvar = tf.keras.backend.variable(np.random.random((2,3)))
    >>> kvar_ones = tf.keras.backend.ones_like(kvar)
    >>> tf.keras.backend.eval(kvar_ones)
    array([[1.,  1.,  1.],
           [1.,  1.,  1.]], dtype=float32)

    """
    return tf.ones_like(x, dtype=dtype, name=name)


def identity(x, name=None):
    """Returns a tensor with the same content as the input tensor.

    Args:
        x: The input tensor.
        name: String, name for the variable to create.

    Returns:
        A tensor of the same shape, type and content.
    """
    return tf.identity(x, name=name)


# Global flag to enforce tf.random.Generator for RandomGenerator.
# When this is enabled, for any caller to RandomGenerator, it will use
# tf.random.Generator to generate random numbers.
# The legacy behavior is to use TF's legacy stateful RNG ops like
# tf.random.uniform.
_USE_GENERATOR_FOR_RNG = False

# The global generator to create the seed when initializing the
# tf.random.Genrator used by RandomGenerator. When tf.random.Generator becomes
# the default solution, we would like the it to be initialized in a controlable
# way, so that each client of the program could start with same seed. This is
# very important for certain use case that requires all the client to have their
# state in sync. This instance will be set when user call
# `tf.keras.utils.set_random_seed()`
_SEED_GENERATOR = threading.local()


@keras_export(
    "keras.backend.experimental.is_tf_random_generator_enabled", v1=[]
)
def is_tf_random_generator_enabled():
    """Check whether `tf.random.Generator` is used for RNG in TF-Keras.

    Compared to existing TF stateful random ops, `tf.random.Generator` uses
    `tf.Variable` and stateless random ops to generate random numbers,
    which leads to better reproducibility in distributed training.
    Note enabling it might introduce some breakage to existing code,
    by producing differently-seeded random number sequences
    and breaking tests that rely on specific random numbers being generated.
    To disable the
    usage of `tf.random.Generator`, please use
    `tf.keras.backend.experimental.disable_random_generator`.

    We expect the `tf.random.Generator` code path to become the default, and
    will remove the legacy stateful random ops such as `tf.random.uniform` in
    the future (see the [TF RNG guide](
    https://www.tensorflow.org/guide/random_numbers)).

    This API will also be removed in a future release as well, together with
    `tf.keras.backend.experimental.enable_tf_random_generator()` and
    `tf.keras.backend.experimental.disable_tf_random_generator()`

    Returns:
      boolean: whether `tf.random.Generator` is used for random number
        generation in TF-Keras.
    """
    return _USE_GENERATOR_FOR_RNG


@keras_export("keras.backend.experimental.enable_tf_random_generator", v1=[])
def enable_tf_random_generator():
    """Enable the `tf.random.Generator` as the RNG for TF-Keras.

    See `tf.keras.backend.experimental.is_tf_random_generator_enabled` for more
    details.
    """

    global _USE_GENERATOR_FOR_RNG
    _USE_GENERATOR_FOR_RNG = True


@keras_export("keras.backend.experimental.disable_tf_random_generator", v1=[])
def disable_tf_random_generator():
    """Disable the `tf.random.Generator` as the RNG for TF-Keras.

    See `tf.keras.backend.experimental.is_tf_random_generator_enabled` for more
    details.
    """
    global _USE_GENERATOR_FOR_RNG
    _USE_GENERATOR_FOR_RNG = False


class RandomGenerator(tf.__internal__.tracking.AutoTrackable):
    """Random generator that selects appropriate random ops.

    This class contains the logic for legacy stateful random ops, as well as the
    new stateless random ops with seeds and tf.random.Generator. Any class that
    relies on RNG (eg initializer, shuffle, dropout) should use this class to
    handle the transition from legacy RNGs to new RNGs.

    Args:
      seed: Optional int seed. When `rng_type` is "stateful", the seed is used
        to create `tf.random.Generator` to produce deterministic sequences.
        When `rng_type` is "stateless", new seed will be created if it is not
        provided by user, and it will be passed down to stateless random ops.
        When `rng_type` is "legacy_stateful", the seed will be passed down to
        stateful random ops.
      rng_type: Type of RNG to use, one of "stateful", "stateless",
        "legacy_stateful". When `None` it uses "stateful" if
        `enable_tf_random_generator` has been activated, or
        "legacy_stateful" otherwise.
        - When using "stateless", the random ops outputs are constant (the same
          inputs result in the same outputs).
        - When using "stateful" or "legacy_stateful", the random ops outputs are
          non-constant, but deterministic: calling the same random op multiple
          times with the same inputs results in a deterministic sequence of
          different outputs.
        - "legacy_stateful" is backed by TF1 stateful RNG ops
          (e.g. `tf.random.uniform`), while "stateful"
          is backed by TF2 APIs (e.g. `tf.random.Generator.uniform`).
        Defaults to `None`.
    """

    RNG_STATELESS = "stateless"
    RNG_STATEFUL = "stateful"
    RNG_LEGACY_STATEFUL = "legacy_stateful"

    def __init__(self, seed=None, rng_type=None, **kwargs):
        self._seed = seed
        self._set_rng_type(rng_type, **kwargs)
        self._built = False

    def _set_rng_type(self, rng_type, **kwargs):
        # Only supported kwargs is "force_generator", which we will remove once
        # we clean up all the caller.
        # TODO(scottzhu): Remove the kwargs for force_generator.
        if kwargs.get("force_generator", False):
            rng_type = self.RNG_STATEFUL
        if rng_type is None:
            if is_tf_random_generator_enabled():
                self._rng_type = self.RNG_STATEFUL
            else:
                self._rng_type = self.RNG_LEGACY_STATEFUL
        else:
            if rng_type not in [
                self.RNG_STATEFUL,
                self.RNG_LEGACY_STATEFUL,
                self.RNG_STATELESS,
            ]:
                raise ValueError(
                    "Invalid `rng_type` received. "
                    'Valid `rng_type` are ["stateless", '
                    '"stateful", "legacy_stateful"].'
                    f" Got: {rng_type}"
                )
            self._rng_type = rng_type

    def _maybe_init(self):
        """Lazily init the RandomGenerator.

        The TF API executing_eagerly_outside_functions() has some side effect,
        and couldn't be used before API like tf.enable_eager_execution(). Some
        of the client side code was creating the initializer at the code load
        time, which triggers the creation of RandomGenerator. Lazy init this
        class to walkaround this issue until it is resolved on TF side.
        """
        # TODO(b/167482354): Change this back to normal init when the bug is
        # fixed.
        if self._built:
            return

        if (
            self._rng_type == self.RNG_STATEFUL
            and not tf.compat.v1.executing_eagerly_outside_functions()
        ):
            # Fall back to legacy stateful since the generator need to work in
            # tf2.
            self._rng_type = self.RNG_LEGACY_STATEFUL

        if self._rng_type == self.RNG_STATELESS:
            self._seed = self._create_seed(self._seed)
            self._generator = None
        elif self._rng_type == self.RNG_STATEFUL:
            with tf_utils.maybe_init_scope(self):
                seed = self._create_seed(self._seed)
                self._generator = tf.random.Generator.from_seed(
                    seed, alg=tf.random.Algorithm.AUTO_SELECT
                )
        else:
            # In legacy stateful, we use stateful op, regardless whether user
            # provide seed or not. Seeded stateful op will ensure generating
            # same sequences.
            self._generator = None
        self._built = True

    def make_seed_for_stateless_op(self):
        """Generate a new seed based on the init config.

        Note that this will not return python ints which will be frozen in the
        graph and cause stateless op to return the same value. It will only
        return value when generator is used, otherwise it will return None.

        Returns:
          A tensor with shape [2,].
        """
        self._maybe_init()
        if self._rng_type == self.RNG_STATELESS:
            return [self._seed, 0]
        elif self._rng_type == self.RNG_STATEFUL:
            return self._generator.make_seeds()[:, 0]
        return None

    def make_legacy_seed(self):
        """Create a new seed for the legacy stateful ops to use.

        When user didn't provide any original seed, this method will return
        None.  Otherwise it will increment the counter and return as the new
        seed.

        Note that it is important to generate different seed for stateful ops in
        the `tf.function`. The random ops will return same value when same seed
        is provided in the `tf.function`.

        Returns:
          int as new seed, or None.
        """
        if self._seed is not None:
            result = self._seed
            self._seed += 1
            return result
        return None

    def _create_seed(self, user_specified_seed):
        if user_specified_seed is not None:
            return user_specified_seed
        elif getattr(_SEED_GENERATOR, "generator", None):
            return _SEED_GENERATOR.generator.randint(1, int(1e9))
        else:
            return random.randint(1, int(1e9))

    def random_normal(
        self, shape, mean=0.0, stddev=1.0, dtype=None, nonce=None
    ):
        """Produce random number based on the normal distribution.

        Args:
          shape: The shape of the random values to generate.
          mean: Floats, default to 0. Mean of the random values to generate.
          stddev: Floats, default to 1. Standard deviation of the random values
            to generate.
          dtype: Optional dtype of the tensor. Only floating point types are
            supported. If not specified, `tf.keras.backend.floatx()` is used,
            which default to `float32` unless you configured it otherwise (via
            `tf.keras.backend.set_floatx(float_dtype)`)
          nonce: Optional integer scalar, that will be folded into the seed in
            the stateless mode.
        """
        self._maybe_init()
        dtype = dtype or floatx()
        if self._rng_type == self.RNG_STATEFUL:
            return self._generator.normal(
                shape=shape, mean=mean, stddev=stddev, dtype=dtype
            )
        elif self._rng_type == self.RNG_STATELESS:
            seed = self.make_seed_for_stateless_op()
            if nonce:
                seed = tf.random.experimental.stateless_fold_in(seed, nonce)
            return tf.random.stateless_normal(
                shape=shape, mean=mean, stddev=stddev, dtype=dtype, seed=seed
            )
        return tf.random.normal(
            shape=shape,
            mean=mean,
            stddev=stddev,
            dtype=dtype,
            seed=self.make_legacy_seed(),
        )

    def random_uniform(
        self, shape, minval=0.0, maxval=None, dtype=None, nonce=None
    ):
        """Produce random number based on the uniform distribution.

        Args:
          shape: The shape of the random values to generate.
          minval: Floats, default to 0. Lower bound of the range of
            random values to generate (inclusive).
          minval: Floats, default to None. Upper bound of the range of
            random values to generate (exclusive).
          dtype: Optional dtype of the tensor. Only floating point types are
            supported. If not specified, `tf.keras.backend.floatx()` is used,
            which default to `float32` unless you configured it otherwise (via
            `tf.keras.backend.set_floatx(float_dtype)`)
          nonce: Optional integer scalar, that will be folded into the seed in
            the stateless mode.
        """
        self._maybe_init()
        dtype = dtype or floatx()
        if self._rng_type == self.RNG_STATEFUL:
            return self._generator.uniform(
                shape=shape, minval=minval, maxval=maxval, dtype=dtype
            )
        elif self._rng_type == self.RNG_STATELESS:
            seed = self.make_seed_for_stateless_op()
            if nonce:
                seed = tf.random.experimental.stateless_fold_in(seed, nonce)
            return tf.random.stateless_uniform(
                shape=shape,
                minval=minval,
                maxval=maxval,
                dtype=dtype,
                seed=seed,
            )
        return tf.random.uniform(
            shape=shape,
            minval=minval,
            maxval=maxval,
            dtype=dtype,
            seed=self.make_legacy_seed(),
        )

    def truncated_normal(
        self, shape, mean=0.0, stddev=1.0, dtype=None, nonce=None
    ):
        """Produce random number based on the truncated normal distribution.

        Args:
          shape: The shape of the random values to generate.
          mean: Floats, default to 0. Mean of the random values to generate.
          stddev: Floats, default to 1. Standard deviation of the random values
            to generate.
          dtype: Optional dtype of the tensor. Only floating point types are
            supported. If not specified, `tf.keras.backend.floatx()` is used,
            which default to `float32` unless you configured it otherwise (via
            `tf.keras.backend.set_floatx(float_dtype)`)
          nonce: Optional integer scalar, that will be folded into the seed in
            the stateless mode.
        """
        self._maybe_init()
        dtype = dtype or floatx()
        if self._rng_type == self.RNG_STATEFUL:
            return self._generator.truncated_normal(
                shape=shape, mean=mean, stddev=stddev, dtype=dtype
            )
        elif self._rng_type == self.RNG_STATELESS:
            seed = self.make_seed_for_stateless_op()
            if nonce:
                seed = tf.random.experimental.stateless_fold_in(seed, nonce)
            return tf.random.stateless_truncated_normal(
                shape=shape, mean=mean, stddev=stddev, dtype=dtype, seed=seed
            )
        return tf.random.truncated_normal(
            shape=shape,
            mean=mean,
            stddev=stddev,
            dtype=dtype,
            seed=self.make_legacy_seed(),
        )

    def dropout(self, inputs, rate, noise_shape=None):
        self._maybe_init()
        if self._rng_type == self.RNG_STATEFUL:
            return tf.nn.experimental.general_dropout(
                inputs,
                rate=rate,
                noise_shape=noise_shape,
                uniform_sampler=self._generator.uniform,
            )
        elif self._rng_type == self.RNG_STATELESS:
            return tf.nn.experimental.stateless_dropout(
                inputs,
                rate=rate,
                noise_shape=noise_shape,
                seed=self.make_seed_for_stateless_op(),
            )
        else:
            return tf.nn.dropout(
                inputs,
                rate=rate,
                noise_shape=noise_shape,
                seed=self.make_legacy_seed(),
            )


@keras_export("keras.backend.random_uniform_variable")
@doc_controls.do_not_generate_docs
def random_uniform_variable(shape, low, high, dtype=None, name=None, seed=None):
    """Instantiates a variable with values drawn from a uniform distribution.

    Args:
        shape: Tuple of integers, shape of returned TF-Keras variable.
        low: Float, lower boundary of the output interval.
        high: Float, upper boundary of the output interval.
        dtype: String, dtype of returned TF-Keras variable.
        name: String, name of returned TF-Keras variable.
        seed: Integer, random seed.

    Returns:
        A TF-Keras variable, filled with drawn samples.

    Example:

    >>> kvar = tf.keras.backend.random_uniform_variable(shape=(2,3),
    ... low=0.0, high=1.0)
    >>> kvar
    <tf.Variable 'Variable:0' shape=(2, 3) dtype=float32, numpy=...,
    dtype=float32)>
    """
    if dtype is None:
        dtype = floatx()
    tf_dtype = tf.as_dtype(dtype)
    if seed is None:
        # ensure that randomness is conditioned by the Numpy RNG
        seed = np.random.randint(10e8)
    value = tf.compat.v1.random_uniform_initializer(
        low, high, dtype=tf_dtype, seed=seed
    )(shape)
    return variable(value, dtype=dtype, name=name)


@keras_export("keras.backend.random_normal_variable")
@doc_controls.do_not_generate_docs
def random_normal_variable(
    shape, mean, scale, dtype=None, name=None, seed=None
):
    """Instantiates a variable with values drawn from a normal distribution.

    Args:
        shape: Tuple of integers, shape of returned TF-Keras variable.
        mean: Float, mean of the normal distribution.
        scale: Float, standard deviation of the normal distribution.
        dtype: String, dtype of returned TF-Keras variable.
        name: String, name of returned TF-Keras variable.
        seed: Integer, random seed.

    Returns:
        A TF-Keras variable, filled with drawn samples.

    Example:

    >>> kvar = tf.keras.backend.random_normal_variable(shape=(2,3),
    ... mean=0.0, scale=1.0)
    >>> kvar
    <tf.Variable 'Variable:0' shape=(2, 3) dtype=float32, numpy=...,
    dtype=float32)>
    """
    if dtype is None:
        dtype = floatx()
    tf_dtype = tf.as_dtype(dtype)
    if seed is None:
        # ensure that randomness is conditioned by the Numpy RNG
        seed = np.random.randint(10e8)
    value = tf.compat.v1.random_normal_initializer(
        mean, scale, dtype=tf_dtype, seed=seed
    )(shape)
    return variable(value, dtype=dtype, name=name)


@keras_export("keras.backend.count_params")
@doc_controls.do_not_generate_docs
def count_params(x):
    """Returns the static number of elements in a variable or tensor.

    Args:
        x: Variable or tensor.

    Returns:
        Integer, the number of scalars in `x`.

    Example:

    >>> kvar = tf.keras.backend.zeros((2,3))
    >>> tf.keras.backend.count_params(kvar)
    6
    >>> tf.keras.backend.eval(kvar)
    array([[0.,  0.,  0.],
           [0.,  0.,  0.]], dtype=float32)

    """
    return np.prod(x.shape.as_list())


@keras_export("keras.backend.cast")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def cast(x, dtype):
    """Casts a tensor to a different dtype and returns it.

    You can cast a TF-Keras variable but it still returns a TF-Keras tensor.

    Args:
        x: TF-Keras tensor (or variable).
        dtype: String, either (`'float16'`, `'float32'`, or `'float64'`).

    Returns:
        TF-Keras tensor with dtype `dtype`.

    Examples:
        Cast a float32 variable to a float64 tensor

    >>> input = tf.keras.backend.ones(shape=(1,3))
    >>> print(input)
    <tf.Variable 'Variable:0' shape=(1, 3) dtype=float32,
    numpy=array([[1., 1., 1.]], dtype=float32)>
    >>> cast_input = tf.keras.backend.cast(input, dtype='float64')
    >>> print(cast_input)
    tf.Tensor([[1. 1. 1.]], shape=(1, 3), dtype=float64)

    """
    return tf.cast(x, dtype)


# UPDATES OPS


@keras_export("keras.backend.update")
@doc_controls.do_not_generate_docs
def update(x, new_x):
    return tf.compat.v1.assign(x, new_x)


@keras_export("keras.backend.update_add")
@doc_controls.do_not_generate_docs
def update_add(x, increment):
    """Update the value of `x` by adding `increment`.

    Args:
        x: A Variable.
        increment: A tensor of same shape as `x`.

    Returns:
        The variable `x` updated.
    """
    return tf.compat.v1.assign_add(x, increment)


@keras_export("keras.backend.update_sub")
@doc_controls.do_not_generate_docs
def update_sub(x, decrement):
    """Update the value of `x` by subtracting `decrement`.

    Args:
        x: A Variable.
        decrement: A tensor of same shape as `x`.

    Returns:
        The variable `x` updated.
    """
    return tf.compat.v1.assign_sub(x, decrement)


@keras_export("keras.backend.moving_average_update")
@doc_controls.do_not_generate_docs
def moving_average_update(x, value, momentum):
    """Compute the exponential moving average of a value.

    The moving average 'x' is updated with 'value' following:

    ```
    x = x * momentum + value * (1 - momentum)
    ```

    For example:

    >>> x = tf.Variable(0.0)
    >>> momentum=0.9
    >>> moving_average_update(x, value = 2.0, momentum=momentum).numpy()
    >>> x.numpy()
    0.2

    The result will be biased towards the initial value of the variable.

    If the variable was initialized to zero, you can divide by
    `1 - momentum ** num_updates` to debias it (Section 3 of
    [Kingma et al., 2015](https://arxiv.org/abs/1412.6980)):

    >>> num_updates = 1.0
    >>> x_zdb = x/(1 - momentum**num_updates)
    >>> x_zdb.numpy()
    2.0

    Args:
        x: A Variable, the moving average.
        value: A tensor with the same shape as `x`, the new value to be
          averaged in.
        momentum: The moving average momentum.

    Returns:
        The updated variable.
    """
    if tf.__internal__.tf2.enabled():
        momentum = tf.cast(momentum, x.dtype)
        value = tf.cast(value, x.dtype)
        return x.assign_sub((x - value) * (1 - momentum))
    else:
        return tf.__internal__.train.assign_moving_average(
            x, value, momentum, zero_debias=True
        )


# LINEAR ALGEBRA


@keras_export("keras.backend.dot")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def dot(x, y):
    """Multiplies 2 tensors (and/or variables) and returns a tensor.

    This operation corresponds to `numpy.dot(a, b, out=None)`.

    Args:
        x: Tensor or variable.
        y: Tensor or variable.

    Returns:
        A tensor, dot product of `x` and `y`.

    Examples:

    If inputs `x` and `y` are 2-D arrays, then it is equivalent to `tf.matmul`.
    >>> x = tf.keras.backend.placeholder(shape=(2, 3))
    >>> y = tf.keras.backend.placeholder(shape=(3, 4))
    >>> xy = tf.keras.backend.dot(x, y)
    >>> xy
    <KerasTensor: shape=(2, 4) dtype=float32 ...>

    >>> x = tf.keras.backend.placeholder(shape=(32, 28, 3))
    >>> y = tf.keras.backend.placeholder(shape=(3, 4))
    >>> xy = tf.keras.backend.dot(x, y)
    >>> xy
    <KerasTensor: shape=(32, 28, 4) dtype=float32 ...>

    If `x` is an N-D array and `y` is an M-D array (where M>=2), it is a sum
    product over the last axis of `x` and the second-to-last axis of `y`.
    >>> x = tf.keras.backend.random_uniform_variable(
    ... shape=(2, 3), low=0., high=1.)
    >>> y = tf.keras.backend.ones((4, 3, 5))
    >>> xy = tf.keras.backend.dot(x, y)
    >>> tf.keras.backend.int_shape(xy)
    (2, 4, 5)
    """
    if ndim(x) is not None and (ndim(x) > 2 or ndim(y) > 2):
        x_shape = []
        for i, s in zip(int_shape(x), tf.unstack(tf.shape(x))):
            if i is not None:
                x_shape.append(i)
            else:
                x_shape.append(s)
        x_shape = tuple(x_shape)
        y_shape = []
        for i, s in zip(int_shape(y), tf.unstack(tf.shape(y))):
            if i is not None:
                y_shape.append(i)
            else:
                y_shape.append(s)
        y_shape = tuple(y_shape)
        y_permute_dim = list(range(ndim(y)))
        y_permute_dim = [y_permute_dim.pop(-2)] + y_permute_dim
        xt = tf.reshape(x, [-1, x_shape[-1]])
        yt = tf.reshape(
            tf.compat.v1.transpose(y, perm=y_permute_dim), [y_shape[-2], -1]
        )
        return tf.reshape(
            tf.matmul(xt, yt), x_shape[:-1] + y_shape[:-2] + y_shape[-1:]
        )
    if is_sparse(x):
        out = tf.sparse.sparse_dense_matmul(x, y)
    else:
        out = tf.matmul(x, y)
    return out


@keras_export("keras.backend.batch_dot")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def batch_dot(x, y, axes=None):
    """Batchwise dot product.

    `batch_dot` is used to compute dot product of `x` and `y` when
    `x` and `y` are data in batch, i.e. in a shape of
    `(batch_size, :)`.
    `batch_dot` results in a tensor or variable with less dimensions
    than the input. If the number of dimensions is reduced to 1,
    we use `expand_dims` to make sure that ndim is at least 2.

    Args:
      x: TF-Keras tensor or variable with `ndim >= 2`.
      y: TF-Keras tensor or variable with `ndim >= 2`.
      axes: Tuple or list of integers with target dimensions, or single integer.
        The sizes of `x.shape[axes[0]]` and `y.shape[axes[1]]` should be equal.

    Returns:
      A tensor with shape equal to the concatenation of `x`'s shape
      (less the dimension that was summed over) and `y`'s shape
      (less the batch dimension and the dimension that was summed over).
      If the final rank is 1, we reshape it to `(batch_size, 1)`.

    Examples:

    >>> x_batch = tf.keras.backend.ones(shape=(32, 20, 1))
    >>> y_batch = tf.keras.backend.ones(shape=(32, 30, 20))
    >>> xy_batch_dot = tf.keras.backend.batch_dot(x_batch, y_batch, axes=(1, 2))
    >>> tf.keras.backend.int_shape(xy_batch_dot)
    (32, 1, 30)

    Shape inference:
      Let `x`'s shape be `(100, 20)` and `y`'s shape be `(100, 30, 20)`.
      If `axes` is (1, 2), to find the output shape of resultant tensor,
          loop through each dimension in `x`'s shape and `y`'s shape:
      * `x.shape[0]` : 100 : append to output shape
      * `x.shape[1]` : 20 : do not append to output shape,
          dimension 1 of `x` has been summed over. (`dot_axes[0]` = 1)
      * `y.shape[0]` : 100 : do not append to output shape,
          always ignore first dimension of `y`
      * `y.shape[1]` : 30 : append to output shape
      * `y.shape[2]` : 20 : do not append to output shape,
          dimension 2 of `y` has been summed over. (`dot_axes[1]` = 2)
      `output_shape` = `(100, 30)`
    """
    x_shape = int_shape(x)
    y_shape = int_shape(y)

    x_ndim = len(x_shape)
    y_ndim = len(y_shape)

    if x_ndim < 2 or y_ndim < 2:
        raise ValueError(
            "Cannot do batch_dot on inputs "
            "with rank < 2. "
            "Received inputs with shapes "
            + str(x_shape)
            + " and "
            + str(y_shape)
            + "."
        )

    x_batch_size = x_shape[0]
    y_batch_size = y_shape[0]

    if x_batch_size is not None and y_batch_size is not None:
        if x_batch_size != y_batch_size:
            raise ValueError(
                "Cannot do batch_dot on inputs "
                "with different batch sizes. "
                "Received inputs with shapes "
                + str(x_shape)
                + " and "
                + str(y_shape)
                + "."
            )
    if isinstance(axes, int):
        axes = [axes, axes]

    if axes is None:
        if y_ndim == 2:
            axes = [x_ndim - 1, y_ndim - 1]
        else:
            axes = [x_ndim - 1, y_ndim - 2]

    if py_any(isinstance(a, (list, tuple)) for a in axes):
        raise ValueError(
            "Multiple target dimensions are not supported. "
            + "Expected: None, int, (int, int), "
            + "Provided: "
            + str(axes)
        )

    # if tuple, convert to list.
    axes = list(axes)

    # convert negative indices.
    if axes[0] < 0:
        axes[0] += x_ndim
    if axes[1] < 0:
        axes[1] += y_ndim

    # sanity checks
    if 0 in axes:
        raise ValueError(
            "Cannot perform batch_dot over axis 0. "
            "If your inputs are not batched, "
            "add a dummy batch dimension to your "
            "inputs using K.expand_dims(x, 0)"
        )
    a0, a1 = axes
    d1 = x_shape[a0]
    d2 = y_shape[a1]

    if d1 is not None and d2 is not None and d1 != d2:
        raise ValueError(
            "Cannot do batch_dot on inputs with shapes "
            + str(x_shape)
            + " and "
            + str(y_shape)
            + " with axes="
            + str(axes)
            + ". x.shape[%d] != y.shape[%d] (%d != %d)."
            % (axes[0], axes[1], d1, d2)
        )

    # backup ndims. Need them later.
    orig_x_ndim = x_ndim
    orig_y_ndim = y_ndim

    # if rank is 2, expand to 3.
    if x_ndim == 2:
        x = tf.expand_dims(x, 1)
        a0 += 1
        x_ndim += 1
    if y_ndim == 2:
        y = tf.expand_dims(y, 2)
        y_ndim += 1

    # bring x's dimension to be reduced to last axis.
    if a0 != x_ndim - 1:
        pattern = list(range(x_ndim))
        for i in range(a0, x_ndim - 1):
            pattern[i] = pattern[i + 1]
        pattern[-1] = a0
        x = tf.compat.v1.transpose(x, pattern)

    # bring y's dimension to be reduced to axis 1.
    if a1 != 1:
        pattern = list(range(y_ndim))
        for i in range(a1, 1, -1):
            pattern[i] = pattern[i - 1]
        pattern[1] = a1
        y = tf.compat.v1.transpose(y, pattern)

    # normalize both inputs to rank 3.
    if x_ndim > 3:
        # squash middle dimensions of x.
        x_shape = shape(x)
        x_mid_dims = x_shape[1:-1]
        x_squashed_shape = tf.stack([x_shape[0], -1, x_shape[-1]])
        x = tf.reshape(x, x_squashed_shape)
        x_squashed = True
    else:
        x_squashed = False

    if y_ndim > 3:
        # squash trailing dimensions of y.
        y_shape = shape(y)
        y_trail_dims = y_shape[2:]
        y_squashed_shape = tf.stack([y_shape[0], y_shape[1], -1])
        y = tf.reshape(y, y_squashed_shape)
        y_squashed = True
    else:
        y_squashed = False

    result = tf.matmul(x, y)

    # if inputs were squashed, we have to reshape the matmul output.
    output_shape = tf.shape(result)
    do_reshape = False

    if x_squashed:
        output_shape = tf.concat(
            [output_shape[:1], x_mid_dims, output_shape[-1:]], 0
        )
        do_reshape = True

    if y_squashed:
        output_shape = tf.concat([output_shape[:-1], y_trail_dims], 0)
        do_reshape = True

    if do_reshape:
        result = tf.reshape(result, output_shape)

    # if the inputs were originally rank 2, we remove the added 1 dim.
    if orig_x_ndim == 2:
        result = tf.squeeze(result, 1)
    elif orig_y_ndim == 2:
        result = tf.squeeze(result, -1)

    return result


@keras_export("keras.backend.transpose")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def transpose(x):
    """Transposes a tensor and returns it.

    Args:
        x: Tensor or variable.

    Returns:
        A tensor.

    Examples:

    >>> var = tf.keras.backend.variable([[1, 2, 3], [4, 5, 6]])
    >>> tf.keras.backend.eval(var)
    array([[1.,  2.,  3.],
           [4.,  5.,  6.]], dtype=float32)
    >>> var_transposed = tf.keras.backend.transpose(var)
    >>> tf.keras.backend.eval(var_transposed)
    array([[1.,  4.],
           [2.,  5.],
           [3.,  6.]], dtype=float32)
    >>> input = tf.keras.backend.placeholder((2, 3))
    >>> input
    <KerasTensor: shape=(2, 3) dtype=float32 ...>
    >>> input_transposed = tf.keras.backend.transpose(input)
    >>> input_transposed
    <KerasTensor: shape=(3, 2) dtype=float32 ...>
    """
    return tf.compat.v1.transpose(x)


@keras_export("keras.backend.gather")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def gather(reference, indices):
    """Retrieves the elements of indices `indices` in the tensor `reference`.

    Args:
        reference: A tensor.
        indices: An integer tensor of indices.

    Returns:
        A tensor of same type as `reference`.

    Examples:

    >>> var = tf.keras.backend.variable([[1, 2, 3], [4, 5, 6]])
    >>> tf.keras.backend.eval(var)
    array([[1., 2., 3.],
           [4., 5., 6.]], dtype=float32)
    >>> var_gathered = tf.keras.backend.gather(var, [0])
    >>> tf.keras.backend.eval(var_gathered)
    array([[1., 2., 3.]], dtype=float32)
    >>> var_gathered = tf.keras.backend.gather(var, [1])
    >>> tf.keras.backend.eval(var_gathered)
    array([[4., 5., 6.]], dtype=float32)
    >>> var_gathered = tf.keras.backend.gather(var, [0,1,0])
    >>> tf.keras.backend.eval(var_gathered)
    array([[1., 2., 3.],
           [4., 5., 6.],
           [1., 2., 3.]], dtype=float32)
    """
    return tf.compat.v1.gather(reference, indices)


# ELEMENT-WISE OPERATIONS


@keras_export("keras.backend.max")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def max(x, axis=None, keepdims=False):
    """Maximum value in a tensor.

    Args:
        x: A tensor or variable.
        axis: An integer, the axis to find maximum values.
        keepdims: A boolean, whether to keep the dimensions or not.
            If `keepdims` is `False`, the rank of the tensor is reduced
            by 1. If `keepdims` is `True`,
            the reduced dimension is retained with length 1.

    Returns:
        A tensor with maximum values of `x`.
    """
    return tf.reduce_max(x, axis, keepdims)


@keras_export("keras.backend.min")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def min(x, axis=None, keepdims=False):
    """Minimum value in a tensor.

    Args:
        x: A tensor or variable.
        axis: An integer, the axis to find minimum values.
        keepdims: A boolean, whether to keep the dimensions or not.
            If `keepdims` is `False`, the rank of the tensor is reduced
            by 1. If `keepdims` is `True`,
            the reduced dimension is retained with length 1.

    Returns:
        A tensor with minimum values of `x`.
    """
    return tf.reduce_min(x, axis, keepdims)


@keras_export("keras.backend.sum")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def sum(x, axis=None, keepdims=False):
    """Sum of the values in a tensor, alongside the specified axis.

    Args:
        x: A tensor or variable.
        axis: An integer, the axis to sum over.
        keepdims: A boolean, whether to keep the dimensions or not.
            If `keepdims` is `False`, the rank of the tensor is reduced
            by 1. If `keepdims` is `True`,
            the reduced dimension is retained with length 1.

    Returns:
        A tensor with sum of `x`.
    """
    return tf.reduce_sum(x, axis, keepdims)


@keras_export("keras.backend.prod")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def prod(x, axis=None, keepdims=False):
    """Multiplies the values in a tensor, alongside the specified axis.

    Args:
        x: A tensor or variable.
        axis: An integer, the axis to compute the product.
        keepdims: A boolean, whether to keep the dimensions or not.
            If `keepdims` is `False`, the rank of the tensor is reduced
            by 1. If `keepdims` is `True`,
            the reduced dimension is retained with length 1.

    Returns:
        A tensor with the product of elements of `x`.
    """
    return tf.reduce_prod(x, axis, keepdims)


@keras_export("keras.backend.cumsum")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def cumsum(x, axis=0):
    """Cumulative sum of the values in a tensor, alongside the specified axis.

    Args:
        x: A tensor or variable.
        axis: An integer, the axis to compute the sum.

    Returns:
        A tensor of the cumulative sum of values of `x` along `axis`.
    """
    return tf.cumsum(x, axis=axis)


@keras_export("keras.backend.cumprod")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def cumprod(x, axis=0):
    """Cumulative product of the values in a tensor alongside `axis`.

    Args:
        x: A tensor or variable.
        axis: An integer, the axis to compute the product.

    Returns:
        A tensor of the cumulative product of values of `x` along `axis`.
    """
    return tf.math.cumprod(x, axis=axis)


@keras_export("keras.backend.var")
@doc_controls.do_not_generate_docs
def var(x, axis=None, keepdims=False):
    """Variance of a tensor, alongside the specified axis.

    Args:
        x: A tensor or variable.
        axis: An integer, the axis to compute the variance.
        keepdims: A boolean, whether to keep the dimensions or not.
            If `keepdims` is `False`, the rank of the tensor is reduced
            by 1. If `keepdims` is `True`,
            the reduced dimension is retained with length 1.

    Returns:
        A tensor with the variance of elements of `x`.
    """
    if x.dtype.base_dtype == tf.bool:
        x = tf.cast(x, floatx())
    return tf.math.reduce_variance(x, axis=axis, keepdims=keepdims)


@keras_export("keras.backend.std")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def std(x, axis=None, keepdims=False):
    """Standard deviation of a tensor, alongside the specified axis.

    It is an alias to `tf.math.reduce_std`.

    Args:
        x: A tensor or variable. It should have numerical dtypes. Boolean type
          inputs will be converted to float.
        axis: An integer, the axis to compute the standard deviation. If `None`
          (the default), reduces all dimensions. Must be in the range
          `[-rank(x), rank(x))`.
        keepdims: A boolean, whether to keep the dimensions or not.
            If `keepdims` is `False`, the rank of the tensor is reduced
            by 1. If `keepdims` is `True`, the reduced dimension is retained
            with length 1.

    Returns:
        A tensor with the standard deviation of elements of `x` with same dtype.
        Boolean type input will be converted to float.
    """
    if x.dtype.base_dtype == tf.bool:
        x = tf.cast(x, floatx())
    return tf.math.reduce_std(x, axis=axis, keepdims=keepdims)


@keras_export("keras.backend.mean")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def mean(x, axis=None, keepdims=False):
    """Mean of a tensor, alongside the specified axis.

    Args:
        x: A tensor or variable.
        axis: A list of integer. Axes to compute the mean.
        keepdims: A boolean, whether to keep the dimensions or not.
            If `keepdims` is `False`, the rank of the tensor is reduced
            by 1 for each entry in `axis`. If `keepdims` is `True`,
            the reduced dimensions are retained with length 1.

    Returns:
        A tensor with the mean of elements of `x`.
    """
    if x.dtype.base_dtype == tf.bool:
        x = tf.cast(x, floatx())
    return tf.reduce_mean(x, axis, keepdims)


@keras_export("keras.backend.any")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def any(x, axis=None, keepdims=False):
    """Bitwise reduction (logical OR).

    Args:
        x: Tensor or variable.
        axis: axis along which to perform the reduction.
        keepdims: whether the drop or broadcast the reduction axes.

    Returns:
        A uint8 tensor (0s and 1s).
    """
    x = tf.cast(x, tf.bool)
    return tf.reduce_any(x, axis, keepdims)


@keras_export("keras.backend.all")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def all(x, axis=None, keepdims=False):
    """Bitwise reduction (logical AND).

    Args:
        x: Tensor or variable.
        axis: axis along which to perform the reduction.
        keepdims: whether the drop or broadcast the reduction axes.

    Returns:
        A uint8 tensor (0s and 1s).
    """
    x = tf.cast(x, tf.bool)
    return tf.reduce_all(x, axis, keepdims)


@keras_export("keras.backend.argmax")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def argmax(x, axis=-1):
    """Returns the index of the maximum value along an axis.

    Args:
        x: Tensor or variable.
        axis: axis along which to perform the reduction.

    Returns:
        A tensor.
    """
    return tf.argmax(x, axis)


@keras_export("keras.backend.argmin")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def argmin(x, axis=-1):
    """Returns the index of the minimum value along an axis.

    Args:
        x: Tensor or variable.
        axis: axis along which to perform the reduction.

    Returns:
        A tensor.
    """
    return tf.argmin(x, axis)


@keras_export("keras.backend.square")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def square(x):
    """Element-wise square.

    Args:
        x: Tensor or variable.

    Returns:
        A tensor.
    """
    return tf.square(x)


@keras_export("keras.backend.abs")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def abs(x):
    """Element-wise absolute value.

    Args:
        x: Tensor or variable.

    Returns:
        A tensor.
    """
    return tf.abs(x)


@keras_export("keras.backend.sqrt")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def sqrt(x):
    """Element-wise square root.

       This function clips negative tensor values to 0 before computing the
       square root.

    Args:
        x: Tensor or variable.

    Returns:
        A tensor.
    """
    zero = _constant_to_tensor(0.0, x.dtype.base_dtype)
    x = tf.maximum(x, zero)
    return tf.sqrt(x)


@keras_export("keras.backend.exp")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def exp(x):
    """Element-wise exponential.

    Args:
        x: Tensor or variable.

    Returns:
        A tensor.
    """
    return tf.exp(x)


@keras_export("keras.backend.log")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def log(x):
    """Element-wise log.

    Args:
        x: Tensor or variable.

    Returns:
        A tensor.
    """
    return tf.math.log(x)


def logsumexp(x, axis=None, keepdims=False):
    """Computes log(sum(exp(elements across dimensions of a tensor))).

    This function is more numerically stable than log(sum(exp(x))).
    It avoids overflows caused by taking the exp of large inputs and
    underflows caused by taking the log of small inputs.

    Args:
        x: A tensor or variable.
        axis: An integer, the axis to reduce over.
        keepdims: A boolean, whether to keep the dimensions or not.
            If `keepdims` is `False`, the rank of the tensor is reduced
            by 1. If `keepdims` is `True`, the reduced dimension is
            retained with length 1.

    Returns:
        The reduced tensor.
    """
    return tf.reduce_logsumexp(x, axis, keepdims)


@keras_export("keras.backend.round")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def round(x):
    """Element-wise rounding to the closest integer.

    In case of tie, the rounding mode used is "half to even".

    Args:
        x: Tensor or variable.

    Returns:
        A tensor.
    """
    return tf.round(x)


@keras_export("keras.backend.sign")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def sign(x):
    """Element-wise sign.

    Args:
        x: Tensor or variable.

    Returns:
        A tensor.
    """
    return tf.sign(x)


@keras_export("keras.backend.pow")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def pow(x, a):
    """Element-wise exponentiation.

    Args:
        x: Tensor or variable.
        a: Python integer.

    Returns:
        A tensor.
    """
    return tf.pow(x, a)


@keras_export("keras.backend.clip")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def clip(x, min_value, max_value):
    """Element-wise value clipping.

    Args:
        x: Tensor or variable.
        min_value: Python float, integer, or tensor.
        max_value: Python float, integer, or tensor.

    Returns:
        A tensor.
    """
    if isinstance(min_value, (int, float)) and isinstance(
        max_value, (int, float)
    ):
        if max_value < min_value:
            max_value = min_value
    if min_value is None:
        min_value = -np.inf
    if max_value is None:
        max_value = np.inf
    return tf.clip_by_value(x, min_value, max_value)


@keras_export("keras.backend.equal")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def equal(x, y):
    """Element-wise equality between two tensors.

    Args:
        x: Tensor or variable.
        y: Tensor or variable.

    Returns:
        A bool tensor.
    """
    return tf.equal(x, y)


@keras_export("keras.backend.not_equal")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def not_equal(x, y):
    """Element-wise inequality between two tensors.

    Args:
        x: Tensor or variable.
        y: Tensor or variable.

    Returns:
        A bool tensor.
    """
    return tf.not_equal(x, y)


@keras_export("keras.backend.greater")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def greater(x, y):
    """Element-wise truth value of (x > y).

    Args:
        x: Tensor or variable.
        y: Tensor or variable.

    Returns:
        A bool tensor.
    """
    return tf.greater(x, y)


@keras_export("keras.backend.greater_equal")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def greater_equal(x, y):
    """Element-wise truth value of (x >= y).

    Args:
        x: Tensor or variable.
        y: Tensor or variable.

    Returns:
        A bool tensor.
    """
    return tf.greater_equal(x, y)


@keras_export("keras.backend.less")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def less(x, y):
    """Element-wise truth value of (x < y).

    Args:
        x: Tensor or variable.
        y: Tensor or variable.

    Returns:
        A bool tensor.
    """
    return tf.less(x, y)


@keras_export("keras.backend.less_equal")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def less_equal(x, y):
    """Element-wise truth value of (x <= y).

    Args:
        x: Tensor or variable.
        y: Tensor or variable.

    Returns:
        A bool tensor.
    """
    return tf.less_equal(x, y)


@keras_export("keras.backend.maximum")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def maximum(x, y):
    """Element-wise maximum of two tensors.

    Args:
        x: Tensor or variable.
        y: Tensor or variable.

    Returns:
        A tensor with the element wise maximum value(s) of `x` and `y`.

    Examples:

    >>> x = tf.Variable([[1, 2], [3, 4]])
    >>> y = tf.Variable([[2, 1], [0, -1]])
    >>> m = tf.keras.backend.maximum(x, y)
    >>> m
    <tf.Tensor: shape=(2, 2), dtype=int32, numpy=
    array([[2, 2],
           [3, 4]], dtype=int32)>
    """
    return tf.maximum(x, y)


@keras_export("keras.backend.minimum")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def minimum(x, y):
    """Element-wise minimum of two tensors.

    Args:
        x: Tensor or variable.
        y: Tensor or variable.

    Returns:
        A tensor.
    """
    return tf.minimum(x, y)


@keras_export("keras.backend.sin")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def sin(x):
    """Computes sin of x element-wise.

    Args:
        x: Tensor or variable.

    Returns:
        A tensor.
    """
    return tf.sin(x)


@keras_export("keras.backend.cos")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def cos(x):
    """Computes cos of x element-wise.

    Args:
        x: Tensor or variable.

    Returns:
        A tensor.
    """
    return tf.cos(x)


def _regular_normalize_batch_in_training(
    x, gamma, beta, reduction_axes, epsilon=1e-3
):
    """Non-fused version of `normalize_batch_in_training`.

    Args:
        x: Input tensor or variable.
        gamma: Tensor by which to scale the input.
        beta: Tensor with which to center the input.
        reduction_axes: iterable of integers,
            axes over which to normalize.
        epsilon: Fuzz factor.

    Returns:
        A tuple length of 3, `(normalized_tensor, mean, variance)`.
    """
    mean, var = tf.compat.v1.nn.moments(x, reduction_axes, None, None, False)
    normed = tf.nn.batch_normalization(x, mean, var, beta, gamma, epsilon)
    return normed, mean, var


def _broadcast_normalize_batch_in_training(
    x, gamma, beta, reduction_axes, epsilon=1e-3
):
    """Non-fused, broadcast version of `normalize_batch_in_training`.

    Args:
        x: Input tensor or variable.
        gamma: Tensor by which to scale the input.
        beta: Tensor with which to center the input.
        reduction_axes: iterable of integers,
            axes over which to normalize.
        epsilon: Fuzz factor.

    Returns:
        A tuple length of 3, `(normalized_tensor, mean, variance)`.
    """
    mean, var = tf.compat.v1.nn.moments(x, reduction_axes, None, None, False)
    target_shape = []
    for axis in range(ndim(x)):
        if axis in reduction_axes:
            target_shape.append(1)
        else:
            target_shape.append(tf.shape(x)[axis])
    target_shape = tf.stack(target_shape)

    broadcast_mean = tf.reshape(mean, target_shape)
    broadcast_var = tf.reshape(var, target_shape)
    if gamma is None:
        broadcast_gamma = None
    else:
        broadcast_gamma = tf.reshape(gamma, target_shape)
    if beta is None:
        broadcast_beta = None
    else:
        broadcast_beta = tf.reshape(beta, target_shape)

    normed = tf.nn.batch_normalization(
        x,
        broadcast_mean,
        broadcast_var,
        broadcast_beta,
        broadcast_gamma,
        epsilon,
    )
    return normed, mean, var


def _fused_normalize_batch_in_training(
    x, gamma, beta, reduction_axes, epsilon=1e-3
):
    """Fused version of `normalize_batch_in_training`.

    Args:
        x: Input tensor or variable.
        gamma: Tensor by which to scale the input.
        beta: Tensor with which to center the input.
        reduction_axes: iterable of integers,
            axes over which to normalize.
        epsilon: Fuzz factor.

    Returns:
        A tuple length of 3, `(normalized_tensor, mean, variance)`.
    """
    if list(reduction_axes) == [0, 1, 2]:
        normalization_axis = 3
        tf_data_format = "NHWC"
    else:
        normalization_axis = 1
        tf_data_format = "NCHW"

    if gamma is None:
        gamma = tf.constant(
            1.0, dtype=x.dtype, shape=[x.shape[normalization_axis]]
        )
    if beta is None:
        beta = tf.constant(
            0.0, dtype=x.dtype, shape=[x.shape[normalization_axis]]
        )

    return tf.compat.v1.nn.fused_batch_norm(
        x, gamma, beta, epsilon=epsilon, data_format=tf_data_format
    )


@keras_export("keras.backend.normalize_batch_in_training")
@doc_controls.do_not_generate_docs
def normalize_batch_in_training(x, gamma, beta, reduction_axes, epsilon=1e-3):
    """Computes mean and std for batch then apply batch_normalization on batch.

    Args:
        x: Input tensor or variable.
        gamma: Tensor by which to scale the input.
        beta: Tensor with which to center the input.
        reduction_axes: iterable of integers,
            axes over which to normalize.
        epsilon: Fuzz factor.

    Returns:
        A tuple length of 3, `(normalized_tensor, mean, variance)`.
    """
    if ndim(x) == 4 and list(reduction_axes) in [[0, 1, 2], [0, 2, 3]]:
        if not _has_nchw_support() and list(reduction_axes) == [0, 2, 3]:
            return _broadcast_normalize_batch_in_training(
                x, gamma, beta, reduction_axes, epsilon=epsilon
            )
        return _fused_normalize_batch_in_training(
            x, gamma, beta, reduction_axes, epsilon=epsilon
        )
    else:
        if sorted(reduction_axes) == list(range(ndim(x)))[:-1]:
            return _regular_normalize_batch_in_training(
                x, gamma, beta, reduction_axes, epsilon=epsilon
            )
        else:
            return _broadcast_normalize_batch_in_training(
                x, gamma, beta, reduction_axes, epsilon=epsilon
            )


@keras_export("keras.backend.batch_normalization")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def batch_normalization(x, mean, var, beta, gamma, axis=-1, epsilon=1e-3):
    """Applies batch normalization on x given mean, var, beta and gamma.

    I.e. returns:
    `output = (x - mean) / (sqrt(var) + epsilon) * gamma + beta`

    Args:
        x: Input tensor or variable.
        mean: Mean of batch.
        var: Variance of batch.
        beta: Tensor with which to center the input.
        gamma: Tensor by which to scale the input.
        axis: Integer, the axis that should be normalized.
            (typically the features axis).
        epsilon: Fuzz factor.

    Returns:
        A tensor.
    """
    if ndim(x) == 4:
        # The CPU implementation of `fused_batch_norm` only supports NHWC
        if axis == 1 or axis == -3:
            tf_data_format = "NCHW"
        elif axis == 3 or axis == -1:
            tf_data_format = "NHWC"
        else:
            tf_data_format = None

        if (
            tf_data_format == "NHWC"
            or tf_data_format == "NCHW"
            and _has_nchw_support()
        ):
            # The mean / var / beta / gamma tensors may be broadcasted
            # so they may have extra axes of size 1, which should be squeezed.
            if ndim(mean) > 1:
                mean = tf.reshape(mean, [-1])
            if ndim(var) > 1:
                var = tf.reshape(var, [-1])
            if beta is None:
                beta = zeros_like(mean)
            elif ndim(beta) > 1:
                beta = tf.reshape(beta, [-1])
            if gamma is None:
                gamma = ones_like(mean)
            elif ndim(gamma) > 1:
                gamma = tf.reshape(gamma, [-1])
        y, _, _ = tf.compat.v1.nn.fused_batch_norm(
            x,
            gamma,
            beta,
            epsilon=epsilon,
            mean=mean,
            variance=var,
            data_format=tf_data_format,
            is_training=False,
        )
        return y
    return tf.nn.batch_normalization(x, mean, var, beta, gamma, epsilon)


# SHAPE OPERATIONS


@keras_export("keras.backend.concatenate")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def concatenate(tensors, axis=-1):
    """Concatenates a list of tensors alongside the specified axis.

    Args:
        tensors: list of tensors to concatenate.
        axis: concatenation axis.

    Returns:
        A tensor.

    Example:

        >>> a = tf.constant([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
        >>> b = tf.constant([[10, 20, 30], [40, 50, 60], [70, 80, 90]])
        >>> tf.keras.backend.concatenate((a, b), axis=-1)
        <tf.Tensor: shape=(3, 6), dtype=int32, numpy=
        array([[ 1,  2,  3, 10, 20, 30],
               [ 4,  5,  6, 40, 50, 60],
               [ 7,  8,  9, 70, 80, 90]], dtype=int32)>

    """
    if axis < 0:
        rank = ndim(tensors[0])
        if rank:
            axis %= rank
        else:
            axis = 0

    if py_all(is_sparse(x) for x in tensors):
        return tf.compat.v1.sparse_concat(axis, tensors)
    elif py_all(isinstance(x, tf.RaggedTensor) for x in tensors):
        return tf.concat(tensors, axis)
    else:
        return tf.concat([to_dense(x) for x in tensors], axis)


@keras_export("keras.backend.reshape")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def reshape(x, shape):
    """Reshapes a tensor to the specified shape.

    Args:
        x: Tensor or variable.
        shape: Target shape tuple.

    Returns:
        A tensor.

    Example:

      >>> a = tf.constant([[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]])
      >>> a
      <tf.Tensor: shape=(4, 3), dtype=int32, numpy=
      array([[ 1,  2,  3],
             [ 4,  5,  6],
             [ 7,  8,  9],
             [10, 11, 12]], dtype=int32)>
      >>> tf.keras.backend.reshape(a, shape=(2, 6))
      <tf.Tensor: shape=(2, 6), dtype=int32, numpy=
      array([[ 1,  2,  3,  4,  5,  6],
             [ 7,  8,  9, 10, 11, 12]], dtype=int32)>

    """
    return tf.reshape(x, shape)


@keras_export("keras.backend.permute_dimensions")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def permute_dimensions(x, pattern):
    """Permutes axes in a tensor.

    Args:
        x: Tensor or variable.
        pattern: A tuple of
            dimension indices, e.g. `(0, 2, 1)`.

    Returns:
        A tensor.

    Example:

      >>> a = tf.constant([[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]])
      >>> a
      <tf.Tensor: shape=(4, 3), dtype=int32, numpy=
      array([[ 1,  2,  3],
             [ 4,  5,  6],
             [ 7,  8,  9],
             [10, 11, 12]], dtype=int32)>
      >>> tf.keras.backend.permute_dimensions(a, pattern=(1, 0))
      <tf.Tensor: shape=(3, 4), dtype=int32, numpy=
      array([[ 1,  4,  7, 10],
             [ 2,  5,  8, 11],
             [ 3,  6,  9, 12]], dtype=int32)>

    """
    return tf.compat.v1.transpose(x, perm=pattern)


@keras_export("keras.backend.resize_images")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def resize_images(
    x, height_factor, width_factor, data_format, interpolation="nearest"
):
    """Resizes the images contained in a 4D tensor.

    Args:
        x: Tensor or variable to resize.
        height_factor: Positive integer.
        width_factor: Positive integer.
        data_format: One of `"channels_first"`, `"channels_last"`.
        interpolation: A string, one of `"area"`, `"bicubic"`, `"bilinear"`,
          `"gaussian"`, `"lanczos3"`, `"lanczos5"`, `"mitchellcubic"`,
          `"nearest"`.

    Returns:
        A tensor.

    Raises:
        ValueError: in case of incorrect value for
          `data_format` or `interpolation`.
    """
    if data_format == "channels_first":
        rows, cols = 2, 3
    elif data_format == "channels_last":
        rows, cols = 1, 2
    else:
        raise ValueError(f"Invalid `data_format` argument: {data_format}")

    new_shape = x.shape[rows : cols + 1]
    if new_shape.is_fully_defined():
        new_shape = tf.constant(new_shape.as_list(), dtype="int32")
    else:
        new_shape = tf.shape(x)[rows : cols + 1]
    new_shape *= tf.constant(
        np.array([height_factor, width_factor], dtype="int32")
    )

    if data_format == "channels_first":
        x = permute_dimensions(x, [0, 2, 3, 1])
    interpolations = {
        "area": tf.image.ResizeMethod.AREA,
        "bicubic": tf.image.ResizeMethod.BICUBIC,
        "bilinear": tf.image.ResizeMethod.BILINEAR,
        "gaussian": tf.image.ResizeMethod.GAUSSIAN,
        "lanczos3": tf.image.ResizeMethod.LANCZOS3,
        "lanczos5": tf.image.ResizeMethod.LANCZOS5,
        "mitchellcubic": tf.image.ResizeMethod.MITCHELLCUBIC,
        "nearest": tf.image.ResizeMethod.NEAREST_NEIGHBOR,
    }
    interploations_list = '"' + '", "'.join(interpolations.keys()) + '"'
    if interpolation in interpolations:
        x = tf.image.resize(x, new_shape, method=interpolations[interpolation])
    else:
        raise ValueError(
            "`interpolation` argument should be one of: "
            f'{interploations_list}. Received: "{interpolation}".'
        )
    if data_format == "channels_first":
        x = permute_dimensions(x, [0, 3, 1, 2])

    return x


@keras_export("keras.backend.resize_volumes")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def resize_volumes(x, depth_factor, height_factor, width_factor, data_format):
    """Resizes the volume contained in a 5D tensor.

    Args:
        x: Tensor or variable to resize.
        depth_factor: Positive integer.
        height_factor: Positive integer.
        width_factor: Positive integer.
        data_format: One of `"channels_first"`, `"channels_last"`.

    Returns:
        A tensor.

    Raises:
        ValueError: if `data_format` is neither
            `channels_last` or `channels_first`.
    """
    if data_format == "channels_first":
        output = repeat_elements(x, depth_factor, axis=2)
        output = repeat_elements(output, height_factor, axis=3)
        output = repeat_elements(output, width_factor, axis=4)
        return output
    elif data_format == "channels_last":
        output = repeat_elements(x, depth_factor, axis=1)
        output = repeat_elements(output, height_factor, axis=2)
        output = repeat_elements(output, width_factor, axis=3)
        return output
    else:
        raise ValueError("Invalid data_format: " + str(data_format))


@keras_export("keras.backend.repeat_elements")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def repeat_elements(x, rep, axis):
    """Repeats the elements of a tensor along an axis, like `np.repeat`.

    If `x` has shape `(s1, s2, s3)` and `axis` is `1`, the output
    will have shape `(s1, s2 * rep, s3)`.

    Args:
        x: Tensor or variable.
        rep: Python integer, number of times to repeat.
        axis: Axis along which to repeat.

    Returns:
        A tensor.

    Example:

        >>> b = tf.constant([1, 2, 3])
        >>> tf.keras.backend.repeat_elements(b, rep=2, axis=0)
        <tf.Tensor: shape=(6,), dtype=int32,
            numpy=array([1, 1, 2, 2, 3, 3], dtype=int32)>

    """
    x_shape = x.shape.as_list()
    # For static axis
    if x_shape[axis] is not None:
        # slices along the repeat axis
        splits = tf.split(value=x, num_or_size_splits=x_shape[axis], axis=axis)
        # repeat each slice the given number of reps
        x_rep = [s for s in splits for _ in range(rep)]
        return concatenate(x_rep, axis)

    # Here we use tf.tile to mimic behavior of np.repeat so that
    # we can handle dynamic shapes (that include None).
    # To do that, we need an auxiliary axis to repeat elements along
    # it and then merge them along the desired axis.

    # Repeating
    auxiliary_axis = axis + 1
    x_shape = tf.shape(x)
    x_rep = tf.expand_dims(x, axis=auxiliary_axis)
    reps = np.ones(len(x.shape) + 1)
    reps[auxiliary_axis] = rep
    x_rep = tf.tile(x_rep, reps)

    # Merging
    reps = np.delete(reps, auxiliary_axis)
    reps[axis] = rep
    reps = tf.constant(reps, dtype="int32")
    x_shape *= reps
    x_rep = tf.reshape(x_rep, x_shape)

    # Fix shape representation
    x_shape = x.shape.as_list()
    x_rep.set_shape(x_shape)
    x_rep._keras_shape = tuple(x_shape)
    return x_rep


@keras_export("keras.backend.repeat")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def repeat(x, n):
    """Repeats a 2D tensor.

    if `x` has shape (samples, dim) and `n` is `2`,
    the output will have shape `(samples, 2, dim)`.

    Args:
        x: Tensor or variable.
        n: Python integer, number of times to repeat.

    Returns:
        A tensor.

    Example:

        >>> b = tf.constant([[1, 2], [3, 4]])
        >>> b
        <tf.Tensor: shape=(2, 2), dtype=int32, numpy=
        array([[1, 2],
               [3, 4]], dtype=int32)>
        >>> tf.keras.backend.repeat(b, n=2)
        <tf.Tensor: shape=(2, 2, 2), dtype=int32, numpy=
        array([[[1, 2],
                [1, 2]],
               [[3, 4],
                [3, 4]]], dtype=int32)>

    """
    assert ndim(x) == 2
    x = tf.expand_dims(x, 1)
    pattern = tf.stack([1, n, 1])
    return tf.tile(x, pattern)


@keras_export("keras.backend.arange")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def arange(start, stop=None, step=1, dtype="int32"):
    """Creates a 1D tensor containing a sequence of integers.

    The function arguments use the same convention as
    Theano's arange: if only one argument is provided,
    it is in fact the "stop" argument and "start" is 0.

    The default type of the returned tensor is `'int32'` to
    match TensorFlow's default.

    Args:
        start: Start value.
        stop: Stop value.
        step: Difference between two successive values.
        dtype: Integer dtype to use.

    Returns:
        An integer tensor.

    Example:

        >>> tf.keras.backend.arange(start=0, stop=10, step=1.5)
        <tf.Tensor: shape=(7,), dtype=float32,
            numpy=array([0. , 1.5, 3. , 4.5, 6. , 7.5, 9. ], dtype=float32)>



    """
    # Match the behavior of numpy and Theano by returning an empty sequence.
    if stop is None and start < 0:
        start = 0
    result = tf.range(start, limit=stop, delta=step, name="arange")
    if dtype != "int32":
        result = cast(result, dtype)
    return result


@keras_export("keras.backend.tile")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def tile(x, n):
    """Creates a tensor by tiling `x` by `n`.

    Args:
        x: A tensor or variable
        n: A list of integer. The length must be the same as the number of
            dimensions in `x`.

    Returns:
        A tiled tensor.
    """
    if isinstance(n, int):
        n = [n]
    return tf.tile(x, n)


@keras_export("keras.backend.flatten")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def flatten(x):
    """Flatten a tensor.

    Args:
        x: A tensor or variable.

    Returns:
        A tensor, reshaped into 1-D

    Example:

        >>> b = tf.constant([[1, 2], [3, 4]])
        >>> b
        <tf.Tensor: shape=(2, 2), dtype=int32, numpy=
        array([[1, 2],
               [3, 4]], dtype=int32)>
        >>> tf.keras.backend.flatten(b)
        <tf.Tensor: shape=(4,), dtype=int32,
            numpy=array([1, 2, 3, 4], dtype=int32)>

    """
    return tf.reshape(x, [-1])


@keras_export("keras.backend.batch_flatten")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def batch_flatten(x):
    """Turn a nD tensor into a 2D tensor with same 0th dimension.

    In other words, it flattens each data samples of a batch.

    Args:
        x: A tensor or variable.

    Returns:
        A tensor.

    Examples:
      Flattening a 3D tensor to 2D by collapsing the last dimension.

    >>> x_batch = tf.keras.backend.ones(shape=(2, 3, 4, 5))
    >>> x_batch_flatten = batch_flatten(x_batch)
    >>> tf.keras.backend.int_shape(x_batch_flatten)
    (2, 60)

    """
    x = tf.reshape(x, tf.stack([-1, prod(shape(x)[1:])]))
    return x


@keras_export("keras.backend.expand_dims")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def expand_dims(x, axis=-1):
    """Adds a 1-sized dimension at index "axis".

    Args:
        x: A tensor or variable.
        axis: Position where to add a new axis.

    Returns:
        A tensor with expanded dimensions.
    """
    return tf.expand_dims(x, axis)


@keras_export("keras.backend.squeeze")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def squeeze(x, axis):
    """Removes a 1-dimension from the tensor at index "axis".

    Args:
        x: A tensor or variable.
        axis: Axis to drop.

    Returns:
        A tensor with the same data as `x` but reduced dimensions.
    """
    return tf.squeeze(x, [axis])


@keras_export("keras.backend.temporal_padding")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def temporal_padding(x, padding=(1, 1)):
    """Pads the middle dimension of a 3D tensor.

    Args:
        x: Tensor or variable.
        padding: Tuple of 2 integers, how many zeros to
            add at the start and end of dim 1.

    Returns:
        A padded 3D tensor.
    """
    assert len(padding) == 2
    pattern = [[0, 0], [padding[0], padding[1]], [0, 0]]
    return tf.compat.v1.pad(x, pattern)


@keras_export("keras.backend.spatial_2d_padding")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def spatial_2d_padding(x, padding=((1, 1), (1, 1)), data_format=None):
    """Pads the 2nd and 3rd dimensions of a 4D tensor.

    Args:
        x: Tensor or variable.
        padding: Tuple of 2 tuples, padding pattern.
        data_format: One of `channels_last` or `channels_first`.

    Returns:
        A padded 4D tensor.

    Raises:
        ValueError: if `data_format` is neither
            `channels_last` or `channels_first`.
    """
    assert len(padding) == 2
    assert len(padding[0]) == 2
    assert len(padding[1]) == 2
    if data_format is None:
        data_format = image_data_format()
    if data_format not in {"channels_first", "channels_last"}:
        raise ValueError("Unknown data_format: " + str(data_format))

    if data_format == "channels_first":
        pattern = [[0, 0], [0, 0], list(padding[0]), list(padding[1])]
    else:
        pattern = [[0, 0], list(padding[0]), list(padding[1]), [0, 0]]
    return tf.compat.v1.pad(x, pattern)


@keras_export("keras.backend.spatial_3d_padding")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def spatial_3d_padding(x, padding=((1, 1), (1, 1), (1, 1)), data_format=None):
    """Pads 5D tensor with zeros along the depth, height, width dimensions.

    Pads these dimensions with respectively
    "padding[0]", "padding[1]" and "padding[2]" zeros left and right.

    For 'channels_last' data_format,
    the 2nd, 3rd and 4th dimension will be padded.
    For 'channels_first' data_format,
    the 3rd, 4th and 5th dimension will be padded.

    Args:
        x: Tensor or variable.
        padding: Tuple of 3 tuples, padding pattern.
        data_format: One of `channels_last` or `channels_first`.

    Returns:
        A padded 5D tensor.

    Raises:
        ValueError: if `data_format` is neither
            `channels_last` or `channels_first`.

    """
    assert len(padding) == 3
    assert len(padding[0]) == 2
    assert len(padding[1]) == 2
    assert len(padding[2]) == 2
    if data_format is None:
        data_format = image_data_format()
    if data_format not in {"channels_first", "channels_last"}:
        raise ValueError("Unknown data_format: " + str(data_format))

    if data_format == "channels_first":
        pattern = [
            [0, 0],
            [0, 0],
            [padding[0][0], padding[0][1]],
            [padding[1][0], padding[1][1]],
            [padding[2][0], padding[2][1]],
        ]
    else:
        pattern = [
            [0, 0],
            [padding[0][0], padding[0][1]],
            [padding[1][0], padding[1][1]],
            [padding[2][0], padding[2][1]],
            [0, 0],
        ]
    return tf.compat.v1.pad(x, pattern)


@keras_export("keras.backend.stack")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def stack(x, axis=0):
    """Stacks a list of rank `R` tensors into a rank `R+1` tensor.

    Args:
        x: List of tensors.
        axis: Axis along which to perform stacking.

    Returns:
        A tensor.

    Example:

        >>> a = tf.constant([[1, 2],[3, 4]])
        >>> b = tf.constant([[10, 20],[30, 40]])
        >>> tf.keras.backend.stack((a, b))
        <tf.Tensor: shape=(2, 2, 2), dtype=int32, numpy=
        array([[[ 1,  2],
                [ 3,  4]],
               [[10, 20],
                [30, 40]]], dtype=int32)>

    """
    return tf.stack(x, axis=axis)


@keras_export("keras.backend.one_hot")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def one_hot(indices, num_classes):
    """Computes the one-hot representation of an integer tensor.

    Args:
        indices: nD integer tensor of shape
            `(batch_size, dim1, dim2, ... dim(n-1))`
        num_classes: Integer, number of classes to consider.

    Returns:
        (n + 1)D one hot representation of the input
        with shape `(batch_size, dim1, dim2, ... dim(n-1), num_classes)`

    Returns:
        The one-hot tensor.
    """
    return tf.one_hot(indices, depth=num_classes, axis=-1)


@keras_export("keras.backend.reverse")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def reverse(x, axes):
    """Reverse a tensor along the specified axes.

    Args:
        x: Tensor to reverse.
        axes: Integer or iterable of integers.
            Axes to reverse.

    Returns:
        A tensor.
    """
    if isinstance(axes, int):
        axes = [axes]
    return tf.reverse(x, axes)


# VALUE MANIPULATION
_VALUE_SET_CODE_STRING = """
  >>> K = tf.keras.backend  # Common keras convention
  >>> v = K.variable(1.)

  >>> # reassign
  >>> K.set_value(v, 2.)
  >>> print(K.get_value(v))
  2.0

  >>> # increment
  >>> K.set_value(v, K.get_value(v) + 1)
  >>> print(K.get_value(v))
  3.0

  Variable semantics in TensorFlow 2 are eager execution friendly. The above
  code is roughly equivalent to:

  >>> v = tf.Variable(1.)

  >>> v.assign(2.)
  >>> print(v.numpy())
  2.0

  >>> v.assign_add(1.)
  >>> print(v.numpy())
  3.0"""[
    3:
]  # Prune first newline and indent to match the docstring template.


@keras_export("keras.backend.get_value")
@doc_controls.do_not_generate_docs
def get_value(x):
    """Returns the value of a variable.

    `backend.get_value` is the complement of `backend.set_value`, and provides
    a generic interface for reading from variables while abstracting away the
    differences between TensorFlow 1.x and 2.x semantics.

    {snippet}

    Args:
        x: input variable.

    Returns:
        A Numpy array.
    """
    if not tf.is_tensor(x):
        return x
    if tf.executing_eagerly() or isinstance(x, tf.__internal__.EagerTensor):
        return x.numpy()
    if not getattr(x, "_in_graph_mode", True):
        # This is a variable which was created in an eager context, but is being
        # evaluated from a Graph.
        with tf.__internal__.eager_context.eager_mode():
            return x.numpy()

    if tf.compat.v1.executing_eagerly_outside_functions():
        # This method of evaluating works inside the TF-Keras FuncGraph.
        with tf.init_scope():
            return x.numpy()

    with x.graph.as_default():
        return x.eval(session=get_session((x,)))


@keras_export("keras.backend.batch_get_value")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def batch_get_value(tensors):
    """Returns the value of more than one tensor variable.

    Args:
        tensors: list of ops to run.

    Returns:
        A list of Numpy arrays.

    Raises:
        RuntimeError: If this method is called inside defun.
    """
    if tf.executing_eagerly():
        return [x.numpy() for x in tensors]
    elif tf.inside_function():
        raise RuntimeError("Cannot get value inside Tensorflow graph function.")
    if tensors:
        return get_session(tensors).run(tensors)
    else:
        return []


@keras_export("keras.backend.set_value")
@doc_controls.do_not_generate_docs
def set_value(x, value):
    """Sets the value of a variable, from a Numpy array.

    `backend.set_value` is the complement of `backend.get_value`, and provides
    a generic interface for assigning to variables while abstracting away the
    differences between TensorFlow 1.x and 2.x semantics.

    {snippet}

    Args:
        x: Variable to set to a new value.
        value: Value to set the tensor to, as a Numpy array
            (of the same shape).
    """
    value = np.asarray(value, dtype=dtype_numpy(x))
    if tf.compat.v1.executing_eagerly_outside_functions():
        _assign_value_to_variable(x, value)
    else:
        with get_graph().as_default():
            tf_dtype = tf.as_dtype(x.dtype.name.split("_")[0])
            if hasattr(x, "_assign_placeholder"):
                assign_placeholder = x._assign_placeholder
                assign_op = x._assign_op
            else:
                # In order to support assigning weights to resizable variables
                # in Keras, we make a placeholder with the correct number of
                # dimensions but with None in each dimension. This way, we can
                # assign weights of any size (as long as they have the correct
                # dimensionality).
                placeholder_shape = tf.TensorShape([None] * value.ndim)
                assign_placeholder = tf.compat.v1.placeholder(
                    tf_dtype, shape=placeholder_shape
                )
                assign_op = x.assign(assign_placeholder)
                x._assign_placeholder = assign_placeholder
                x._assign_op = assign_op
            get_session().run(assign_op, feed_dict={assign_placeholder: value})


@keras_export("keras.backend.batch_set_value")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def batch_set_value(tuples):
    """Sets the values of many tensor variables at once.

    Args:
        tuples: a list of tuples `(tensor, value)`.
            `value` should be a Numpy array.
    """
    if tf.executing_eagerly() or tf.inside_function():
        for x, value in tuples:
            value = np.asarray(value, dtype=dtype_numpy(x))
            _assign_value_to_variable(x, value)
    else:
        with get_graph().as_default():
            if tuples:
                assign_ops = []
                feed_dict = {}
                for x, value in tuples:
                    value = np.asarray(value, dtype=dtype_numpy(x))
                    tf_dtype = tf.as_dtype(x.dtype.name.split("_")[0])
                    if hasattr(x, "_assign_placeholder"):
                        assign_placeholder = x._assign_placeholder
                        assign_op = x._assign_op
                    else:
                        # In order to support assigning weights to resizable
                        # variables in Keras, we make a placeholder with the
                        # correct number of dimensions but with None in each
                        # dimension. This way, we can assign weights of any size
                        # (as long as they have the correct dimensionality).
                        placeholder_shape = tf.TensorShape([None] * value.ndim)
                        assign_placeholder = tf.compat.v1.placeholder(
                            tf_dtype, shape=placeholder_shape
                        )
                        assign_op = x.assign(assign_placeholder)
                        x._assign_placeholder = assign_placeholder
                        x._assign_op = assign_op
                    assign_ops.append(assign_op)
                    feed_dict[assign_placeholder] = value
                get_session().run(assign_ops, feed_dict=feed_dict)


get_value.__doc__ = get_value.__doc__.format(snippet=_VALUE_SET_CODE_STRING)
set_value.__doc__ = set_value.__doc__.format(snippet=_VALUE_SET_CODE_STRING)


def _assign_value_to_variable(variable, value):
    # Helper function to assign value to variable. It handles normal tf.Variable
    # as well as DTensor variable.
    if isinstance(variable, dtensor.DVariable):
        mesh = variable.layout.mesh
        replicate_layout = dtensor.Layout.replicated(
            rank=variable.shape.rank, mesh=mesh
        )
        # TODO(b/262894693): Avoid the broadcast of tensor to all devices.
        d_value = dtensor.copy_to_mesh(value, replicate_layout)
        d_value = dtensor.relayout(d_value, variable.layout)
        variable.assign(d_value)
    else:
        # For the normal tf.Variable assign
        variable.assign(value)


@keras_export("keras.backend.print_tensor")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def print_tensor(x, message="", summarize=3):
    """Prints `message` and the tensor value when evaluated.

    Note that `print_tensor` returns a new tensor identical to `x`
    which should be used in the following code. Otherwise the
    print operation is not taken into account during evaluation.

    Example:

    >>> x = tf.constant([[1.0, 2.0], [3.0, 4.0]])
    >>> tf.keras.backend.print_tensor(x)
    <tf.Tensor: shape=(2, 2), dtype=float32, numpy=
      array([[1., 2.],
             [3., 4.]], dtype=float32)>

    Args:
        x: Tensor to print.
        message: Message to print jointly with the tensor.
        summarize: The first and last `summarize` elements within each dimension
            are recursively printed per Tensor. If None, then the first 3 and
            last 3 elements of each dimension are printed for each tensor. If
            set to -1, it will print all elements of every tensor.

    Returns:
        The same tensor `x`, unchanged.
    """
    if isinstance(x, tf.Tensor) and hasattr(x, "graph"):
        with get_graph().as_default():
            op = tf.print(
                message, x, output_stream=sys.stdout, summarize=summarize
            )
            with tf.control_dependencies([op]):
                return tf.identity(x)
    else:
        tf.print(message, x, output_stream=sys.stdout, summarize=summarize)
        return x


# GRAPH MANIPULATION


class GraphExecutionFunction:
    """Runs a computation graph.

    It's possible to pass arguments to `tf.Session.run()` via `session_kwargs`.
    In particular additional operations via `fetches` argument and additional
    tensor substitutions via `feed_dict` arguments. Note that given
    substitutions are merged with substitutions from `inputs`. Even though
    `feed_dict` is passed once in the constructor (called in `model.compile()`)
    we can modify the values in the dictionary. Through this feed_dict we can
    provide additional substitutions besides TF-Keras inputs.

    Args:
        inputs: Feed placeholders to the computation graph.
        outputs: Output tensors to fetch.
        updates: Additional update ops to be run at function call.
        name: A name to help users identify what this function does.
        session_kwargs: Arguments to `tf.Session.run()`:
                        `fetches`, `feed_dict`, `options`, `run_metadata`.
    """

    def __init__(
        self, inputs, outputs, updates=None, name=None, **session_kwargs
    ):
        updates = updates or []
        if not isinstance(updates, (list, tuple)):
            raise TypeError(
                "`updates` in a TF-Keras backend function "
                "should be a list or tuple."
            )

        self.inputs = tf.nest.flatten(
            tf_utils.convert_variables_to_tensors(inputs),
            expand_composites=True,
        )
        self._outputs_structure = tf_utils.convert_variables_to_tensors(outputs)
        self.outputs = tf.nest.flatten(
            self._outputs_structure, expand_composites=True
        )
        # TODO(b/127668432): Consider using autograph to generate these
        # dependencies in call.
        # Index 0 = total loss or model output for `predict`.
        with tf.control_dependencies([self.outputs[0]]):
            updates_ops = []
            for update in updates:
                if isinstance(update, tuple):
                    p, new_p = update
                    updates_ops.append(tf.compat.v1.assign(p, new_p))
                else:
                    # assumed already an op
                    updates_ops.append(update)
            self.updates_op = tf.group(*updates_ops)
        self.name = name
        # additional tensor substitutions
        self.feed_dict = session_kwargs.pop("feed_dict", None)
        # additional operations
        self.fetches = session_kwargs.pop("fetches", [])
        if not isinstance(self.fetches, list):
            self.fetches = [self.fetches]
        self.run_options = session_kwargs.pop("options", None)
        self.run_metadata = session_kwargs.pop("run_metadata", None)
        # The main use case of `fetches` being passed to a model is the ability
        # to run custom updates
        # This requires us to wrap fetches in `identity` ops.
        self.fetches = [tf.identity(x) for x in self.fetches]
        self.session_kwargs = session_kwargs
        # This mapping keeps track of the function that should receive the
        # output from a fetch in `fetches`: { fetch: function(fetch_output) }
        # A Callback can use this to register a function with access to the
        # output values for a fetch it added.
        self.fetch_callbacks = {}

        if session_kwargs:
            raise ValueError(
                "Some keys in session_kwargs are not supported at this time: %s"
                % (session_kwargs.keys(),)
            )

        self._callable_fn = None
        self._feed_arrays = None
        self._feed_symbols = None
        self._symbol_vals = None
        self._fetches = None
        self._session = None

    def _make_callable(self, feed_arrays, feed_symbols, symbol_vals, session):
        """Generates a callable that runs the graph.

        Args:
          feed_arrays: List of input tensors to be fed Numpy arrays at runtime.
          feed_symbols: List of input tensors to be fed symbolic tensors at
            runtime.
          symbol_vals: List of symbolic tensors to be fed to `feed_symbols`.
          session: Session to use to generate the callable.

        Returns:
          Function that runs the graph according to the above options.
        """
        # Prepare callable options.
        callable_opts = config_pb2.CallableOptions()
        # Handle external-data feed.
        for x in feed_arrays:
            callable_opts.feed.append(x.name)
        if self.feed_dict:
            for key in sorted(self.feed_dict.keys()):
                callable_opts.feed.append(key.name)
        # Handle symbolic feed.
        for x, y in zip(feed_symbols, symbol_vals):
            connection = callable_opts.tensor_connection.add()
            if x.dtype != y.dtype:
                y = tf.cast(y, dtype=x.dtype)
            from_tensor = _as_graph_element(y)
            if from_tensor is None:
                from_tensor = y
            connection.from_tensor = from_tensor.name  # Data tensor
            connection.to_tensor = x.name  # Placeholder
        # Handle fetches.
        for x in self.outputs + self.fetches:
            callable_opts.fetch.append(x.name)
        # Handle updates.
        callable_opts.target.append(self.updates_op.name)
        # Handle run_options.
        if self.run_options:
            callable_opts.run_options.CopyFrom(self.run_options)
        # Create callable.
        callable_fn = session._make_callable_from_options(callable_opts)
        # Cache parameters corresponding to the generated callable, so that
        # we can detect future mismatches and refresh the callable.
        self._callable_fn = callable_fn
        self._feed_arrays = feed_arrays
        self._feed_symbols = feed_symbols
        self._symbol_vals = symbol_vals
        self._fetches = list(self.fetches)
        self._session = session

    def _call_fetch_callbacks(self, fetches_output):
        for fetch, output in zip(self._fetches, fetches_output):
            if fetch in self.fetch_callbacks:
                self.fetch_callbacks[fetch](output)

    def _eval_if_composite(self, tensor):
        """Helper method which evaluates any CompositeTensors passed to it."""
        # We need to evaluate any composite tensor objects that have been
        # reconstructed in 'pack_sequence_as', since otherwise they'll be output
        # as actual CompositeTensor objects instead of the value(s) contained in
        # the CompositeTensors. E.g., if output_structure contains a
        # SparseTensor, then this ensures that we return its value as a
        # SparseTensorValue rather than a SparseTensor.

        if tf_utils.is_extension_type(tensor):
            return self._session.run(tensor)
        else:
            return tensor

    def __call__(self, inputs):
        inputs = tf.nest.flatten(
            tf_utils.convert_variables_to_tensors(inputs),
            expand_composites=True,
        )

        session = get_session(inputs)
        feed_arrays = []
        array_vals = []
        feed_symbols = []
        symbol_vals = []
        for tensor, value in zip(self.inputs, inputs):
            if value is None:
                continue

            if tf.is_tensor(value):
                # Case: feeding symbolic tensor.
                feed_symbols.append(tensor)
                symbol_vals.append(value)
            else:
                # Case: feeding Numpy array.
                feed_arrays.append(tensor)
                # We need to do array conversion and type casting at this level,
                # since `callable_fn` only supports exact matches.
                tensor_type = tf.as_dtype(tensor.dtype)
                array_vals.append(
                    np.asarray(value, dtype=tensor_type.as_numpy_dtype)
                )

        if self.feed_dict:
            for key in sorted(self.feed_dict.keys()):
                array_vals.append(
                    np.asarray(
                        self.feed_dict[key], dtype=key.dtype.as_numpy_dtype
                    )
                )

        # Refresh callable if anything has changed.
        if (
            self._callable_fn is None
            or feed_arrays != self._feed_arrays
            or symbol_vals != self._symbol_vals
            or feed_symbols != self._feed_symbols
            or self.fetches != self._fetches
            or session != self._session
        ):
            self._make_callable(feed_arrays, feed_symbols, symbol_vals, session)

        fetched = self._callable_fn(*array_vals, run_metadata=self.run_metadata)
        self._call_fetch_callbacks(fetched[-len(self._fetches) :])
        output_structure = tf.nest.pack_sequence_as(
            self._outputs_structure,
            fetched[: len(self.outputs)],
            expand_composites=True,
        )
        # We need to evaluate any composite tensor objects that have been
        # reconstructed in 'pack_sequence_as', since otherwise they'll be output
        # as actual CompositeTensor objects instead of the value(s) contained in
        # the CompositeTensors. E.g., if output_structure contains a
        # SparseTensor, then this ensures that we return its value as a
        # SparseTensorValue rather than a SparseTensor.
        return tf.nest.map_structure(self._eval_if_composite, output_structure)


@keras_export("keras.backend.function")
@doc_controls.do_not_generate_docs
def function(inputs, outputs, updates=None, name=None, **kwargs):
    """Instantiates a TF-Keras function.

    Args:
        inputs: List of placeholder tensors.
        outputs: List of output tensors.
        updates: List of update ops.
        name: String, name of function.
        **kwargs: Passed to `tf.Session.run`.

    Returns:
        Output values as Numpy arrays.

    Raises:
        ValueError: if invalid kwargs are passed in or if in eager execution.
    """
    if tf.compat.v1.executing_eagerly_outside_functions():
        if kwargs:
            raise ValueError(
                "Session keyword arguments are not supported during "
                "eager execution. You passed: %s" % (kwargs,)
            )
        if updates:
            raise ValueError(
                "`updates` argument is not supported during "
                "eager execution. You passed: %s" % (updates,)
            )
        from tf_keras import models

        model = models.Model(inputs=inputs, outputs=outputs)

        wrap_outputs = isinstance(outputs, list) and len(outputs) == 1

        def func(model_inputs):
            outs = model(model_inputs)
            if wrap_outputs:
                outs = [outs]
            return tf_utils.sync_to_numpy_or_python_type(outs)

        return func

    if kwargs:
        for key in kwargs:
            if key not in tf_inspect.getfullargspec(tf.compat.v1.Session.run)[
                0
            ] and key not in ["inputs", "outputs", "updates", "name"]:
                msg = (
                    'Invalid argument "%s" passed to K.function with '
                    "TensorFlow backend" % key
                )
                raise ValueError(msg)
    return GraphExecutionFunction(
        inputs, outputs, updates=updates, name=name, **kwargs
    )


@keras_export("keras.backend.gradients")
@doc_controls.do_not_generate_docs
def gradients(loss, variables):
    """Returns the gradients of `loss` w.r.t. `variables`.

    Args:
        loss: Scalar tensor to minimize.
        variables: List of variables.

    Returns:
        A gradients tensor.
    """
    return tf.compat.v1.gradients(
        loss, variables, colocate_gradients_with_ops=True
    )


@keras_export("keras.backend.stop_gradient")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def stop_gradient(variables):
    """Returns `variables` but with zero gradient w.r.t. every other variable.

    Args:
        variables: Tensor or list of tensors to consider constant with respect
          to any other variable.


    Returns:
        A single tensor or a list of tensors (depending on the passed argument)
        that has no gradient with respect to any other variable.
    """
    if isinstance(variables, (list, tuple)):
        return map(tf.stop_gradient, variables)
    return tf.stop_gradient(variables)


# CONTROL FLOW


@keras_export("keras.backend.rnn")
@tf.__internal__.dispatch.add_dispatch_support
def rnn(
    step_function,
    inputs,
    initial_states,
    go_backwards=False,
    mask=None,
    constants=None,
    unroll=False,
    input_length=None,
    time_major=False,
    zero_output_for_mask=False,
    return_all_outputs=True,
):
    """Iterates over the time dimension of a tensor.

    Args:
        step_function: RNN step function.
            Args;
                input; Tensor with shape `(samples, ...)` (no time dimension),
                    representing input for the batch of samples at a certain
                    time step.
                states; List of tensors.
            Returns;
                output; Tensor with shape `(samples, output_dim)`
                    (no time dimension).
                new_states; List of tensors, same length and shapes
                    as 'states'. The first state in the list must be the
                    output tensor at the previous timestep.
        inputs: Tensor of temporal data of shape `(samples, time, ...)`
            (at least 3D), or nested tensors, and each of which has shape
            `(samples, time, ...)`.
        initial_states: Tensor with shape `(samples, state_size)`
            (no time dimension), containing the initial values for the states
            used in the step function. In the case that state_size is in a
            nested shape, the shape of initial_states will also follow the
            nested structure.
        go_backwards: Boolean. If True, do the iteration over the time
            dimension in reverse order and return the reversed sequence.
        mask: Binary tensor with shape `(samples, time, 1)`,
            with a zero for every element that is masked.
        constants: List of constant values passed at each step.
        unroll: Whether to unroll the RNN or to use a symbolic `while_loop`.
        input_length: An integer or a 1-D Tensor, depending on whether
            the time dimension is fixed-length or not. In case of variable
            length input, it is used for masking in case there's no mask
            specified.
        time_major: Boolean. If true, the inputs and outputs will be in shape
            `(timesteps, batch, ...)`, whereas in the False case, it will be
            `(batch, timesteps, ...)`. Using `time_major = True` is a bit more
            efficient because it avoids transposes at the beginning and end of
            the RNN calculation. However, most TensorFlow data is batch-major,
            so by default this function accepts input and emits output in
            batch-major form.
        zero_output_for_mask: Boolean. If True, the output for masked timestep
            will be zeros, whereas in the False case, output from previous
            timestep is returned.
        return_all_outputs: Boolean. If True, return the recurrent outputs for
            all timesteps in the sequence. If False, only return the output for
            the last timestep (which consumes less memory).

    Returns:
        A tuple, `(last_output, outputs, new_states)`.
            last_output: the latest output of the rnn, of shape `(samples, ...)`
            outputs:
                - If `return_all_outputs=True`: a tensor with shape
                  `(samples, time, ...)` where each entry `outputs[s, t]` is the
                  output of the step function at time `t` for sample `s`
                - Else, a tensor equal to `last_output` with shape
                  `(samples, 1, ...)`
            new_states: list of tensors, latest states returned by
                the step function, of shape `(samples, ...)`.

    Raises:
        ValueError: if input dimension is less than 3.
        ValueError: if `unroll` is `True` but input timestep is not a fixed
            number.
        ValueError: if `mask` is provided (not `None`) but states is not
            provided (`len(states)` == 0).
    """
    if not tf.__internal__.tf2.enabled():
        return_all_outputs = True  # Not supported in TF1.

    def swap_batch_timestep(input_t):
        # Swap the batch and timestep dim for the incoming tensor.
        axes = list(range(len(input_t.shape)))
        axes[0], axes[1] = 1, 0
        return tf.compat.v1.transpose(input_t, axes)

    if not time_major:
        inputs = tf.nest.map_structure(swap_batch_timestep, inputs)

    flatted_inputs = tf.nest.flatten(inputs)
    time_steps = flatted_inputs[0].shape[0]
    batch = flatted_inputs[0].shape[1]
    time_steps_t = tf.shape(flatted_inputs[0])[0]

    for input_ in flatted_inputs:
        input_.shape.with_rank_at_least(3)

    if mask is not None:
        if mask.dtype != tf.bool:
            mask = tf.cast(mask, tf.bool)
        if len(mask.shape) == 2:
            mask = expand_dims(mask)
        if not time_major:
            mask = swap_batch_timestep(mask)

    if constants is None:
        constants = []

    # tf.where needs its condition tensor to be the same shape as its two
    # result tensors, but in our case the condition (mask) tensor is
    # (nsamples, 1), and inputs are (nsamples, ndimensions) or even more.
    # So we need to broadcast the mask to match the shape of inputs.
    # That's what the tile call does, it just repeats the mask along its
    # second dimension n times.
    def _expand_mask(mask_t, input_t, fixed_dim=1):
        if tf.nest.is_nested(mask_t):
            raise ValueError(
                f"mask_t is expected to be tensor, but got {mask_t}"
            )
        if tf.nest.is_nested(input_t):
            raise ValueError(
                f"input_t is expected to be tensor, but got {input_t}"
            )
        rank_diff = len(input_t.shape) - len(mask_t.shape)
        for _ in range(rank_diff):
            mask_t = tf.expand_dims(mask_t, -1)
        multiples = [1] * fixed_dim + input_t.shape.as_list()[fixed_dim:]
        return tf.tile(mask_t, multiples)

    if unroll:
        if not time_steps:
            raise ValueError("Unrolling requires a fixed number of timesteps.")
        states = tuple(initial_states)
        successive_states = []
        successive_outputs = []

        # Process the input tensors. The input tensor need to be split on the
        # time_step dim, and reverse if go_backwards is True. In the case of
        # nested input, the input is flattened and then transformed
        # individually.  The result of this will be a tuple of lists, each of
        # the item in tuple is list of the tensor with shape (batch, feature)
        def _process_single_input_t(input_t):
            input_t = tf.unstack(input_t)  # unstack for time_step dim
            if go_backwards:
                input_t.reverse()
            return input_t

        if tf.nest.is_nested(inputs):
            processed_input = tf.nest.map_structure(
                _process_single_input_t, inputs
            )
        else:
            processed_input = (_process_single_input_t(inputs),)

        def _get_input_tensor(time):
            inp = [t_[time] for t_ in processed_input]
            return tf.nest.pack_sequence_as(inputs, inp)

        if mask is not None:
            mask_list = tf.unstack(mask)
            if go_backwards:
                mask_list.reverse()

            for i in range(time_steps):
                inp = _get_input_tensor(i)
                mask_t = mask_list[i]
                output, new_states = step_function(
                    inp, tuple(states) + tuple(constants)
                )
                tiled_mask_t = _expand_mask(mask_t, output)

                if not successive_outputs:
                    prev_output = zeros_like(output)
                else:
                    prev_output = successive_outputs[-1]

                output = tf.where(tiled_mask_t, output, prev_output)

                flat_states = tf.nest.flatten(states)
                flat_new_states = tf.nest.flatten(new_states)
                tiled_mask_t = tuple(
                    _expand_mask(mask_t, s) for s in flat_states
                )
                flat_final_states = tuple(
                    tf.where(m, s, ps)
                    for m, s, ps in zip(
                        tiled_mask_t, flat_new_states, flat_states
                    )
                )
                states = tf.nest.pack_sequence_as(states, flat_final_states)

                if return_all_outputs:
                    successive_outputs.append(output)
                    successive_states.append(states)
                else:
                    successive_outputs = [output]
                    successive_states = [states]
            last_output = successive_outputs[-1]
            new_states = successive_states[-1]
            outputs = tf.stack(successive_outputs)

            if zero_output_for_mask:
                last_output = tf.where(
                    _expand_mask(mask_list[-1], last_output),
                    last_output,
                    zeros_like(last_output),
                )
                outputs = tf.where(
                    _expand_mask(mask, outputs, fixed_dim=2),
                    outputs,
                    zeros_like(outputs),
                )

        else:  # mask is None
            for i in range(time_steps):
                inp = _get_input_tensor(i)
                output, states = step_function(
                    inp, tuple(states) + tuple(constants)
                )
                if return_all_outputs:
                    successive_outputs.append(output)
                    successive_states.append(states)
                else:
                    successive_outputs = [output]
                    successive_states = [states]
            last_output = successive_outputs[-1]
            new_states = successive_states[-1]
            outputs = tf.stack(successive_outputs)

    else:  # Unroll == False
        states = tuple(initial_states)

        # Create input tensor array, if the inputs is nested tensors, then it
        # will be flattened first, and tensor array will be created one per
        # flattened tensor.
        input_ta = tuple(
            tf.TensorArray(
                dtype=inp.dtype,
                size=time_steps_t,
                tensor_array_name=f"input_ta_{i}",
            )
            for i, inp in enumerate(flatted_inputs)
        )
        input_ta = tuple(
            ta.unstack(input_)
            if not go_backwards
            else ta.unstack(reverse(input_, 0))
            for ta, input_ in zip(input_ta, flatted_inputs)
        )

        # Get the time(0) input and compute the output for that, the output will
        # be used to determine the dtype of output tensor array. Don't read from
        # input_ta due to TensorArray clear_after_read default to True.
        input_time_zero = tf.nest.pack_sequence_as(
            inputs, [inp[0] for inp in flatted_inputs]
        )
        # output_time_zero is used to determine the cell output shape and its
        # dtype.  the value is discarded.
        output_time_zero, _ = step_function(
            input_time_zero, tuple(initial_states) + tuple(constants)
        )

        output_ta_size = time_steps_t if return_all_outputs else 1
        output_ta = tuple(
            tf.TensorArray(
                dtype=out.dtype,
                size=output_ta_size,
                element_shape=out.shape,
                tensor_array_name=f"output_ta_{i}",
            )
            for i, out in enumerate(tf.nest.flatten(output_time_zero))
        )

        time = tf.constant(0, dtype="int32", name="time")

        # We only specify the 'maximum_iterations' when building for XLA since
        # that causes slowdowns on GPU in TF.
        if (
            not tf.executing_eagerly()
            and control_flow_util.GraphOrParentsInXlaContext(
                tf.compat.v1.get_default_graph()
            )
        ):
            if input_length is None:
                max_iterations = time_steps_t
            else:
                max_iterations = tf.reduce_max(input_length)
        else:
            max_iterations = None

        while_loop_kwargs = {
            "cond": lambda time, *_: time < time_steps_t,
            "maximum_iterations": max_iterations,
            "parallel_iterations": 32,
            "swap_memory": True,
        }
        if mask is not None:
            if go_backwards:
                mask = reverse(mask, 0)

            mask_ta = tf.TensorArray(
                dtype=tf.bool, size=time_steps_t, tensor_array_name="mask_ta"
            )
            mask_ta = mask_ta.unstack(mask)

            def masking_fn(time):
                return mask_ta.read(time)

            def compute_masked_output(mask_t, flat_out, flat_mask):
                tiled_mask_t = tuple(
                    _expand_mask(mask_t, o, fixed_dim=len(mask_t.shape))
                    for o in flat_out
                )
                return tuple(
                    tf.where(m, o, fm)
                    for m, o, fm in zip(tiled_mask_t, flat_out, flat_mask)
                )

        elif isinstance(input_length, tf.Tensor):
            if go_backwards:
                max_len = tf.reduce_max(input_length, axis=0)
                rev_input_length = tf.subtract(max_len - 1, input_length)

                def masking_fn(time):
                    return tf.less(rev_input_length, time)

            else:

                def masking_fn(time):
                    return tf.greater(input_length, time)

            def compute_masked_output(mask_t, flat_out, flat_mask):
                return tuple(
                    tf.compat.v1.where(mask_t, o, zo)
                    for (o, zo) in zip(flat_out, flat_mask)
                )

        else:
            masking_fn = None

        if masking_fn is not None:
            # Mask for the T output will be base on the output of T - 1. In the
            # case T = 0, a zero filled tensor will be used.
            flat_zero_output = tuple(
                tf.zeros_like(o) for o in tf.nest.flatten(output_time_zero)
            )

            def _step(time, output_ta_t, prev_output, *states):
                """RNN step function.

                Args:
                    time: Current timestep value.
                    output_ta_t: TensorArray.
                    prev_output: tuple of outputs from time - 1.
                    *states: List of states.

                Returns:
                    Tuple: `(time + 1, output_ta_t, output) + tuple(new_states)`
                """
                current_input = tuple(ta.read(time) for ta in input_ta)
                # maybe set shape.
                current_input = tf.nest.pack_sequence_as(inputs, current_input)
                mask_t = masking_fn(time)
                output, new_states = step_function(
                    current_input, tuple(states) + tuple(constants)
                )
                # mask output
                flat_output = tf.nest.flatten(output)
                flat_mask_output = (
                    flat_zero_output
                    if zero_output_for_mask
                    else tf.nest.flatten(prev_output)
                )
                flat_new_output = compute_masked_output(
                    mask_t, flat_output, flat_mask_output
                )

                # mask states
                flat_state = tf.nest.flatten(states)
                flat_new_state = tf.nest.flatten(new_states)
                for state, new_state in zip(flat_state, flat_new_state):
                    if isinstance(new_state, tf.Tensor):
                        new_state.set_shape(state.shape)
                flat_final_state = compute_masked_output(
                    mask_t, flat_new_state, flat_state
                )
                new_states = tf.nest.pack_sequence_as(
                    new_states, flat_final_state
                )

                ta_index_to_write = time if return_all_outputs else 0
                output_ta_t = tuple(
                    ta.write(ta_index_to_write, out)
                    for ta, out in zip(output_ta_t, flat_new_output)
                )

                return (time + 1, output_ta_t, tuple(flat_new_output)) + tuple(
                    new_states
                )

            final_outputs = tf.compat.v1.while_loop(
                body=_step,
                loop_vars=(time, output_ta, flat_zero_output) + states,
                **while_loop_kwargs,
            )
            # Skip final_outputs[2] which is the output for final timestep.
            new_states = final_outputs[3:]
        else:

            def _step(time, output_ta_t, *states):
                """RNN step function.

                Args:
                    time: Current timestep value.
                    output_ta_t: TensorArray.
                    *states: List of states.

                Returns:
                    Tuple: `(time + 1,output_ta_t) + tuple(new_states)`
                """
                current_input = tuple(ta.read(time) for ta in input_ta)
                current_input = tf.nest.pack_sequence_as(inputs, current_input)
                output, new_states = step_function(
                    current_input, tuple(states) + tuple(constants)
                )
                flat_state = tf.nest.flatten(states)
                flat_new_state = tf.nest.flatten(new_states)
                for state, new_state in zip(flat_state, flat_new_state):
                    if isinstance(new_state, tf.Tensor):
                        new_state.set_shape(state.shape)

                flat_output = tf.nest.flatten(output)
                ta_index_to_write = time if return_all_outputs else 0
                output_ta_t = tuple(
                    ta.write(ta_index_to_write, out)
                    for ta, out in zip(output_ta_t, flat_output)
                )

                new_states = tf.nest.pack_sequence_as(
                    initial_states, flat_new_state
                )
                return (time + 1, output_ta_t) + tuple(new_states)

            final_outputs = tf.compat.v1.while_loop(
                body=_step,
                loop_vars=(time, output_ta) + states,
                **while_loop_kwargs,
            )
            new_states = final_outputs[2:]

        output_ta = final_outputs[1]

        outputs = tuple(o.stack() for o in output_ta)
        last_output = tuple(o[-1] for o in outputs)

        outputs = tf.nest.pack_sequence_as(output_time_zero, outputs)
        last_output = tf.nest.pack_sequence_as(output_time_zero, last_output)

    # static shape inference
    def set_shape(output_):
        if isinstance(output_, tf.Tensor):
            shape = output_.shape.as_list()
            if return_all_outputs:
                shape[0] = time_steps
            else:
                shape[0] = 1
            shape[1] = batch
            output_.set_shape(shape)
        return output_

    outputs = tf.nest.map_structure(set_shape, outputs)

    if not time_major:
        outputs = tf.nest.map_structure(swap_batch_timestep, outputs)

    return last_output, outputs, new_states


@keras_export("keras.backend.switch")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def switch(condition, then_expression, else_expression):
    """Switches between two operations depending on a scalar value.

    Note that both `then_expression` and `else_expression`
    should be symbolic tensors of the *same shape*.

    Args:
        condition: tensor (`int` or `bool`).
        then_expression: either a tensor, or a callable that returns a tensor.
        else_expression: either a tensor, or a callable that returns a tensor.

    Returns:
        The selected tensor.

    Raises:
        ValueError: If rank of `condition` is greater than rank of expressions.
    """
    if condition.dtype != tf.bool:
        condition = tf.cast(condition, "bool")
    cond_ndim = ndim(condition)
    if not cond_ndim:
        if not callable(then_expression):

            def then_expression_fn():
                return then_expression

        else:
            then_expression_fn = then_expression
        if not callable(else_expression):

            def else_expression_fn():
                return else_expression

        else:
            else_expression_fn = else_expression
        x = tf.compat.v1.cond(condition, then_expression_fn, else_expression_fn)
    else:
        # tf.where needs its condition tensor
        # to be the same shape as its two
        # result tensors
        if callable(then_expression):
            then_expression = then_expression()
        if callable(else_expression):
            else_expression = else_expression()
        expr_ndim = ndim(then_expression)
        if cond_ndim > expr_ndim:
            raise ValueError(
                "Rank of `condition` should be less than or"
                " equal to rank of `then_expression` and "
                "`else_expression`. ndim(condition)="
                + str(cond_ndim)
                + ", ndim(then_expression)="
                + str(expr_ndim)
            )
        if cond_ndim > 1:
            ndim_diff = expr_ndim - cond_ndim
            cond_shape = tf.concat(
                [tf.shape(condition), [1] * ndim_diff], axis=0
            )
            condition = tf.reshape(condition, cond_shape)
            expr_shape = tf.shape(then_expression)
            shape_diff = expr_shape - cond_shape
            tile_shape = tf.where(
                shape_diff > 0, expr_shape, tf.ones_like(expr_shape)
            )
            condition = tf.tile(condition, tile_shape)
        x = tf.where(condition, then_expression, else_expression)
    return x


@keras_export("keras.backend.in_train_phase")
@doc_controls.do_not_generate_docs
def in_train_phase(x, alt, training=None):
    """Selects `x` in train phase, and `alt` otherwise.

    Note that `alt` should have the *same shape* as `x`.

    Args:
        x: What to return in train phase
            (tensor or callable that returns a tensor).
        alt: What to return otherwise
            (tensor or callable that returns a tensor).
        training: Optional scalar tensor
            (or Python boolean, or Python integer)
            specifying the learning phase.

    Returns:
        Either `x` or `alt` based on the `training` flag.
        the `training` flag defaults to `K.learning_phase()`.
    """
    from tf_keras.engine import (
        base_layer_utils,
    )

    if training is None:
        training = base_layer_utils.call_context().training

    if training is None:
        training = learning_phase()

    # TODO(b/138862903): Handle the case when training is tensor.
    if not tf.is_tensor(training):
        if training == 1 or training is True:
            if callable(x):
                return x()
            else:
                return x

        elif training == 0 or training is False:
            if callable(alt):
                return alt()
            else:
                return alt

    # else: assume learning phase is a placeholder tensor.
    x = switch(training, x, alt)
    return x


@keras_export("keras.backend.in_test_phase")
@doc_controls.do_not_generate_docs
def in_test_phase(x, alt, training=None):
    """Selects `x` in test phase, and `alt` otherwise.

    Note that `alt` should have the *same shape* as `x`.

    Args:
        x: What to return in test phase
            (tensor or callable that returns a tensor).
        alt: What to return otherwise
            (tensor or callable that returns a tensor).
        training: Optional scalar tensor
            (or Python boolean, or Python integer)
            specifying the learning phase.

    Returns:
        Either `x` or `alt` based on `K.learning_phase`.
    """
    return in_train_phase(alt, x, training=training)


# NN OPERATIONS


@keras_export("keras.backend.relu")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def relu(x, alpha=0.0, max_value=None, threshold=0.0):
    """Rectified linear unit.

    With default values, it returns element-wise `max(x, 0)`.

    Otherwise, it follows:
    `f(x) = max_value` for `x >= max_value`,
    `f(x) = x` for `threshold <= x < max_value`,
    `f(x) = alpha * (x - threshold)` otherwise.

    Args:
        x: A tensor or variable.
        alpha: A scalar, slope of negative section (default=`0.`).
        max_value: float. Saturation threshold.
        threshold: float. Threshold value for thresholded activation.

    Returns:
        A tensor.
    """
    # While x can be a tensor or variable, we also see cases where
    # numpy arrays, lists, tuples are passed as well.
    # lists, tuples do not have 'dtype' attribute.
    dtype = getattr(x, "dtype", floatx())
    if alpha != 0.0:
        if max_value is None and threshold == 0:
            return tf.nn.leaky_relu(x, alpha=alpha)

        if threshold != 0:
            negative_part = tf.nn.relu(-x + threshold)
        else:
            negative_part = tf.nn.relu(-x)

    clip_max = max_value is not None

    if threshold != 0:
        # computes x for x > threshold else 0
        x = x * tf.cast(tf.greater(x, threshold), dtype=dtype)
    elif max_value == 6:
        # if no threshold, then can use nn.relu6 native TF op for performance
        x = tf.nn.relu6(x)
        clip_max = False
    else:
        x = tf.nn.relu(x)

    if clip_max:
        max_value = _constant_to_tensor(max_value, x.dtype.base_dtype)
        zero = _constant_to_tensor(0, x.dtype.base_dtype)
        x = tf.clip_by_value(x, zero, max_value)

    if alpha != 0.0:
        alpha = _to_tensor(alpha, x.dtype.base_dtype)
        x -= alpha * negative_part
    return x


@keras_export("keras.backend.elu")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def elu(x, alpha=1.0):
    """Exponential linear unit.

    Args:
        x: A tensor or variable to compute the activation function for.
        alpha: A scalar, slope of negative section.

    Returns:
        A tensor.
    """
    res = tf.nn.elu(x)
    if alpha == 1:
        return res
    else:
        return tf.where(x > 0, res, alpha * res)


@keras_export("keras.backend.softmax")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def softmax(x, axis=-1):
    """Softmax of a tensor.

    Args:
        x: A tensor or variable.
        axis: The dimension softmax would be performed on.
            The default is -1 which indicates the last dimension.

    Returns:
        A tensor.
    """
    if x.shape.rank <= 1:
        raise ValueError(
            f"Cannot apply softmax to a tensor that is 1D. Received input: {x}"
        )

    if isinstance(axis, int):
        output = tf.nn.softmax(x, axis=axis)
    else:
        # nn.softmax does not support tuple axis.
        numerator = tf.exp(x - tf.reduce_max(x, axis=axis, keepdims=True))
        denominator = tf.reduce_sum(numerator, axis=axis, keepdims=True)
        output = numerator / denominator

    # Cache the logits to use for crossentropy loss.
    output._keras_logits = x
    return output


@keras_export("keras.backend.softplus")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def softplus(x):
    """Softplus of a tensor.

    Args:
        x: A tensor or variable.

    Returns:
        A tensor.
    """
    return tf.math.softplus(x)


@keras_export("keras.backend.softsign")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def softsign(x):
    """Softsign of a tensor.

    Args:
        x: A tensor or variable.

    Returns:
        A tensor.
    """
    return tf.math.softsign(x)


def _get_logits(output, from_logits, op_type, fn_name):
    output_ = output
    from_logits_ = from_logits

    has_keras_logits = hasattr(output, "_keras_logits")
    if has_keras_logits:
        output_ = output._keras_logits
        from_logits_ = True

    from_expected_op_type = (
        not isinstance(output, (tf.__internal__.EagerTensor, tf.Variable))
        and output.op.type == op_type
    ) and not has_keras_logits

    if from_expected_op_type:
        # When softmax activation function is used for output operation, we
        # use logits from the softmax function directly to compute loss in order
        # to prevent collapsing zero when training.
        # See b/117284466
        assert len(output.op.inputs) == 1
        output_ = output.op.inputs[0]
        from_logits_ = True

    if from_logits and (has_keras_logits or from_expected_op_type):
        warnings.warn(
            f'"`{fn_name}` received `from_logits=True`, but '
            f"the `output` argument was produced by a {op_type} "
            "activation and thus does not represent logits. "
            "Was this intended?",
            stacklevel=2,
        )

    return output_, from_logits_


@keras_export("keras.backend.categorical_crossentropy")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def categorical_crossentropy(target, output, from_logits=False, axis=-1):
    """Categorical crossentropy between an output tensor and a target tensor.

    Args:
        target: A tensor of the same shape as `output`.
        output: A tensor resulting from a softmax
            (unless `from_logits` is True, in which
            case `output` is expected to be the logits).
        from_logits: Boolean, whether `output` is the
            result of a softmax, or is a tensor of logits.
        axis: Int specifying the channels axis. `axis=-1` corresponds to data
            format `channels_last`, and `axis=1` corresponds to data format
            `channels_first`.

    Returns:
        Output tensor.

    Raises:
        ValueError: if `axis` is neither -1 nor one of the axes of `output`.

    Example:

    >>> a = tf.constant([1., 0., 0., 0., 1., 0., 0., 0., 1.], shape=[3,3])
    >>> print(a)
    tf.Tensor(
      [[1. 0. 0.]
       [0. 1. 0.]
       [0. 0. 1.]], shape=(3, 3), dtype=float32)
    >>> b = tf.constant([.9, .05, .05, .05, .89, .06, .05, .01, .94],
    ...                 shape=[3, 3])
    >>> print(b)
    tf.Tensor(
      [[0.9  0.05 0.05]
       [0.05 0.89 0.06]
       [0.05 0.01 0.94]], shape=(3, 3), dtype=float32)
    >>> loss = tf.keras.backend.categorical_crossentropy(a, b)
    >>> print(np.around(loss, 5))
    [0.10536 0.11653 0.06188]
    >>> loss = tf.keras.backend.categorical_crossentropy(a, a)
    >>> print(np.around(loss, 5))
    [0. 0. 0.]

    """
    target = tf.convert_to_tensor(target)
    output = tf.convert_to_tensor(output)
    target.shape.assert_is_compatible_with(output.shape)

    output, from_logits = _get_logits(
        output, from_logits, "Softmax", "categorical_crossentropy"
    )
    if from_logits:
        return tf.nn.softmax_cross_entropy_with_logits(
            labels=target, logits=output, axis=axis
        )

    # Adjust the predictions so that the probability of
    # each class for every sample adds up to 1
    # This is needed to ensure that the cross entropy is
    # computed correctly.
    output = output / tf.reduce_sum(output, axis, True)

    # Compute cross entropy from probabilities.
    epsilon_ = _constant_to_tensor(epsilon(), output.dtype.base_dtype)
    output = tf.clip_by_value(output, epsilon_, 1.0 - epsilon_)
    return -tf.reduce_sum(target * tf.math.log(output), axis)


@keras_export("keras.backend.categorical_focal_crossentropy")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def categorical_focal_crossentropy(
    target,
    output,
    alpha=0.25,
    gamma=2.0,
    from_logits=False,
    axis=-1,
):
    """Computes the alpha balanced focal crossentropy loss.

    According to [Lin et al., 2018](https://arxiv.org/pdf/1708.02002.pdf), it
    helps to apply a focal factor to down-weight easy examples and focus more on
    hard examples. The general formula for the focal loss (FL)
    is as follows:

    `FL(p_t) = (1 − p_t)^gamma * log(p_t)`

    where `p_t` is defined as follows:
    `p_t = output if y_true == 1, else 1 - output`

    `(1 − p_t)^gamma` is the `modulating_factor`, where `gamma` is a focusing
    parameter. When `gamma` = 0, there is no focal effect on the cross entropy.
    `gamma` reduces the importance given to simple examples in a smooth manner.

    The authors use alpha-balanced variant of focal loss (FL) in the paper:
    `FL(p_t) = −alpha * (1 − p_t)^gamma * log(p_t)`

    where `alpha` is the weight factor for the classes. If `alpha` = 1, the
    loss won't be able to handle class imbalance properly as all
    classes will have the same weight. This can be a constant or a list of
    constants. If alpha is a list, it must have the same length as the number
    of classes.

    The formula above can be generalized to:
    `FL(p_t) = alpha * (1 − p_t)^gamma * CrossEntropy(target, output)`

    where minus comes from `CrossEntropy(target, output)` (CE).

    Extending this to multi-class case is straightforward:
    `FL(p_t) = alpha * (1 − p_t)^gamma * CategoricalCE(target, output)`

    Args:
        target: Ground truth values from the dataset.
        output: Predictions of the model.
        alpha: A weight balancing factor for all classes, default is `0.25` as
            mentioned in the reference. It can be a list of floats or a scalar.
            In the multi-class case, alpha may be set by inverse class
            frequency by using `compute_class_weight` from `sklearn.utils`.
        gamma: A focusing parameter, default is `2.0` as mentioned in the
            reference. It helps to gradually reduce the importance given to
            simple examples in a smooth manner.
        from_logits: Whether `output` is expected to be a logits tensor. By
            default, we consider that `output` encodes a probability
            distribution.
        axis: Int specifying the channels axis. `axis=-1` corresponds to data
             format `channels_last`, and `axis=1` corresponds to data format
             `channels_first`.

    Returns:
        A tensor.
    """
    target = tf.convert_to_tensor(target)
    output = tf.convert_to_tensor(output)
    target.shape.assert_is_compatible_with(output.shape)

    output, from_logits = _get_logits(
        output, from_logits, "Softmax", "categorical_focal_crossentropy"
    )

    if from_logits:
        output = tf.nn.softmax(output, axis=axis)

    # Adjust the predictions so that the probability of
    # each class for every sample adds up to 1
    # This is needed to ensure that the cross entropy is
    # computed correctly.
    output = output / tf.reduce_sum(output, axis=axis, keepdims=True)

    epsilon_ = _constant_to_tensor(epsilon(), output.dtype.base_dtype)
    output = tf.clip_by_value(output, epsilon_, 1.0 - epsilon_)

    # Calculate cross entropy
    cce = -target * tf.math.log(output)

    # Calculate factors
    modulating_factor = tf.pow(1.0 - output, gamma)
    weighting_factor = tf.multiply(modulating_factor, alpha)

    # Apply weighting factor
    focal_cce = tf.multiply(weighting_factor, cce)
    focal_cce = tf.reduce_sum(focal_cce, axis=axis)
    return focal_cce


@keras_export("keras.backend.sparse_categorical_crossentropy")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def sparse_categorical_crossentropy(
    target, output, from_logits=False, axis=-1, ignore_class=None
):
    """Categorical crossentropy with integer targets.

    Args:
        target: An integer tensor.
        output: A tensor resulting from a softmax
            (unless `from_logits` is True, in which
            case `output` is expected to be the logits).
        from_logits: Boolean, whether `output` is the
            result of a softmax, or is a tensor of logits.
        axis: Int specifying the channels axis. `axis=-1` corresponds to data
            format `channels_last`, and `axis=1` corresponds to data format
            `channels_first`.
        ignore_class: Optional integer. The ID of a class to be ignored
            during loss computation. This is useful, for example, in
            segmentation problems featuring a "void" class (commonly -1
            or 255) in segmentation maps.
            By default (`ignore_class=None`), all classes are considered.

    Returns:
        Output tensor.

    Raises:
        ValueError: if `axis` is neither -1 nor one of the axes of `output`.
    """
    target = tf.convert_to_tensor(target)
    output = tf.convert_to_tensor(output)

    target = cast(target, "int64")

    output, from_logits = _get_logits(
        output, from_logits, "Softmax", "sparse_categorical_crossentropy"
    )
    if not from_logits:
        epsilon_ = _constant_to_tensor(epsilon(), output.dtype.base_dtype)
        output = tf.clip_by_value(output, epsilon_, 1 - epsilon_)
        output = tf.math.log(output)

    # Permute output so that the last axis contains the logits/probabilities.
    if isinstance(output.shape, (tuple, list)):
        output_rank = len(output.shape)
    else:
        output_rank = output.shape.ndims
    if output_rank is not None:
        axis %= output_rank
        if axis != output_rank - 1:
            permutation = list(
                itertools.chain(
                    range(axis), range(axis + 1, output_rank), [axis]
                )
            )
            output = tf.compat.v1.transpose(output, perm=permutation)
    elif axis != -1:
        raise ValueError(
            "Cannot compute sparse categorical crossentropy with `axis={}` "
            "on an output tensor with unknown rank".format(axis)
        )

    # Try to adjust the shape so that rank of labels = rank of logits - 1.
    output_shape = tf.shape(output)
    target_rank = target.shape.ndims

    update_shape = (
        target_rank is not None
        and output_rank is not None
        and target_rank != output_rank - 1
    )
    if update_shape:
        target = flatten(target)
        output = tf.reshape(output, [-1, output_shape[-1]])

    if ignore_class is not None:
        valid_mask = tf.not_equal(target, cast(ignore_class, target.dtype))
        target = target[valid_mask]
        output = output[valid_mask]

    if py_any(_is_symbolic_tensor(v) for v in [target, output]):
        with get_graph().as_default():
            res = tf.nn.sparse_softmax_cross_entropy_with_logits(
                labels=target, logits=output
            )
    else:
        res = tf.nn.sparse_softmax_cross_entropy_with_logits(
            labels=target, logits=output
        )

    if ignore_class is not None:
        res_shape = cast(output_shape[:-1], "int64")
        valid_mask = tf.reshape(valid_mask, res_shape)
        res = tf.scatter_nd(tf.where(valid_mask), res, res_shape)
        res._keras_mask = valid_mask

        return res

    if update_shape and output_rank >= 3:
        # If our output includes timesteps or
        # spatial dimensions we need to reshape
        res = tf.reshape(res, output_shape[:-1])

    return res


@keras_export("keras.backend.binary_crossentropy")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def binary_crossentropy(target, output, from_logits=False):
    """Binary crossentropy between an output tensor and a target tensor.

    Args:
        target: A tensor with the same shape as `output`.
        output: A tensor.
        from_logits: Whether `output` is expected to be a logits tensor.
            By default, we consider that `output`
            encodes a probability distribution.

    Returns:
        A tensor.
    """
    target = tf.convert_to_tensor(target)
    output = tf.convert_to_tensor(output)

    output, from_logits = _get_logits(
        output, from_logits, "Sigmoid", "binary_crossentropy"
    )
    if from_logits:
        return tf.nn.sigmoid_cross_entropy_with_logits(
            labels=target, logits=output
        )

    epsilon_ = _constant_to_tensor(epsilon(), output.dtype.base_dtype)
    output = tf.clip_by_value(output, epsilon_, 1.0 - epsilon_)

    # Compute cross entropy from probabilities.
    bce = target * tf.math.log(output + epsilon())
    bce += (1 - target) * tf.math.log(1 - output + epsilon())
    return -bce


@keras_export("keras.backend.binary_focal_crossentropy")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def binary_focal_crossentropy(
    target,
    output,
    apply_class_balancing=False,
    alpha=0.25,
    gamma=2.0,
    from_logits=False,
):
    """Binary focal crossentropy between an output tensor and a target tensor.

    According to [Lin et al., 2018](https://arxiv.org/pdf/1708.02002.pdf), it
    helps to apply a focal factor to down-weight easy examples and focus more on
    hard examples. By default, the focal tensor is computed as follows:

    `focal_factor = (1 - output) ** gamma` for class 1
    `focal_factor = output ** gamma` for class 0
    where `gamma` is a focusing parameter. When `gamma` = 0, there is no focal
    effect on the binary crossentropy.

    If `apply_class_balancing == True`, this function also takes into account a
    weight balancing factor for the binary classes 0 and 1 as follows:

    `weight = alpha` for class 1 (`target == 1`)
    `weight = 1 - alpha` for class 0
    where `alpha` is a float in the range of `[0, 1]`.

    Args:
        target: A tensor with the same shape as `output`.
        output: A tensor.
        apply_class_balancing: A bool, whether to apply weight balancing on the
            binary classes 0 and 1.
        alpha: A weight balancing factor for class 1, default is `0.25` as
            mentioned in the reference. The weight for class 0 is `1.0 - alpha`.
        gamma: A focusing parameter, default is `2.0` as mentioned in the
            reference.
        from_logits: Whether `output` is expected to be a logits tensor. By
            default, we consider that `output` encodes a probability
            distribution.

    Returns:
        A tensor.
    """

    sigmoidal = sigmoid(output) if from_logits else output

    p_t = target * sigmoidal + (1 - target) * (1 - sigmoidal)

    # Calculate focal factor
    focal_factor = tf.pow(1.0 - p_t, gamma)

    # Binary crossentropy
    bce = binary_crossentropy(
        target=target,
        output=output,
        from_logits=from_logits,
    )
    focal_bce = focal_factor * bce

    if apply_class_balancing:
        weight = target * alpha + (1 - target) * (1 - alpha)
        focal_bce = weight * focal_bce

    return focal_bce


@keras_export("keras.backend.sigmoid")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def sigmoid(x):
    """Element-wise sigmoid.

    Args:
        x: A tensor or variable.

    Returns:
        A tensor.
    """
    output = tf.sigmoid(x)
    # Cache the logits to use for crossentropy loss.
    output._keras_logits = x
    return output


@keras_export("keras.backend.hard_sigmoid")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def hard_sigmoid(x):
    """Segment-wise linear approximation of sigmoid.

    Faster than sigmoid.
    Returns `0.` if `x < -2.5`, `1.` if `x > 2.5`.
    In `-2.5 <= x <= 2.5`, returns `0.2 * x + 0.5`.

    Args:
        x: A tensor or variable.

    Returns:
        A tensor.
    """
    point_two = _constant_to_tensor(0.2, x.dtype.base_dtype)
    point_five = _constant_to_tensor(0.5, x.dtype.base_dtype)
    x = tf.multiply(x, point_two)
    x = tf.add(x, point_five)
    x = tf.clip_by_value(x, 0.0, 1.0)
    return x


@keras_export("keras.backend.tanh")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def tanh(x):
    """Element-wise tanh.

    Args:
        x: A tensor or variable.

    Returns:
        A tensor.
    """
    return tf.tanh(x)


@keras_export("keras.backend.dropout")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def dropout(x, level, noise_shape=None, seed=None):
    """Sets entries in `x` to zero at random, while scaling the entire tensor.

    Args:
        x: tensor
        level: fraction of the entries in the tensor
            that will be set to 0.
        noise_shape: shape for randomly generated keep/drop flags,
            must be broadcastable to the shape of `x`
        seed: random seed to ensure determinism.

    Returns:
        A tensor.
    """
    if seed is None:
        seed = np.random.randint(10e6)
    return tf.nn.dropout(x, rate=level, noise_shape=noise_shape, seed=seed)


@keras_export("keras.backend.l2_normalize")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def l2_normalize(x, axis=None):
    """Normalizes a tensor wrt the L2 norm alongside the specified axis.

    Args:
        x: Tensor or variable.
        axis: axis along which to perform normalization.

    Returns:
        A tensor.
    """
    return tf.linalg.l2_normalize(x, axis=axis)


@keras_export("keras.backend.in_top_k")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def in_top_k(predictions, targets, k):
    """Returns whether the `targets` are in the top `k` `predictions`.

    Args:
        predictions: A tensor of shape `(batch_size, classes)` and type
          `float32`.
        targets: A 1D tensor of length `batch_size` and type `int32` or `int64`.
        k: An `int`, number of top elements to consider.

    Returns:
        A 1D tensor of length `batch_size` and type `bool`.
        `output[i]` is `True` if `predictions[i, targets[i]]` is within top-`k`
        values of `predictions[i]`.
    """
    return tf.compat.v1.math.in_top_k(predictions, targets, k)


# CONVOLUTIONS


def _preprocess_conv1d_input(x, data_format):
    """Transpose and cast the input before the conv1d.

    Args:
        x: input tensor.
        data_format: string, `"channels_last"` or `"channels_first"`.

    Returns:
        A tensor.
    """
    tf_data_format = "NWC"  # to pass TF Conv2dNative operations
    if data_format == "channels_first":
        if not _has_nchw_support():
            x = tf.compat.v1.transpose(x, (0, 2, 1))  # NCW -> NWC
        else:
            tf_data_format = "NCW"
    return x, tf_data_format


def _preprocess_conv2d_input(x, data_format, force_transpose=False):
    """Transpose and cast the input before the conv2d.

    Args:
        x: input tensor.
        data_format: string, `"channels_last"` or `"channels_first"`.
        force_transpose: Boolean. If True, the input will always be transposed
            from NCHW to NHWC if `data_format` is `"channels_first"`.
            If False, the transposition only occurs on CPU (GPU ops are
            assumed to support NCHW).

    Returns:
        A tensor.
    """
    tf_data_format = "NHWC"
    if data_format == "channels_first":
        if not _has_nchw_support() or force_transpose:
            x = tf.compat.v1.transpose(x, (0, 2, 3, 1))  # NCHW -> NHWC
        else:
            tf_data_format = "NCHW"
    return x, tf_data_format


def _preprocess_conv3d_input(x, data_format):
    """Transpose and cast the input before the conv3d.

    Args:
        x: input tensor.
        data_format: string, `"channels_last"` or `"channels_first"`.

    Returns:
        A tensor.
    """
    tf_data_format = "NDHWC"
    if data_format == "channels_first":
        if not _has_nchw_support():
            x = tf.compat.v1.transpose(x, (0, 2, 3, 4, 1))
        else:
            tf_data_format = "NCDHW"
    return x, tf_data_format


def _preprocess_padding(padding):
    """Convert keras' padding to TensorFlow's padding.

    Args:
        padding: string, one of 'same' , 'valid'

    Returns:
        a string, one of 'SAME', 'VALID'.

    Raises:
        ValueError: if invalid `padding'`
    """
    if padding == "same":
        padding = "SAME"
    elif padding == "valid":
        padding = "VALID"
    else:
        raise ValueError("Invalid padding: " + str(padding))
    return padding


@keras_export("keras.backend.conv1d")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def conv1d(
    x, kernel, strides=1, padding="valid", data_format=None, dilation_rate=1
):
    """1D convolution.

    Args:
        x: Tensor or variable.
        kernel: kernel tensor.
        strides: stride integer.
        padding: string, `"same"`, `"causal"` or `"valid"`.
        data_format: string, one of "channels_last", "channels_first".
        dilation_rate: integer dilate rate.

    Returns:
        A tensor, result of 1D convolution.

    Raises:
        ValueError: if `data_format` is neither `channels_last` or
        `channels_first`.
    """
    if data_format is None:
        data_format = image_data_format()
    if data_format not in {"channels_first", "channels_last"}:
        raise ValueError("Unknown data_format: " + str(data_format))

    kernel_shape = kernel.shape.as_list()
    if padding == "causal":
        # causal (dilated) convolution:
        left_pad = dilation_rate * (kernel_shape[0] - 1)
        x = temporal_padding(x, (left_pad, 0))
        padding = "valid"
    padding = _preprocess_padding(padding)

    x, tf_data_format = _preprocess_conv1d_input(x, data_format)
    x = tf.compat.v1.nn.convolution(
        input=x,
        filter=kernel,
        dilation_rate=dilation_rate,
        strides=strides,
        padding=padding,
        data_format=tf_data_format,
    )
    if data_format == "channels_first" and tf_data_format == "NWC":
        x = tf.compat.v1.transpose(x, (0, 2, 1))  # NWC -> NCW
    return x


@keras_export("keras.backend.conv2d")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def conv2d(
    x,
    kernel,
    strides=(1, 1),
    padding="valid",
    data_format=None,
    dilation_rate=(1, 1),
):
    """2D convolution.

    Args:
        x: Tensor or variable.
        kernel: kernel tensor.
        strides: strides tuple.
        padding: string, `"same"` or `"valid"`.
        data_format: `"channels_last"` or `"channels_first"`.
        dilation_rate: tuple of 2 integers.

    Returns:
        A tensor, result of 2D convolution.

    Raises:
        ValueError: if `data_format` is neither `channels_last` or
        `channels_first`.
    """
    if data_format is None:
        data_format = image_data_format()
    if data_format not in {"channels_first", "channels_last"}:
        raise ValueError("Unknown data_format: " + str(data_format))

    x, tf_data_format = _preprocess_conv2d_input(x, data_format)
    padding = _preprocess_padding(padding)
    x = tf.compat.v1.nn.convolution(
        input=x,
        filter=kernel,
        dilation_rate=dilation_rate,
        strides=strides,
        padding=padding,
        data_format=tf_data_format,
    )
    if data_format == "channels_first" and tf_data_format == "NHWC":
        x = tf.compat.v1.transpose(x, (0, 3, 1, 2))  # NHWC -> NCHW
    return x


@keras_export("keras.backend.conv2d_transpose")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def conv2d_transpose(
    x,
    kernel,
    output_shape,
    strides=(1, 1),
    padding="valid",
    data_format=None,
    dilation_rate=(1, 1),
):
    """2D deconvolution (i.e.

    transposed convolution).

    Args:
        x: Tensor or variable.
        kernel: kernel tensor.
        output_shape: 1D int tensor for the output shape.
        strides: strides tuple.
        padding: string, `"same"` or `"valid"`.
        data_format: string, `"channels_last"` or `"channels_first"`.
        dilation_rate: Tuple of 2 integers.

    Returns:
        A tensor, result of transposed 2D convolution.

    Raises:
        ValueError: if `data_format` is neither `channels_last` or
        `channels_first`.
    """
    if data_format is None:
        data_format = image_data_format()
    if data_format not in {"channels_first", "channels_last"}:
        raise ValueError("Unknown data_format: " + str(data_format))

    # `atrous_conv2d_transpose` only supports NHWC format, even on GPU.
    if data_format == "channels_first" and dilation_rate != (1, 1):
        force_transpose = True
    else:
        force_transpose = False

    x, tf_data_format = _preprocess_conv2d_input(
        x, data_format, force_transpose
    )

    if data_format == "channels_first" and tf_data_format == "NHWC":
        output_shape = (
            output_shape[0],
            output_shape[2],
            output_shape[3],
            output_shape[1],
        )
    if output_shape[0] is None:
        output_shape = (shape(x)[0],) + tuple(output_shape[1:])

    if isinstance(output_shape, (tuple, list)):
        output_shape = tf.stack(list(output_shape))

    padding = _preprocess_padding(padding)
    if tf_data_format == "NHWC":
        strides = (1,) + strides + (1,)
    else:
        strides = (1, 1) + strides

    if dilation_rate == (1, 1):
        x = tf.compat.v1.nn.conv2d_transpose(
            x,
            kernel,
            output_shape,
            strides,
            padding=padding,
            data_format=tf_data_format,
        )
    else:
        if dilation_rate[0] != dilation_rate[1]:
            raise ValueError(
                "Expected the 2 dimensions of the `dilation_rate` argument "
                "to be equal to each other. "
                f"Received: dilation_rate={dilation_rate}"
            )
        x = tf.nn.atrous_conv2d_transpose(
            x, kernel, output_shape, rate=dilation_rate[0], padding=padding
        )
    if data_format == "channels_first" and tf_data_format == "NHWC":
        x = tf.compat.v1.transpose(x, (0, 3, 1, 2))  # NHWC -> NCHW
    return x


def separable_conv1d(
    x,
    depthwise_kernel,
    pointwise_kernel,
    strides=1,
    padding="valid",
    data_format=None,
    dilation_rate=1,
):
    """1D convolution with separable filters.

    Args:
        x: input tensor
        depthwise_kernel: convolution kernel for the depthwise convolution.
        pointwise_kernel: kernel for the 1x1 convolution.
        strides: stride integer.
        padding: string, `"same"` or `"valid"`.
        data_format: string, `"channels_last"` or `"channels_first"`.
        dilation_rate: integer dilation rate.

    Returns:
        Output tensor.

    Raises:
        ValueError: if `data_format` is neither `channels_last` or
        `channels_first`.
    """
    if data_format is None:
        data_format = image_data_format()
    if data_format not in {"channels_first", "channels_last"}:
        raise ValueError("Unknown data_format: " + str(data_format))

    if isinstance(strides, int):
        strides = (strides,)
    if isinstance(dilation_rate, int):
        dilation_rate = (dilation_rate,)

    x, tf_data_format = _preprocess_conv1d_input(x, data_format)
    padding = _preprocess_padding(padding)
    if not isinstance(strides, tuple):
        strides = tuple(strides)
    if tf_data_format == "NWC":
        spatial_start_dim = 1
        strides = (1,) + strides * 2 + (1,)
    else:
        spatial_start_dim = 2
        strides = (1, 1) + strides * 2
    x = tf.expand_dims(x, spatial_start_dim)
    depthwise_kernel = tf.expand_dims(depthwise_kernel, 0)
    pointwise_kernel = tf.expand_dims(pointwise_kernel, 0)
    dilation_rate = (1,) + dilation_rate

    x = tf.nn.separable_conv2d(
        x,
        depthwise_kernel,
        pointwise_kernel,
        strides=strides,
        padding=padding,
        dilations=dilation_rate,
        data_format=tf_data_format,
    )

    x = tf.squeeze(x, [spatial_start_dim])

    if data_format == "channels_first" and tf_data_format == "NWC":
        x = tf.compat.v1.transpose(x, (0, 2, 1))  # NWC -> NCW

    return x


@keras_export("keras.backend.separable_conv2d")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def separable_conv2d(
    x,
    depthwise_kernel,
    pointwise_kernel,
    strides=(1, 1),
    padding="valid",
    data_format=None,
    dilation_rate=(1, 1),
):
    """2D convolution with separable filters.

    Args:
        x: input tensor
        depthwise_kernel: convolution kernel for the depthwise convolution.
        pointwise_kernel: kernel for the 1x1 convolution.
        strides: strides tuple (length 2).
        padding: string, `"same"` or `"valid"`.
        data_format: string, `"channels_last"` or `"channels_first"`.
        dilation_rate: tuple of integers,
            dilation rates for the separable convolution.

    Returns:
        Output tensor.

    Raises:
        ValueError: if `data_format` is neither `channels_last` or
        `channels_first`.
        ValueError: if `strides` is not a tuple of 2 integers.
    """
    if data_format is None:
        data_format = image_data_format()
    if data_format not in {"channels_first", "channels_last"}:
        raise ValueError("Unknown data_format: " + str(data_format))
    if len(strides) != 2:
        raise ValueError("`strides` must be a tuple of 2 integers.")

    x, tf_data_format = _preprocess_conv2d_input(x, data_format)
    padding = _preprocess_padding(padding)
    if not isinstance(strides, tuple):
        strides = tuple(strides)
    if tf_data_format == "NHWC":
        strides = (1,) + strides + (1,)
    else:
        strides = (1, 1) + strides

    x = tf.nn.separable_conv2d(
        x,
        depthwise_kernel,
        pointwise_kernel,
        strides=strides,
        padding=padding,
        dilations=dilation_rate,
        data_format=tf_data_format,
    )
    if data_format == "channels_first" and tf_data_format == "NHWC":
        x = tf.compat.v1.transpose(x, (0, 3, 1, 2))  # NHWC -> NCHW
    return x


@keras_export("keras.backend.depthwise_conv2d")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def depthwise_conv2d(
    x,
    depthwise_kernel,
    strides=(1, 1),
    padding="valid",
    data_format=None,
    dilation_rate=(1, 1),
):
    """2D convolution with separable filters.

    Args:
        x: input tensor
        depthwise_kernel: convolution kernel for the depthwise convolution.
        strides: strides tuple (length 2).
        padding: string, `"same"` or `"valid"`.
        data_format: string, `"channels_last"` or `"channels_first"`.
        dilation_rate: tuple of integers,
            dilation rates for the separable convolution.

    Returns:
        Output tensor.

    Raises:
        ValueError: if `data_format` is neither `channels_last` or
        `channels_first`.
    """
    if data_format is None:
        data_format = image_data_format()
    if data_format not in {"channels_first", "channels_last"}:
        raise ValueError("Unknown data_format: " + str(data_format))

    x, tf_data_format = _preprocess_conv2d_input(x, data_format)
    padding = _preprocess_padding(padding)
    if tf_data_format == "NHWC":
        strides = (1,) + strides + (1,)
    else:
        strides = (1, 1) + strides

    x = tf.nn.depthwise_conv2d(
        x,
        depthwise_kernel,
        strides=strides,
        padding=padding,
        dilations=dilation_rate,
        data_format=tf_data_format,
    )
    if data_format == "channels_first" and tf_data_format == "NHWC":
        x = tf.compat.v1.transpose(x, (0, 3, 1, 2))  # NHWC -> NCHW
    return x


@keras_export("keras.backend.conv3d")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def conv3d(
    x,
    kernel,
    strides=(1, 1, 1),
    padding="valid",
    data_format=None,
    dilation_rate=(1, 1, 1),
):
    """3D convolution.

    Args:
        x: Tensor or variable.
        kernel: kernel tensor.
        strides: strides tuple.
        padding: string, `"same"` or `"valid"`.
        data_format: string, `"channels_last"` or `"channels_first"`.
        dilation_rate: tuple of 3 integers.

    Returns:
        A tensor, result of 3D convolution.

    Raises:
        ValueError: if `data_format` is neither `channels_last` or
        `channels_first`.
    """
    if data_format is None:
        data_format = image_data_format()
    if data_format not in {"channels_first", "channels_last"}:
        raise ValueError("Unknown data_format: " + str(data_format))

    x, tf_data_format = _preprocess_conv3d_input(x, data_format)
    padding = _preprocess_padding(padding)
    x = tf.compat.v1.nn.convolution(
        input=x,
        filter=kernel,
        dilation_rate=dilation_rate,
        strides=strides,
        padding=padding,
        data_format=tf_data_format,
    )
    if data_format == "channels_first" and tf_data_format == "NDHWC":
        x = tf.compat.v1.transpose(x, (0, 4, 1, 2, 3))
    return x


def conv3d_transpose(
    x,
    kernel,
    output_shape,
    strides=(1, 1, 1),
    padding="valid",
    data_format=None,
):
    """3D deconvolution (i.e.

    transposed convolution).

    Args:
        x: input tensor.
        kernel: kernel tensor.
        output_shape: 1D int tensor for the output shape.
        strides: strides tuple.
        padding: string, "same" or "valid".
        data_format: string, `"channels_last"` or `"channels_first"`.

    Returns:
        A tensor, result of transposed 3D convolution.

    Raises:
        ValueError: if `data_format` is neither `channels_last` or
        `channels_first`.
    """
    if data_format is None:
        data_format = image_data_format()
    if data_format not in {"channels_first", "channels_last"}:
        raise ValueError("Unknown data_format: " + str(data_format))
    if isinstance(output_shape, (tuple, list)):
        output_shape = tf.stack(output_shape)

    x, tf_data_format = _preprocess_conv3d_input(x, data_format)

    if data_format == "channels_first" and tf_data_format == "NDHWC":
        output_shape = (
            output_shape[0],
            output_shape[2],
            output_shape[3],
            output_shape[4],
            output_shape[1],
        )
    if output_shape[0] is None:
        output_shape = (tf.shape(x)[0],) + tuple(output_shape[1:])
        output_shape = tf.stack(list(output_shape))

    padding = _preprocess_padding(padding)
    if tf_data_format == "NDHWC":
        strides = (1,) + strides + (1,)
    else:
        strides = (1, 1) + strides

    x = tf.compat.v1.nn.conv3d_transpose(
        x,
        kernel,
        output_shape,
        strides,
        padding=padding,
        data_format=tf_data_format,
    )
    if data_format == "channels_first" and tf_data_format == "NDHWC":
        x = tf.compat.v1.transpose(x, (0, 4, 1, 2, 3))
    return x


@keras_export("keras.backend.pool2d")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def pool2d(
    x,
    pool_size,
    strides=(1, 1),
    padding="valid",
    data_format=None,
    pool_mode="max",
):
    """2D Pooling.

    Args:
        x: Tensor or variable.
        pool_size: tuple of 2 integers.
        strides: tuple of 2 integers.
        padding: string, `"same"` or `"valid"`.
        data_format: string, `"channels_last"` or `"channels_first"`.
        pool_mode: string, `"max"` or `"avg"`.

    Returns:
        A tensor, result of 2D pooling.

    Raises:
        ValueError: if `data_format` is neither `"channels_last"` or
        `"channels_first"`.
        ValueError: if `pool_size` is not a tuple of 2 integers.
        ValueError: if `strides` is not a tuple of 2 integers.
        ValueError: if `pool_mode` is neither `"max"` or `"avg"`.
    """
    if data_format is None:
        data_format = image_data_format()
    if data_format not in {"channels_first", "channels_last"}:
        raise ValueError("Unknown data_format: " + str(data_format))
    if len(pool_size) != 2:
        raise ValueError("`pool_size` must be a tuple of 2 integers.")
    if len(strides) != 2:
        raise ValueError("`strides` must be a tuple of 2 integers.")

    x, tf_data_format = _preprocess_conv2d_input(x, data_format)
    padding = _preprocess_padding(padding)
    if tf_data_format == "NHWC":
        strides = (1,) + strides + (1,)
        pool_size = (1,) + pool_size + (1,)
    else:
        strides = (1, 1) + strides
        pool_size = (1, 1) + pool_size

    if pool_mode == "max":
        x = tf.compat.v1.nn.max_pool(
            x, pool_size, strides, padding=padding, data_format=tf_data_format
        )
    elif pool_mode == "avg":
        x = tf.compat.v1.nn.avg_pool(
            x, pool_size, strides, padding=padding, data_format=tf_data_format
        )
    else:
        raise ValueError("Invalid pooling mode: " + str(pool_mode))

    if data_format == "channels_first" and tf_data_format == "NHWC":
        x = tf.compat.v1.transpose(x, (0, 3, 1, 2))  # NHWC -> NCHW
    return x


@keras_export("keras.backend.pool3d")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def pool3d(
    x,
    pool_size,
    strides=(1, 1, 1),
    padding="valid",
    data_format=None,
    pool_mode="max",
):
    """3D Pooling.

    Args:
        x: Tensor or variable.
        pool_size: tuple of 3 integers.
        strides: tuple of 3 integers.
        padding: string, `"same"` or `"valid"`.
        data_format: string, `"channels_last"` or `"channels_first"`.
        pool_mode: string, `"max"` or `"avg"`.

    Returns:
        A tensor, result of 3D pooling.

    Raises:
        ValueError: if `data_format` is neither `"channels_last"` or
        `"channels_first"`.
        ValueError: if `pool_mode` is neither `"max"` or `"avg"`.
    """
    if data_format is None:
        data_format = image_data_format()
    if data_format not in {"channels_first", "channels_last"}:
        raise ValueError("Unknown data_format: " + str(data_format))

    x, tf_data_format = _preprocess_conv3d_input(x, data_format)
    padding = _preprocess_padding(padding)
    if tf_data_format == "NDHWC":
        strides = (1,) + strides + (1,)
        pool_size = (1,) + pool_size + (1,)
    else:
        strides = (1, 1) + strides
        pool_size = (1, 1) + pool_size

    if pool_mode == "max":
        x = tf.nn.max_pool3d(
            x, pool_size, strides, padding=padding, data_format=tf_data_format
        )
    elif pool_mode == "avg":
        x = tf.nn.avg_pool3d(
            x, pool_size, strides, padding=padding, data_format=tf_data_format
        )
    else:
        raise ValueError("Invalid pooling mode: " + str(pool_mode))

    if data_format == "channels_first" and tf_data_format == "NDHWC":
        x = tf.compat.v1.transpose(x, (0, 4, 1, 2, 3))
    return x


def local_conv(
    inputs, kernel, kernel_size, strides, output_shape, data_format=None
):
    """Apply N-D convolution with un-shared weights.

    Args:
        inputs: (N+2)-D tensor with shape
            (batch_size, channels_in, d_in1, ..., d_inN)
            if data_format='channels_first', or
            (batch_size, d_in1, ..., d_inN, channels_in)
            if data_format='channels_last'.
        kernel: the unshared weight for N-D convolution,
            with shape (output_items, feature_dim, channels_out), where
            feature_dim = np.prod(kernel_size) * channels_in,
            output_items = np.prod(output_shape).
        kernel_size: a tuple of N integers, specifying the
            spatial dimensions of the N-D convolution window.
        strides: a tuple of N integers, specifying the strides
            of the convolution along the spatial dimensions.
        output_shape: a tuple of (d_out1, ..., d_outN) specifying the spatial
            dimensionality of the output.
        data_format: string, "channels_first" or "channels_last".

    Returns:
        An (N+2)-D tensor with shape:
        (batch_size, channels_out) + output_shape
        if data_format='channels_first', or:
        (batch_size,) + output_shape + (channels_out,)
        if data_format='channels_last'.

    Raises:
        ValueError: if `data_format` is neither
        `channels_last` nor `channels_first`.
    """
    if data_format is None:
        data_format = image_data_format()
    if data_format not in {"channels_first", "channels_last"}:
        raise ValueError("Unknown data_format: " + str(data_format))

    kernel_shape = int_shape(kernel)
    feature_dim = kernel_shape[1]
    channels_out = kernel_shape[-1]
    ndims = len(output_shape)
    spatial_dimensions = list(range(ndims))

    xs = []
    output_axes_ticks = [range(axis_max) for axis_max in output_shape]
    for position in itertools.product(*output_axes_ticks):
        slices = [slice(None)]

        if data_format == "channels_first":
            slices.append(slice(None))

        slices.extend(
            slice(
                position[d] * strides[d],
                position[d] * strides[d] + kernel_size[d],
            )
            for d in spatial_dimensions
        )

        if data_format == "channels_last":
            slices.append(slice(None))

        xs.append(reshape(inputs[slices], (1, -1, feature_dim)))

    x_aggregate = concatenate(xs, axis=0)
    output = batch_dot(x_aggregate, kernel)
    output = reshape(output, output_shape + (-1, channels_out))

    if data_format == "channels_first":
        permutation = [ndims, ndims + 1] + spatial_dimensions
    else:
        permutation = [ndims] + spatial_dimensions + [ndims + 1]

    return permute_dimensions(output, permutation)


@keras_export("keras.backend.local_conv1d")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def local_conv1d(inputs, kernel, kernel_size, strides, data_format=None):
    """Apply 1D conv with un-shared weights.

    Args:
        inputs: 3D tensor with shape:
            (batch_size, steps, input_dim)
            if data_format is "channels_last" or
            (batch_size, input_dim, steps)
            if data_format is "channels_first".
        kernel: the unshared weight for convolution,
            with shape (output_length, feature_dim, filters).
        kernel_size: a tuple of a single integer,
            specifying the length of the 1D convolution window.
        strides: a tuple of a single integer,
            specifying the stride length of the convolution.
        data_format: the data format, channels_first or channels_last.

    Returns:
        A 3d tensor with shape:
        (batch_size, output_length, filters)
        if data_format='channels_first'
        or 3D tensor with shape:
        (batch_size, filters, output_length)
        if data_format='channels_last'.
    """
    output_shape = (kernel.shape[0],)
    return local_conv(
        inputs, kernel, kernel_size, strides, output_shape, data_format
    )


@keras_export("keras.backend.local_conv2d")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def local_conv2d(
    inputs, kernel, kernel_size, strides, output_shape, data_format=None
):
    """Apply 2D conv with un-shared weights.

    Args:
        inputs: 4D tensor with shape:
            (batch_size, filters, new_rows, new_cols)
            if data_format='channels_first'
            or 4D tensor with shape:
            (batch_size, new_rows, new_cols, filters)
            if data_format='channels_last'.
        kernel: the unshared weight for convolution,
            with shape (output_items, feature_dim, filters).
        kernel_size: a tuple of 2 integers, specifying the
            width and height of the 2D convolution window.
        strides: a tuple of 2 integers, specifying the strides
            of the convolution along the width and height.
        output_shape: a tuple with (output_row, output_col).
        data_format: the data format, channels_first or channels_last.

    Returns:
        A 4D tensor with shape:
        (batch_size, filters, new_rows, new_cols)
        if data_format='channels_first'
        or 4D tensor with shape:
        (batch_size, new_rows, new_cols, filters)
        if data_format='channels_last'.
    """
    return local_conv(
        inputs, kernel, kernel_size, strides, output_shape, data_format
    )


@keras_export("keras.backend.bias_add")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def bias_add(x, bias, data_format=None):
    """Adds a bias vector to a tensor.

    Args:
        x: Tensor or variable.
        bias: Bias tensor to add.
        data_format: string, `"channels_last"` or `"channels_first"`.

    Returns:
        Output tensor.

    Raises:
        ValueError: In one of the two cases below:
                    1. invalid `data_format` argument.
                    2. invalid bias shape.
                       the bias should be either a vector or
                       a tensor with ndim(x) - 1 dimension
    """
    if data_format is None:
        data_format = image_data_format()
    if data_format not in {"channels_first", "channels_last"}:
        raise ValueError("Unknown data_format: " + str(data_format))
    bias_shape = int_shape(bias)
    if len(bias_shape) != 1 and len(bias_shape) != ndim(x) - 1:
        raise ValueError(
            "Unexpected bias dimensions %d, expect to be 1 or %d dimensions"
            % (len(bias_shape), ndim(x) - 1)
        )

    if len(bias_shape) == 1:
        if data_format == "channels_first":
            return tf.nn.bias_add(x, bias, data_format="NCHW")
        return tf.nn.bias_add(x, bias, data_format="NHWC")
    if ndim(x) in (3, 4, 5):
        if data_format == "channels_first":
            bias_reshape_axis = (1, bias_shape[-1]) + bias_shape[:-1]
            return x + reshape(bias, bias_reshape_axis)
        return x + reshape(bias, (1,) + bias_shape)
    return tf.nn.bias_add(x, bias)


# RANDOMNESS


@keras_export("keras.backend.random_normal")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def random_normal(shape, mean=0.0, stddev=1.0, dtype=None, seed=None):
    """Returns a tensor with normal distribution of values.

    It is an alias to `tf.random.normal`.

    Args:
        shape: A tuple of integers, the shape of tensor to create.
        mean: A float, the mean value of the normal distribution to draw
          samples. Defaults to `0.0`.
        stddev: A float, the standard deviation of the normal distribution
          to draw samples. Defaults to `1.0`.
        dtype: `tf.dtypes.DType`, dtype of returned tensor. None uses Keras
          backend dtype which is float32. Defaults to `None`.
        seed: Integer, random seed. Will use a random numpy integer when not
          specified.

    Returns:
        A tensor with normal distribution of values.

    Example:

    >>> random_normal_tensor = tf.keras.backend.random_normal(shape=(2,3),
    ... mean=0.0, stddev=1.0)
    >>> random_normal_tensor
    <tf.Tensor: shape=(2, 3), dtype=float32, numpy=...,
    dtype=float32)>
    """
    if dtype is None:
        dtype = floatx()
    if seed is None:
        seed = np.random.randint(10e6)
    return tf.random.normal(
        shape, mean=mean, stddev=stddev, dtype=dtype, seed=seed
    )


@keras_export("keras.backend.random_uniform")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def random_uniform(shape, minval=0.0, maxval=1.0, dtype=None, seed=None):
    """Returns a tensor with uniform distribution of values.

    Args:
        shape: A tuple of integers, the shape of tensor to create.
        minval: A float, lower boundary of the uniform distribution
            to draw samples.
        maxval: A float, upper boundary of the uniform distribution
            to draw samples.
        dtype: String, dtype of returned tensor.
        seed: Integer, random seed.

    Returns:
        A tensor.

    Example:

    >>> random_uniform_tensor = tf.keras.backend.random_uniform(shape=(2,3),
    ... minval=0.0, maxval=1.0)
    >>> random_uniform_tensor
    <tf.Tensor: shape=(2, 3), dtype=float32, numpy=...,
    dtype=float32)>
    """
    if dtype is None:
        dtype = floatx()
    if seed is None:
        seed = np.random.randint(10e6)
    return tf.random.uniform(
        shape, minval=minval, maxval=maxval, dtype=dtype, seed=seed
    )


@keras_export("keras.backend.random_binomial")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def random_binomial(shape, p=0.0, dtype=None, seed=None):
    """Returns a tensor with random binomial distribution of values.

    DEPRECATED, use `tf.keras.backend.random_bernoulli` instead.

    The binomial distribution with parameters `n` and `p` is the probability
    distribution of the number of successful Bernoulli process. Only supports
    `n` = 1 for now.

    Args:
        shape: A tuple of integers, the shape of tensor to create.
        p: A float, `0. <= p <= 1`, probability of binomial distribution.
        dtype: String, dtype of returned tensor.
        seed: Integer, random seed.

    Returns:
        A tensor.

    Example:

    >>> random_binomial_tensor = tf.keras.backend.random_binomial(shape=(2,3),
    ... p=0.5)
    >>> random_binomial_tensor
    <tf.Tensor: shape=(2, 3), dtype=float32, numpy=...,
    dtype=float32)>
    """
    warnings.warn(
        "`tf.keras.backend.random_binomial` is deprecated, "
        "and will be removed in a future version."
        "Please use `tf.keras.backend.random_bernoulli` instead.",
        stacklevel=2,
    )
    return random_bernoulli(shape, p, dtype, seed)


@keras_export("keras.backend.random_bernoulli")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def random_bernoulli(shape, p=0.0, dtype=None, seed=None):
    """Returns a tensor with random bernoulli distribution of values.

    Args:
        shape: A tuple of integers, the shape of tensor to create.
        p: A float, `0. <= p <= 1`, probability of bernoulli distribution.
        dtype: String, dtype of returned tensor.
        seed: Integer, random seed.

    Returns:
        A tensor.
    """
    if dtype is None:
        dtype = floatx()
    if seed is None:
        seed = np.random.randint(10e6)
    return tf.where(
        tf.random.uniform(shape, dtype=dtype, seed=seed) <= p,
        tf.ones(shape, dtype=dtype),
        tf.zeros(shape, dtype=dtype),
    )


@keras_export("keras.backend.truncated_normal")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def truncated_normal(shape, mean=0.0, stddev=1.0, dtype=None, seed=None):
    """Returns a tensor with truncated random normal distribution of values.

    The generated values follow a normal distribution
    with specified mean and standard deviation,
    except that values whose magnitude is more than
    two standard deviations from the mean are dropped and re-picked.

    Args:
        shape: A tuple of integers, the shape of tensor to create.
        mean: Mean of the values.
        stddev: Standard deviation of the values.
        dtype: String, dtype of returned tensor.
        seed: Integer, random seed.

    Returns:
        A tensor.
    """
    if dtype is None:
        dtype = floatx()
    if seed is None:
        seed = np.random.randint(10e6)
    return tf.random.truncated_normal(
        shape, mean, stddev, dtype=dtype, seed=seed
    )


# CTC
# TensorFlow has a native implementation, but it uses sparse tensors
# and therefore requires a wrapper for TF-Keras. The functions below convert
# dense to sparse tensors and also wraps up the beam search code that is
# in TensorFlow's CTC implementation


@keras_export("keras.backend.ctc_label_dense_to_sparse")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def ctc_label_dense_to_sparse(labels, label_lengths):
    """Converts CTC labels from dense to sparse.

    Args:
        labels: dense CTC labels.
        label_lengths: length of the labels.

    Returns:
        A sparse tensor representation of the labels.
    """
    label_shape = tf.shape(labels)
    num_batches_tns = tf.stack([label_shape[0]])
    max_num_labels_tns = tf.stack([label_shape[1]])

    def range_less_than(old_input, current_input):
        return tf.expand_dims(tf.range(tf.shape(old_input)[1]), 0) < tf.fill(
            max_num_labels_tns, current_input
        )

    init = tf.cast(tf.fill([1, label_shape[1]], 0), tf.bool)
    dense_mask = tf.compat.v1.scan(
        range_less_than, label_lengths, initializer=init, parallel_iterations=1
    )
    dense_mask = dense_mask[:, 0, :]

    label_array = tf.reshape(
        tf.tile(tf.range(0, label_shape[1]), num_batches_tns), label_shape
    )
    label_ind = tf.compat.v1.boolean_mask(label_array, dense_mask)

    batch_array = tf.compat.v1.transpose(
        tf.reshape(
            tf.tile(tf.range(0, label_shape[0]), max_num_labels_tns),
            reverse(label_shape, 0),
        )
    )
    batch_ind = tf.compat.v1.boolean_mask(batch_array, dense_mask)
    indices = tf.compat.v1.transpose(
        tf.reshape(concatenate([batch_ind, label_ind], axis=0), [2, -1])
    )

    vals_sparse = tf.compat.v1.gather_nd(labels, indices)

    return tf.SparseTensor(
        tf.cast(indices, tf.int64), vals_sparse, tf.cast(label_shape, tf.int64)
    )


@keras_export("keras.backend.ctc_batch_cost")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def ctc_batch_cost(y_true, y_pred, input_length, label_length):
    """Runs CTC loss algorithm on each batch element.

    Args:
        y_true: tensor `(samples, max_string_length)`
            containing the truth labels.
        y_pred: tensor `(samples, time_steps, num_categories)`
            containing the prediction, or output of the softmax.
        input_length: tensor `(samples, 1)` containing the sequence length for
            each batch item in `y_pred`.
        label_length: tensor `(samples, 1)` containing the sequence length for
            each batch item in `y_true`.

    Returns:
        Tensor with shape (samples,1) containing the
            CTC loss of each element.
    """
    label_length = tf.cast(tf.squeeze(label_length, axis=-1), tf.int32)
    input_length = tf.cast(tf.squeeze(input_length, axis=-1), tf.int32)
    sparse_labels = tf.cast(
        ctc_label_dense_to_sparse(y_true, label_length), tf.int32
    )

    y_pred = tf.math.log(
        tf.compat.v1.transpose(y_pred, perm=[1, 0, 2]) + epsilon()
    )

    return tf.expand_dims(
        tf.compat.v1.nn.ctc_loss(
            inputs=y_pred, labels=sparse_labels, sequence_length=input_length
        ),
        1,
    )


@keras_export("keras.backend.ctc_decode")
@tf.__internal__.dispatch.add_dispatch_support
@doc_controls.do_not_generate_docs
def ctc_decode(y_pred, input_length, greedy=True, beam_width=100, top_paths=1):
    """Decodes the output of a softmax.

    Can use either greedy search (also known as best path)
    or a constrained dictionary search.

    Args:
        y_pred: tensor `(samples, time_steps, num_categories)`
            containing the prediction, or output of the softmax.
        input_length: tensor `(samples, )` containing the sequence length for
            each batch item in `y_pred`.
        greedy: perform much faster best-path search if `true`.
            This does not use a dictionary.
        beam_width: if `greedy` is `false`: a beam search decoder will be used
            with a beam of this width.
        top_paths: if `greedy` is `false`,
            how many of the most probable paths will be returned.

    Returns:
        Tuple:
            List: if `greedy` is `true`, returns a list of one element that
                contains the decoded sequence.
                If `false`, returns the `top_paths` most probable
                decoded sequences.
                Each decoded sequence has shape (samples, time_steps).
                Important: blank labels are returned as `-1`.
            Tensor `(top_paths, )` that contains
                the log probability of each decoded sequence.
    """
    input_shape = shape(y_pred)
    num_samples, num_steps = input_shape[0], input_shape[1]
    y_pred = tf.math.log(
        tf.compat.v1.transpose(y_pred, perm=[1, 0, 2]) + epsilon()
    )
    input_length = tf.cast(input_length, tf.int32)

    if greedy:
        (decoded, log_prob) = tf.nn.ctc_greedy_decoder(
            inputs=y_pred, sequence_length=input_length
        )
    else:
        (decoded, log_prob) = tf.compat.v1.nn.ctc_beam_search_decoder(
            inputs=y_pred,
            sequence_length=input_length,
            beam_width=beam_width,
            top_paths=top_paths,
        )
    decoded_dense = []
    for st in decoded:
        st = tf.SparseTensor(st.indices, st.values, (num_samples, num_steps))
        decoded_dense.append(tf.sparse.to_dense(sp_input=st, default_value=-1))
    return (decoded_dense, log_prob)


# HIGH ORDER FUNCTIONS


@keras_export("keras.backend.map_fn")
@doc_controls.do_not_generate_docs
def map_fn(fn, elems, name=None, dtype=None):
    """Map the function fn over the elements elems and return the outputs.

    Args:
        fn: Callable that will be called upon each element in elems
        elems: tensor
        name: A string name for the map node in the graph
        dtype: Output data type.

    Returns:
        Tensor with dtype `dtype`.
    """
    return tf.compat.v1.map_fn(fn, elems, name=name, dtype=dtype)


@keras_export("keras.backend.foldl")
@doc_controls.do_not_generate_docs
def foldl(fn, elems, initializer=None, name=None):
    """Reduce elems using fn to combine them from left to right.

    Args:
        fn: Callable that will be called upon each element in elems and an
            accumulator, for instance `lambda acc, x: acc + x`
        elems: tensor
        initializer: The first value used (`elems[0]` in case of None)
        name: A string name for the foldl node in the graph

    Returns:
        Tensor with same type and shape as `initializer`.
    """
    return tf.compat.v1.foldl(fn, elems, initializer=initializer, name=name)


@keras_export("keras.backend.foldr")
@doc_controls.do_not_generate_docs
def foldr(fn, elems, initializer=None, name=None):
    """Reduce elems using fn to combine them from right to left.

    Args:
        fn: Callable that will be called upon each element in elems and an
            accumulator, for instance `lambda acc, x: acc + x`
        elems: tensor
        initializer: The first value used (`elems[-1]` in case of None)
        name: A string name for the foldr node in the graph

    Returns:
        Same type and shape as initializer
    """
    return tf.compat.v1.foldr(fn, elems, initializer=initializer, name=name)


# Load TF-Keras default configuration from config file if present.
# Set TF-Keras base dir path given KERAS_HOME env variable, if applicable.
# Otherwise either ~/.keras or /tmp.
if "KERAS_HOME" in os.environ:
    _keras_dir = os.environ.get("KERAS_HOME")
else:
    _keras_base_dir = os.path.expanduser("~")
    _keras_dir = os.path.join(_keras_base_dir, ".keras")
_config_path = os.path.expanduser(os.path.join(_keras_dir, "keras.json"))
if os.path.exists(_config_path):
    try:
        with open(_config_path) as fh:
            _config = json.load(fh)
    except ValueError:
        _config = {}
    _floatx = _config.get("floatx", floatx())
    assert _floatx in {"float16", "float32", "float64"}
    _epsilon = _config.get("epsilon", epsilon())
    assert isinstance(_epsilon, float)
    _image_data_format = _config.get("image_data_format", image_data_format())
    assert _image_data_format in {"channels_last", "channels_first"}
    set_floatx(_floatx)
    set_epsilon(_epsilon)
    set_image_data_format(_image_data_format)

# Save config file.
if not os.path.exists(_keras_dir):
    try:
        os.makedirs(_keras_dir)
    except OSError:
        # Except permission denied and potential race conditions
        # in multi-threaded environments.
        pass

if not os.path.exists(_config_path):
    _config = {
        "floatx": floatx(),
        "epsilon": epsilon(),
        "backend": "tensorflow",
        "image_data_format": image_data_format(),
    }
    try:
        with open(_config_path, "w") as f:
            f.write(json.dumps(_config, indent=4))
    except IOError:
        # Except permission denied.
        pass


def configure_and_create_distributed_session(distribution_strategy):
    """Configure session config and create a session with it."""

    def _create_session(distribution_strategy):
        """Create the Distributed Strategy session."""
        session_config = get_default_session_config()

        # If a session already exists, merge in its config; in the case there is
        # a conflict, take values of the existing config.
        global _SESSION
        if getattr(_SESSION, "session", None) and _SESSION.session._config:
            session_config.MergeFrom(_SESSION.session._config)

        if is_tpu_strategy(distribution_strategy):
            # TODO(priyag, yuefengz): Remove this workaround when Distribute
            # Coordinator is integrated with keras and we can create a session
            # from there.
            distribution_strategy.configure(session_config)
            master = (
                distribution_strategy.extended._tpu_cluster_resolver.master()
            )
            session = tf.compat.v1.Session(config=session_config, target=master)
        else:
            worker_context = dc.get_current_worker_context()
            if worker_context:
                dc_session_config = worker_context.session_config
                # Merge the default session config to the one from distribute
                # coordinator, which is fine for now since they don't have
                # conflicting configurations.
                dc_session_config.MergeFrom(session_config)
                session = tf.compat.v1.Session(
                    config=dc_session_config,
                    target=worker_context.master_target,
                )
            else:
                distribution_strategy.configure(session_config)
                session = tf.compat.v1.Session(config=session_config)

        set_session(session)

    if distribution_strategy.extended._in_multi_worker_mode():
        dc.run_distribute_coordinator(_create_session, distribution_strategy)
    else:
        _create_session(distribution_strategy)


def _is_tpu_strategy_class(clz):
    is_tpu_strat = lambda k: k.__name__.startswith("TPUStrategy")
    if is_tpu_strat(clz):
        return True
    return py_any(map(_is_tpu_strategy_class, clz.__bases__))


def is_tpu_strategy(strategy):
    """Returns whether input is a TPUStrategy instance or subclass instance."""
    return _is_tpu_strategy_class(strategy.__class__)


def _is_symbolic_tensor(x):
    return tf.is_tensor(x) and not isinstance(x, tf.__internal__.EagerTensor)


def convert_inputs_if_ragged(inputs):
    """Converts any ragged tensors to dense."""

    def _convert_ragged_input(inputs):
        if isinstance(inputs, tf.RaggedTensor):
            return inputs.to_tensor()
        return inputs

    flat_inputs = tf.nest.flatten(inputs)
    contains_ragged = py_any(
        isinstance(i, tf.RaggedTensor) for i in flat_inputs
    )

    if not contains_ragged:
        return inputs, None

    inputs = tf.nest.map_structure(_convert_ragged_input, inputs)
    # Multiple mask are not yet supported, so one mask is used on all inputs.
    # We approach this similarly when using row lengths to ignore steps.
    nested_row_lengths = tf.cast(
        flat_inputs[0].nested_row_lengths()[0], "int32"
    )
    return inputs, nested_row_lengths


def maybe_convert_to_ragged(
    is_ragged_input, output, nested_row_lengths, go_backwards=False
):
    """Converts any ragged input back to its initial structure."""
    if not is_ragged_input:
        return output

    if go_backwards:
        # Reverse based on the timestep dim, so that nested_row_lengths will
        # mask from the correct direction. Return the reverse ragged tensor.
        output = reverse(output, [1])
        ragged = tf.RaggedTensor.from_tensor(output, nested_row_lengths)
        return reverse(ragged, [1])
    else:
        return tf.RaggedTensor.from_tensor(output, nested_row_lengths)


class ContextValueCache(weakref.WeakKeyDictionary):
    """Container that caches (possibly tensor) values based on the context.

    This class is similar to defaultdict, where values may be produced by the
    default factory specified during initialization. This class also has a
    default value for the key (when key is `None`) -- the key is set to the
    current graph or eager context. The default factories for key and value are
    only used in `__getitem__` and `setdefault`. The `.get()` behavior remains
    the same.

    This object will return the value of the current graph or closest parent
    graph if the current graph is a function. This is to reflect the fact that
    if a tensor is created in eager/graph, child functions may capture that
    tensor.

    The default factory method may accept keyword arguments (unlike defaultdict,
    which only accepts callables with 0 arguments). To pass keyword arguments to
    `default_factory`, use the `setdefault` method instead of `__getitem__`.

    An example of how this class can be used in different contexts:

    ```
    cache = ContextValueCache(int)

    # Eager mode
    cache[None] += 2
    cache[None] += 4
    assert cache[None] == 6

    # Graph mode
    with tf.Graph().as_default() as g:
      cache[None] += 5
      cache[g] += 3
    assert cache[g] == 8
    ```

    Example of a default factory with arguments:

    ```
    cache = ContextValueCache(lambda x: x + 1)
    g = tf.get_default_graph()

    # Example with keyword argument.
    value = cache.setdefault(key=g, kwargs={'x': 3})
    assert cache[g] == 4
    ```
    """

    def __init__(self, default_factory):
        self.default_factory = default_factory
        weakref.WeakKeyDictionary.__init__(self)

    def _key(self):
        if tf.executing_eagerly():
            return _DUMMY_EAGER_GRAPH.key
        else:
            return tf.compat.v1.get_default_graph()

    def _get_parent_graph(self, graph):
        """Returns the parent graph or dummy eager object."""
        # TODO(b/149317164): Currently FuncGraphs use ops.get_default_graph() as
        # the outer graph. This results in outer_graph always being a Graph,
        # even in eager mode (get_default_graph will create a new Graph if there
        # isn't a default graph). Because of this bug, we have to specially set
        # the key when eager execution is enabled.
        parent_graph = graph.outer_graph
        if (
            not isinstance(parent_graph, tf.__internal__.FuncGraph)
            and tf.compat.v1.executing_eagerly_outside_functions()
        ):
            return _DUMMY_EAGER_GRAPH.key
        return parent_graph

    def _get_recursive(self, key):
        """Gets the value at key or the closest parent graph."""
        value = self.get(key)
        if value is not None:
            return value

        # Since FuncGraphs are able to capture tensors and variables from their
        # parent graphs, recursively search to see if there is a value stored
        # for one of the parent graphs.
        if isinstance(key, tf.__internal__.FuncGraph):
            return self._get_recursive(self._get_parent_graph(key))
        return None

    def __getitem__(self, key):
        """Gets the value at key (or current context), or sets default value.

        Args:
          key: May be `None` or `Graph`object. When `None`, the key is set to
            the current context.

        Returns:
          Either the cached or default value.
        """
        if key is None:
            key = self._key()

        value = self._get_recursive(key)
        if value is None:
            value = self[key] = self.default_factory()
        return value

    def setdefault(self, key=None, default=None, kwargs=None):
        """Sets the default value if key is not in dict, and returns the
        value."""
        if key is None:
            key = self._key()
        kwargs = kwargs or {}

        if default is None and key not in self:
            default = self.default_factory(**kwargs)
        return weakref.WeakKeyDictionary.setdefault(self, key, default)


# This dictionary holds a mapping {graph: learning_phase}. In eager mode, a
# dummy object is used.
# A learning phase is a bool tensor used to run TF-Keras models in
# either train mode (learning_phase == 1) or test mode (learning_phase == 0).
_GRAPH_LEARNING_PHASES = ContextValueCache(
    object_identity.ObjectIdentityWeakSet
)

# This dictionary holds a mapping between a graph and variables to initialize
# in the graph.
_GRAPH_VARIABLES = ContextValueCache(object_identity.ObjectIdentityWeakSet)

# This dictionary holds a mapping between a graph and TF optimizers created in
# the graph.
_GRAPH_TF_OPTIMIZERS = ContextValueCache(object_identity.ObjectIdentityWeakSet)