"""Commonly-used neural network operations not included in NumPy."""

import warnings

from keras.src import backend
from keras.src.api_export import keras_export
from keras.src.backend import KerasTensor
from keras.src.backend import any_symbolic_tensors
from keras.src.backend import standardize_data_format
from keras.src.backend.common.backend_utils import (
    compute_conv_transpose_output_shape,
)
from keras.src.ops import operation_utils
from keras.src.ops.operation import Operation
from keras.src.ops.operation_utils import reduce_shape


class Relu(Operation):
    def call(self, x):
        return backend.nn.relu(x)

    def compute_output_spec(self, x):
        return KerasTensor(x.shape, dtype=x.dtype)


@keras_export(["keras.ops.relu", "keras.ops.nn.relu"])
def relu(x):
    """Rectified linear unit activation function.

    It is defined as `f(x) = max(0, x)`.

    Args:
        x: Input tensor.

    Returns:
        A tensor with the same shape as `x`.

    Example:

    >>> x1 = keras.ops.convert_to_tensor([-1.0, 0.0, 1.0, 0.2])
    >>> keras.ops.relu(x1)
    array([0.0, 0.0, 1.0, 0.2], dtype=float32)
    """
    if any_symbolic_tensors((x,)):
        return Relu().symbolic_call(x)
    return backend.nn.relu(x)


class Relu6(Operation):
    def call(self, x):
        return backend.nn.relu6(x)

    def compute_output_spec(self, x):
        return KerasTensor(x.shape, dtype=x.dtype)


@keras_export(["keras.ops.relu6", "keras.ops.nn.relu6"])
def relu6(x):
    """Rectified linear unit activation function with upper bound of 6.

    It is defined as `f(x) = np.clip(x, 0, 6)`.

    Args:
        x: Input tensor.

    Returns:
        A tensor with the same shape as `x`.

    Example:

    >>> x = keras.ops.convert_to_tensor([-3.0, -2.0, 0.1, 0.2, 6.0, 8.0])
    >>> keras.ops.relu6(x)
    array([0.0, 0.0, 0.1, 0.2, 6.0, 6.0], dtype=float32)
    """
    if any_symbolic_tensors((x,)):
        return Relu6().symbolic_call(x)
    return backend.nn.relu6(x)


class Sigmoid(Operation):
    def call(self, x):
        return backend.nn.sigmoid(x)

    def compute_output_spec(self, x):
        return KerasTensor(x.shape, dtype=x.dtype)


@keras_export(["keras.ops.sigmoid", "keras.ops.nn.sigmoid"])
def sigmoid(x):
    """Sigmoid activation function.

    It is defined as `f(x) = 1 / (1 + exp(-x))`.

    Args:
        x: Input tensor.

    Returns:
        A tensor with the same shape as `x`.

    Example:

    >>> x = keras.ops.convert_to_tensor([-6.0, 1.0, 0.0, 1.0, 6.0])
    >>> keras.ops.sigmoid(x)
    array([0.00247262, 0.7310586, 0.5, 0.7310586, 0.9975274], dtype=float32)

    """
    if any_symbolic_tensors((x,)):
        return Sigmoid().symbolic_call(x)
    return backend.nn.sigmoid(x)


class Softplus(Operation):
    def call(self, x):
        return backend.nn.softplus(x)

    def compute_output_spec(self, x):
        return KerasTensor(x.shape, dtype=x.dtype)


@keras_export(["keras.ops.softplus", "keras.ops.nn.softplus"])
def softplus(x):
    """Softplus activation function.

    It is defined as `f(x) = log(exp(x) + 1)`, where `log` is the natural
    logarithm and `exp` is the exponential function.

    Args:
        x: Input tensor.

    Returns:
        A tensor with the same shape as `x`.

    Example:

    >>> x = keras.ops.convert_to_tensor([-0.555, 0.0, 0.555])
    >>> keras.ops.softplus(x)
    array([0.45366603, 0.6931472, 1.008666], dtype=float32)

    """
    if any_symbolic_tensors((x,)):
        return Softplus().symbolic_call(x)
    return backend.nn.softplus(x)


class Softsign(Operation):
    def call(self, x):
        return backend.nn.softsign(x)

    def compute_output_spec(self, x):
        return KerasTensor(x.shape, dtype=x.dtype)


@keras_export(["keras.ops.softsign", "keras.ops.nn.softsign"])
def softsign(x):
    """Softsign activation function.

    It is defined as `f(x) = x / (abs(x) + 1)`.

    Args:
        x: Input tensor.

    Returns:
        A tensor with the same shape as `x`.

    Example:

    >>> x = keras.ops.convert_to_tensor([-0.100, -10.0, 1.0, 0.0, 100.0])
    >>> keras.ops.softsign(x)
    Array([-0.09090909, -0.90909094, 0.5, 0.0, 0.990099], dtype=float32)

    """
    if any_symbolic_tensors((x,)):
        return Softsign().symbolic_call(x)
    return backend.nn.softsign(x)


class Silu(Operation):
    def call(self, x):
        return backend.nn.silu(x)

    def compute_output_spec(self, x):
        return KerasTensor(x.shape, dtype=x.dtype)


@keras_export(
    [
        "keras.ops.silu",
        "keras.ops.nn.silu",
        "keras.ops.swish",
        "keras.ops.nn.swish",
    ]
)
def silu(x):
    """Sigmoid Linear Unit (SiLU) activation function, also known as Swish.

    The SiLU activation function is computed by the sigmoid function multiplied
    by its input. It is defined as `f(x) = x * sigmoid(x)`.

    Args:
        x: Input tensor.

    Returns:
        A tensor with the same shape as `x`.

    Example:

    >>> x = keras.ops.convert_to_tensor([-6.0, 1.0, 0.0, 1.0, 6.0])
    >>> keras.ops.sigmoid(x)
    array([0.00247262, 0.7310586, 0.5, 0.7310586, 0.9975274], dtype=float32)
    >>> keras.ops.silu(x)
    array([-0.0148357, 0.7310586, 0.0, 0.7310586, 5.9851646], dtype=float32)

    """
    if any_symbolic_tensors((x,)):
        return Silu().symbolic_call(x)
    return backend.nn.silu(x)


class LogSigmoid(Operation):
    def call(self, x):
        return backend.nn.log_sigmoid(x)

    def compute_output_spec(self, x):
        return KerasTensor(x.shape, dtype=x.dtype)


@keras_export(
    [
        "keras.ops.log_sigmoid",
        "keras.ops.nn.log_sigmoid",
    ]
)
def log_sigmoid(x):
    """Logarithm of the sigmoid activation function.

    It is defined as `f(x) = log(1 / (1 + exp(-x)))`.

    Args:
        x: Input tensor.

    Returns:
        A tensor with the same shape as `x`.

    Example:

    >>> x = keras.ops.convert_to_tensor([-0.541391, 0.0, 0.50, 5.0])
    >>> keras.ops.log_sigmoid(x)
    array([-1.0000418, -0.6931472, -0.474077, -0.00671535], dtype=float32)

    """
    if any_symbolic_tensors((x,)):
        return LogSigmoid().symbolic_call(x)
    return backend.nn.log_sigmoid(x)


class LeakyRelu(Operation):
    def __init__(self, negative_slope=0.2):
        super().__init__()
        self.negative_slope = negative_slope

    def call(self, x):
        return backend.nn.leaky_relu(x, self.negative_slope)

    def compute_output_spec(self, x):
        return KerasTensor(x.shape, dtype=x.dtype)


@keras_export(["keras.ops.leaky_relu", "keras.ops.nn.leaky_relu"])
def leaky_relu(x, negative_slope=0.2):
    """Leaky version of a Rectified Linear Unit activation function.

    It allows a small gradient when the unit is not active, it is defined as:

    `f(x) = alpha * x for x < 0` or `f(x) = x for x >= 0`.

    Args:
        x: Input tensor.
        negative_slope: Slope of the activation function at x < 0.
            Defaults to `0.2`.

    Returns:
        A tensor with the same shape as `x`.

    Example:

    >>> x = np.array([-1., 0., 1.])
    >>> x_leaky_relu = keras.ops.leaky_relu(x)
    >>> print(x_leaky_relu)
    array([-0.2,  0. ,  1. ], shape=(3,), dtype=float64)

    """
    if any_symbolic_tensors((x,)):
        return LeakyRelu(negative_slope).symbolic_call(x)
    return backend.nn.leaky_relu(x, negative_slope=negative_slope)


class HardSigmoid(Operation):
    def call(self, x):
        return backend.nn.hard_sigmoid(x)

    def compute_output_spec(self, x):
        return KerasTensor(x.shape, dtype=x.dtype)


@keras_export(
    [
        "keras.ops.hard_sigmoid",
        "keras.ops.nn.hard_sigmoid",
    ]
)
def hard_sigmoid(x):
    """Hard sigmoid activation function.

    It is defined as:

    `0 if x < -2.5`, `1 if x > 2.5`, `(0.2 * x) + 0.5 if -2.5 <= x <= 2.5`.

    Args:
        x: Input tensor.

    Returns:
        A tensor with the same shape as `x`.

    Example:

    >>> x = np.array([-1., 0., 1.])
    >>> x_hard_sigmoid = keras.ops.hard_sigmoid(x)
    >>> print(x_hard_sigmoid)
    array([0.3, 0.5, 0.7], shape=(3,), dtype=float64)

    """
    if any_symbolic_tensors((x,)):
        return HardSigmoid().symbolic_call(x)
    return backend.nn.hard_sigmoid(x)


class HardSilu(Operation):
    def call(self, x):
        return backend.nn.hard_silu(x)

    def compute_output_spec(self, x):
        return KerasTensor(x.shape, dtype=x.dtype)


@keras_export(
    [
        "keras.ops.hard_silu",
        "keras.ops.nn.hard_silu",
        "keras.ops.hard_swish",
        "keras.ops.nn.hard_swish",
    ]
)
def hard_silu(x):
    """Hard SiLU activation function, also known as Hard Swish.

    It is defined as:

    - `0` if `if x < -3`
    - `x` if `x > 3`
    - `x * (x + 3) / 6` if `-3 <= x <= 3`

    It's a faster, piecewise linear approximation of the silu activation.

    Args:
        x: Input tensor.

    Returns:
        A tensor with the same shape as `x`.

    Example:

    >>> x = keras.ops.convert_to_tensor([-3.0, -1.0, 0.0, 1.0, 3.0])
    >>> keras.ops.hard_silu(x)
    array([-0.0, -0.3333333, 0.0, 0.6666667, 3.0], shape=(5,), dtype=float32)

    """
    if any_symbolic_tensors((x,)):
        return HardSilu().symbolic_call(x)
    return backend.nn.hard_silu(x)


class Elu(Operation):
    def __init__(self, alpha=1.0):
        super().__init__()
        self.alpha = alpha

    def call(self, x):
        return backend.nn.elu(x, alpha=self.alpha)

    def compute_output_spec(self, x):
        return KerasTensor(x.shape, dtype=x.dtype)


@keras_export(["keras.ops.elu", "keras.ops.nn.elu"])
def elu(x, alpha=1.0):
    """Exponential Linear Unit activation function.

    It is defined as:

    `f(x) =  alpha * (exp(x) - 1.) for x < 0`, `f(x) = x for x >= 0`.

    Args:
        x: Input tensor.
        alpha: A scalar, slope of positive section. Defaults to `1.0`.

    Returns:
        A tensor with the same shape as `x`.

    Example:

    >>> x = np.array([-1., 0., 1.])
    >>> x_elu = keras.ops.elu(x)
    >>> print(x_elu)
    array([-0.63212055, 0., 1.], shape=(3,), dtype=float64)

    """
    if any_symbolic_tensors((x,)):
        return Elu(alpha).symbolic_call(x)
    return backend.nn.elu(x, alpha=alpha)


class Selu(Operation):
    def call(self, x):
        return backend.nn.selu(x)

    def compute_output_spec(self, x):
        return KerasTensor(x.shape, dtype=x.dtype)


@keras_export(["keras.ops.selu", "keras.ops.nn.selu"])
def selu(x):
    """Scaled Exponential Linear Unit (SELU) activation function.

    It is defined as:

    `f(x) =  scale * alpha * (exp(x) - 1.) for x < 0`,
    `f(x) = scale * x for x >= 0`.

    Args:
        x: Input tensor.

    Returns:
        A tensor with the same shape as `x`.

    Example:

    >>> x = np.array([-1., 0., 1.])
    >>> x_selu = keras.ops.selu(x)
    >>> print(x_selu)
    array([-1.11133055, 0., 1.05070098], shape=(3,), dtype=float64)

    """
    if any_symbolic_tensors((x,)):
        return Selu().symbolic_call(x)
    return backend.nn.selu(x)


class Gelu(Operation):
    def __init__(self, approximate=True):
        super().__init__()
        self.approximate = approximate

    def call(self, x):
        return backend.nn.gelu(x, self.approximate)

    def compute_output_spec(self, x):
        return KerasTensor(x.shape, dtype=x.dtype)


@keras_export(["keras.ops.gelu", "keras.ops.nn.gelu"])
def gelu(x, approximate=True):
    """Gaussian Error Linear Unit (GELU) activation function.

    If `approximate` is `True`, it is defined as:
    `f(x) = 0.5 * x * (1 + tanh(sqrt(2 / pi) * (x + 0.044715 * x^3)))`

    Or if `approximate` is `False`, it is defined as:
    `f(x) = x * P(X <= x) = 0.5 * x * (1 + erf(x / sqrt(2)))`,
    where `P(X) ~ N(0, 1)`.

    Args:
        x: Input tensor.
        approximate: Approximate version of GELU activation. Defaults to `True`.

    Returns:
        A tensor with the same shape as `x`.

    Example:

    >>> x = np.array([-1., 0., 1.])
    >>> x_gelu = keras.ops.gelu(x)
    >>> print(x_gelu)
    array([-0.15865525, 0., 0.84134475], shape=(3,), dtype=float64)

    """
    if any_symbolic_tensors((x,)):
        return Gelu(approximate).symbolic_call(x)
    return backend.nn.gelu(x, approximate)


class Softmax(Operation):
    def __init__(self, axis=-1):
        super().__init__()
        self.axis = axis

    def call(self, x):
        return backend.nn.softmax(x, axis=self.axis)

    def compute_output_spec(self, x):
        return KerasTensor(x.shape, dtype=x.dtype)


@keras_export(["keras.ops.softmax", "keras.ops.nn.softmax"])
def softmax(x, axis=-1):
    """Softmax activation function.

    The elements of the output vector lie within the range `(0, 1)`, and their
    total sum is exactly 1 (excluding the floating point rounding error).

    Each vector is processed independently. The `axis` argument specifies the
    axis along which the function is applied within the input.

    It is defined as:
    `f(x) = exp(x) / sum(exp(x))`

    Args:
        x: Input tensor.
        axis: Integer, axis along which the softmax is applied.

    Returns:
        A tensor with the same shape as `x`.

    Example:

    >>> x = np.array([-1., 0., 1.])
    >>> x_softmax = keras.ops.softmax(x)
    >>> print(x_softmax)
    array([0.09003057, 0.24472847, 0.66524096], shape=(3,), dtype=float64)

    """
    # Don't use `backend.shape` since TensorFlow returns
    # symbolic tensors for unknown shape which can trigger
    # an error in TensorFlow graph execution.
    if isinstance(axis, int) and x.shape[axis] == 1:
        warnings.warn(
            f"You are using a softmax over axis {axis} "
            f"of a tensor of shape {x.shape}. This axis "
            "has size 1. The softmax operation will always return "
            "the value 1, which is likely not what you intended. "
            "Did you mean to use a sigmoid instead?"
        )
    if any_symbolic_tensors((x,)):
        return Softmax(axis).symbolic_call(x)
    if isinstance(axis, tuple):
        original_shape = x.shape
        new_shape = []
        skip_dims = set(axis)
        i = 0
        while i < len(original_shape):
            if i in skip_dims:
                size = 1
                while i in skip_dims:
                    size *= original_shape[i]
                    i += 1
                new_shape.append(size)
            else:
                new_shape.append(original_shape[i])
                i += 1
        x = backend.numpy.reshape(x, new_shape)
        x = backend.nn.softmax(x, axis=-1)
        x = backend.numpy.reshape(x, original_shape)
        return x
    else:
        return backend.nn.softmax(x, axis=axis)


class LogSoftmax(Operation):
    def __init__(self, axis=-1):
        super().__init__()
        self.axis = axis

    def call(self, x):
        return backend.nn.log_softmax(x, axis=self.axis)

    def compute_output_spec(self, x):
        return KerasTensor(x.shape, dtype=x.dtype)


@keras_export(
    [
        "keras.ops.log_softmax",
        "keras.ops.nn.log_softmax",
    ]
)
def log_softmax(x, axis=-1):
    """Log-softmax activation function.

    It is defined as:
    `f(x) = x - max(x) - log(sum(exp(x - max(x))))`

    Args:
        x: Input tensor.
        axis: Integer, axis along which the log-softmax is applied.
            Defaults to `-1`.

    Returns:
        A tensor with the same shape as `x`.

    Example:

    >>> x = np.array([-1., 0., 1.])
    >>> x_log_softmax = keras.ops.log_softmax(x)
    >>> print(x_log_softmax)
    array([-2.40760596, -1.40760596, -0.40760596], shape=(3,), dtype=float64)

    """
    if any_symbolic_tensors((x,)):
        return LogSoftmax(axis).symbolic_call(x)
    if isinstance(axis, tuple):
        original_shape = x.shape
        new_shape = []
        skip_dims = set(axis)
        i = 0
        while i < len(original_shape):
            if i in skip_dims:
                size = 1
                while i in skip_dims:
                    size *= original_shape[i]
                    i += 1
                new_shape.append(size)
            else:
                new_shape.append(original_shape[i])
                i += 1
        x = backend.numpy.reshape(x, new_shape)
        x = backend.nn.log_softmax(x, axis=-1)
        x = backend.numpy.reshape(x, original_shape)
        return x
    else:
        return backend.nn.log_softmax(x, axis=axis)


class MaxPool(Operation):
    def __init__(
        self,
        pool_size,
        strides=None,
        padding="valid",
        data_format=None,
    ):
        super().__init__()
        self.pool_size = pool_size
        self.strides = strides
        self.padding = padding.lower()
        self.data_format = data_format

    def call(self, inputs):
        return backend.nn.max_pool(
            inputs,
            self.pool_size,
            self.strides,
            self.padding,
            self.data_format,
        )

    def compute_output_spec(self, inputs):
        output_shape = operation_utils.compute_pooling_output_shape(
            inputs.shape,
            self.pool_size,
            self.strides,
            self.padding,
            self.data_format,
        )
        return KerasTensor(output_shape, dtype=inputs.dtype)


@keras_export(["keras.ops.max_pool", "keras.ops.nn.max_pool"])
def max_pool(
    inputs,
    pool_size,
    strides=None,
    padding="valid",
    data_format=None,
):
    """Max pooling operation.

    Args:
        inputs: Tensor of rank N+2. `inputs` has shape
            `(batch_size,) + inputs_spatial_shape + (num_channels,)` if
            `data_format="channels_last"`, or
            `(batch_size, num_channels) + inputs_spatial_shape` if
            `data_format="channels_first"`. Pooling happens over the spatial
            dimensions only.
        pool_size: int or tuple/list of integers of size
            `len(inputs_spatial_shape)`, specifying the size of the pooling
            window for each spatial dimension of the input tensor. If
            `pool_size` is int, then every spatial dimension shares the same
            `pool_size`.
        strides: int or tuple/list of integers of size
            `len(inputs_spatial_shape)`. The stride of the sliding window for
            each spatial dimension of the input tensor. If `strides` is int,
            then every spatial dimension shares the same `strides`.
        padding: string, either `"valid"` or `"same"`. `"valid"` means no
            padding is applied, and `"same"` results in padding evenly to the
            left/right or up/down of the input such that output has the
            same height/width dimension as the input when `strides=1`.
        data_format: A string, either `"channels_last"` or `"channels_first"`.
            `data_format` determines the ordering of the dimensions in the
            inputs. If `data_format="channels_last"`, `inputs` is of shape
            `(batch_size, ..., channels)` while if
            `data_format="channels_first"`, `inputs` is of shape
            `(batch_size, channels, ...)`.

    Returns:
        A tensor of rank N+2, the result of the max pooling operation.
    """
    data_format = standardize_data_format(data_format)
    padding = padding.lower()
    if any_symbolic_tensors((inputs,)):
        return MaxPool(
            pool_size,
            strides,
            padding,
            data_format,
        ).symbolic_call(inputs)
    return backend.nn.max_pool(inputs, pool_size, strides, padding, data_format)


class AveragePool(Operation):
    def __init__(
        self,
        pool_size,
        strides=None,
        padding="valid",
        data_format=None,
    ):
        super().__init__()
        self.pool_size = pool_size
        self.strides = strides
        self.padding = padding.lower()
        self.data_format = data_format

    def call(self, inputs):
        return backend.nn.average_pool(
            inputs,
            self.pool_size,
            self.strides,
            self.padding,
            self.data_format,
        )

    def compute_output_spec(self, inputs):
        output_shape = operation_utils.compute_pooling_output_shape(
            inputs.shape,
            self.pool_size,
            self.strides,
            self.padding,
            self.data_format,
        )
        return KerasTensor(output_shape, dtype=inputs.dtype)


@keras_export(
    [
        "keras.ops.average_pool",
        "keras.ops.nn.average_pool",
    ]
)
def average_pool(
    inputs,
    pool_size,
    strides=None,
    padding="valid",
    data_format=None,
):
    """Average pooling operation.

    Args:
        inputs: Tensor of rank N+2. `inputs` has shape
            `(batch_size,) + inputs_spatial_shape + (num_channels,)` if
            `data_format="channels_last"`, or
            `(batch_size, num_channels) + inputs_spatial_shape` if
            `data_format="channels_first"`. Pooling happens over the spatial
            dimensions only.
        pool_size: int or tuple/list of integers of size
            `len(inputs_spatial_shape)`, specifying the size of the pooling
            window for each spatial dimension of the input tensor. If
            `pool_size` is int, then every spatial dimension shares the same
            `pool_size`.
        strides: int or tuple/list of integers of size
            `len(inputs_spatial_shape)`. The stride of the sliding window for
            each spatial dimension of the input tensor. If `strides` is int,
            then every spatial dimension shares the same `strides`.
        padding: string, either `"valid"` or `"same"`. `"valid"` means no
            padding is applied, and `"same"` results in padding evenly to the
            left/right or up/down of the input such that output has the
            same height/width dimension as the input when `strides=1`.
        data_format: A string, either `"channels_last"` or `"channels_first"`.
            `data_format` determines the ordering of the dimensions in the
            inputs. If `data_format="channels_last"`, `inputs` is of shape
            `(batch_size, ..., channels)` while if
            `data_format="channels_first"`, `inputs` is of shape
            `(batch_size, channels, ...)`.

    Returns:
        A tensor of rank N+2, the result of the average pooling operation.
    """
    data_format = standardize_data_format(data_format)
    padding = padding.lower()
    if any_symbolic_tensors((inputs,)):
        return AveragePool(
            pool_size,
            strides,
            padding,
            data_format,
        ).symbolic_call(inputs)
    return backend.nn.average_pool(
        inputs, pool_size, strides, padding, data_format
    )


class Conv(Operation):
    def __init__(
        self,
        strides=1,
        padding="valid",
        data_format=None,
        dilation_rate=1,
    ):
        super().__init__()
        self.strides = strides
        self.padding = padding.lower()
        self.data_format = data_format
        self.dilation_rate = dilation_rate

    def call(self, inputs, kernel):
        return backend.nn.conv(
            inputs,
            kernel,
            strides=self.strides,
            padding=self.padding,
            data_format=self.data_format,
            dilation_rate=self.dilation_rate,
        )

    def compute_output_spec(self, inputs, kernel):
        output_shape = operation_utils.compute_conv_output_shape(
            inputs.shape,
            kernel.shape[-1],
            kernel.shape[:-2],
            self.strides,
            self.padding,
            self.data_format,
            self.dilation_rate,
        )
        return KerasTensor(output_shape, dtype=inputs.dtype)


@keras_export(["keras.ops.conv", "keras.ops.nn.conv"])
def conv(
    inputs,
    kernel,
    strides=1,
    padding="valid",
    data_format=None,
    dilation_rate=1,
):
    """General N-D convolution.

    This ops supports 1D, 2D and 3D convolution.

    Args:
        inputs: Tensor of rank N+2. `inputs` has shape
            `(batch_size,) + inputs_spatial_shape + (num_channels,)` if
            `data_format="channels_last"`, or
            `(batch_size, num_channels) + inputs_spatial_shape` if
            `data_format="channels_first"`.
        kernel: Tensor of rank N+2. `kernel` has shape
            `(kernel_spatial_shape, num_input_channels, num_output_channels)`.
            `num_input_channels` should match the number of channels in
            `inputs`.
        strides: int or int tuple/list of `len(inputs_spatial_shape)`,
            specifying the strides of the convolution along each spatial
            dimension. If `strides` is int, then every spatial dimension shares
            the same `strides`.
        padding: string, either `"valid"` or `"same"`. `"valid"` means no
            padding is applied, and `"same"` results in padding evenly to the
            left/right or up/down of the input such that output has the
            same height/width dimension as the input when `strides=1`.
        data_format: A string, either `"channels_last"` or `"channels_first"`.
            `data_format` determines the ordering of the dimensions in the
            inputs. If `data_format="channels_last"`, `inputs` is of shape
            `(batch_size, ..., channels)` while if
            `data_format="channels_first"`, `inputs` is of shape
            `(batch_size, channels, ...)`.
        dilation_rate: int or int tuple/list of `len(inputs_spatial_shape)`,
            specifying the dilation rate to use for dilated convolution. If
            `dilation_rate` is int, then every spatial dimension shares
            the same `dilation_rate`.

    Returns:
        A tensor of rank N+2, the result of the conv operation.
    """
    data_format = standardize_data_format(data_format)
    padding = padding.lower()
    if any_symbolic_tensors((inputs,)):
        return Conv(strides, padding, data_format, dilation_rate).symbolic_call(
            inputs, kernel
        )
    return backend.nn.conv(
        inputs, kernel, strides, padding, data_format, dilation_rate
    )


class DepthwiseConv(Operation):
    def __init__(
        self,
        strides=1,
        padding="valid",
        data_format=None,
        dilation_rate=1,
    ):
        super().__init__()
        self.strides = strides
        self.padding = padding.lower()
        self.data_format = data_format
        self.dilation_rate = dilation_rate

    def call(self, inputs, kernel):
        return backend.nn.depthwise_conv(
            inputs,
            kernel,
            self.strides,
            self.padding,
            self.data_format,
            self.dilation_rate,
        )

    def compute_output_spec(self, inputs, kernel):
        output_shape = operation_utils.compute_conv_output_shape(
            inputs.shape,
            kernel.shape[-1] * kernel.shape[-2],
            kernel.shape[:-2],
            self.strides,
            self.padding,
            self.data_format,
            self.dilation_rate,
        )
        return KerasTensor(output_shape, dtype=inputs.dtype)


@keras_export(
    [
        "keras.ops.depthwise_conv",
        "keras.ops.nn.depthwise_conv",
    ]
)
def depthwise_conv(
    inputs,
    kernel,
    strides=1,
    padding="valid",
    data_format=None,
    dilation_rate=1,
):
    """General N-D depthwise convolution.

    This ops supports 1D and 2D depthwise convolution.

    Args:
        inputs: Tensor of rank N+2. `inputs` has shape
            `(batch_size,) + inputs_spatial_shape + (num_channels,)` if
            `data_format="channels_last"`, or
            `(batch_size, num_channels) + inputs_spatial_shape` if
            `data_format="channels_first"`.
        kernel: Tensor of rank N+2. `kernel` has shape
            [kernel_spatial_shape, num_input_channels, num_channels_multiplier],
            `num_input_channels` should match the number of channels in
            `inputs`.
        strides: int or int tuple/list of `len(inputs_spatial_shape)`,
            specifying the strides of the convolution along each spatial
            dimension. If `strides` is int, then every spatial dimension shares
            the same `strides`.
        padding: string, either `"valid"` or `"same"`. `"valid"` means no
            padding is applied, and `"same"` results in padding evenly to the
            left/right or up/down of the input such that output has the
            same height/width dimension as the input when `strides=1`.
        data_format: A string, either `"channels_last"` or `"channels_first"`.
            `data_format` determines the ordering of the dimensions in the
            inputs. If `data_format="channels_last"`, `inputs` is of shape
            `(batch_size, ..., channels)` while if
            `data_format="channels_first"`, `inputs` is of shape
            `(batch_size, channels, ...)`.
        dilation_rate: int or int tuple/list of `len(inputs_spatial_shape)`,
            specifying the dilation rate to use for dilated convolution. If
            `dilation_rate` is int, then every spatial dimension shares
            the same `dilation_rate`.

    Returns:
        A tensor of rank N+2, the result of the depthwise conv operation.
    """
    data_format = standardize_data_format(data_format)
    padding = padding.lower()
    if any_symbolic_tensors((inputs,)):
        return DepthwiseConv(
            strides, padding, data_format, dilation_rate
        ).symbolic_call(inputs, kernel)
    return backend.nn.depthwise_conv(
        inputs,
        kernel,
        strides,
        padding,
        data_format,
        dilation_rate,
    )


class SeparableConv(Operation):
    def __init__(
        self,
        strides=1,
        padding="valid",
        data_format=None,
        dilation_rate=1,
    ):
        super().__init__()
        self.strides = strides
        self.padding = padding.lower()
        self.data_format = data_format
        self.dilation_rate = dilation_rate

    def call(self, inputs, depthwise_kernel, pointwise_kernel):
        return backend.nn.separable_conv(
            inputs,
            depthwise_kernel,
            pointwise_kernel,
            self.strides,
            self.padding,
            self.data_format,
            self.dilation_rate,
        )

    def compute_output_spec(self, inputs, depthwise_kernel, pointwise_kernel):
        output_shape = list(
            depthwise_conv(
                inputs,
                depthwise_kernel,
                self.strides,
                self.padding,
                self.data_format,
                self.dilation_rate,
            ).shape
        )
        if self.data_format == "channels_last":
            output_shape[-1] = pointwise_kernel.shape[-1]
        else:
            output_shape[1] = pointwise_kernel.shape[-1]
        return KerasTensor(output_shape, dtype=inputs.dtype)


@keras_export(
    [
        "keras.ops.separable_conv",
        "keras.ops.nn.separable_conv",
    ]
)
def separable_conv(
    inputs,
    depthwise_kernel,
    pointwise_kernel,
    strides=1,
    padding="valid",
    data_format=None,
    dilation_rate=1,
):
    """General N-D separable convolution.

    This ops supports 1D and 2D separable convolution. `separable_conv` is
    a depthwise conv followed by a pointwise conv.

    Args:
        inputs: Tensor of rank N+2. `inputs` has shape
            `(batch_size,) + inputs_spatial_shape + (num_channels,)` if
            `data_format="channels_last"`, or
            `(batch_size, num_channels) + inputs_spatial_shape` if
            `data_format="channels_first"`.
        depthwise_kernel: Tensor of rank N+2. `depthwise_kernel` has shape
            [kernel_spatial_shape, num_input_channels, num_channels_multiplier],
            `num_input_channels` should match the number of channels in
            `inputs`.
        pointwise_kernel: Tensor of rank N+2. `pointwise_kernel` has shape
            `(*ones_like(kernel_spatial_shape),
            num_input_channels * num_channels_multiplier, num_output_channels)`.
        strides: int or int tuple/list of `len(inputs_spatial_shape)`,
            specifying the strides of the convolution along each spatial
            dimension. If `strides` is int, then every spatial dimension shares
            the same `strides`.
        padding: string, either `"valid"` or `"same"`. `"valid"` means no
            padding is applied, and `"same"` results in padding evenly to the
            left/right or up/down of the input such that output has the
            same height/width dimension as the input when `strides=1`.
        data_format: A string, either `"channels_last"` or `"channels_first"`.
            `data_format` determines the ordering of the dimensions in the
            inputs. If `data_format="channels_last"`, `inputs` is of shape
            `(batch_size, ..., channels)` while if
            `data_format="channels_first"`, `inputs` is of shape
            `(batch_size, channels, ...)`.
        dilation_rate: int or int tuple/list of `len(inputs_spatial_shape)`,
            specifying the dilation rate to use for dilated convolution. If
            `dilation_rate` is int, then every spatial dimension shares
            the same `dilation_rate`.

    Returns:
        A tensor of rank N+2, the result of the depthwise conv operation.
    """
    data_format = standardize_data_format(data_format)
    padding = padding.lower()
    if any_symbolic_tensors((inputs,)):
        return SeparableConv(
            strides,
            padding,
            data_format,
            dilation_rate,
        ).symbolic_call(inputs, depthwise_kernel, pointwise_kernel)
    return backend.nn.separable_conv(
        inputs,
        depthwise_kernel,
        pointwise_kernel,
        strides,
        padding,
        data_format,
        dilation_rate,
    )


class ConvTranspose(Operation):
    def __init__(
        self,
        strides,
        padding="valid",
        output_padding=None,
        data_format=None,
        dilation_rate=1,
    ):
        super().__init__()
        self.strides = strides
        self.output_padding = output_padding
        self.padding = padding.lower()
        self.data_format = data_format
        self.dilation_rate = dilation_rate

    def call(
        self,
        inputs,
        kernel,
    ):
        return backend.nn.conv_transpose(
            inputs,
            kernel,
            self.strides,
            self.output_padding,
            self.padding,
            self.data_format,
            self.dilation_rate,
        )

    def compute_output_spec(self, inputs, kernel):
        kernel_size = kernel.shape[:-2]
        filters = kernel.shape[-2]
        output_shape = compute_conv_transpose_output_shape(
            inputs.shape,
            kernel_size,
            filters,
            self.strides,
            self.padding,
            self.output_padding,
            self.data_format,
            self.dilation_rate,
        )
        return KerasTensor(output_shape, dtype=inputs.dtype)


@keras_export(
    [
        "keras.ops.conv_transpose",
        "keras.ops.nn.conv_transpose",
    ]
)
def conv_transpose(
    inputs,
    kernel,
    strides,
    padding="valid",
    output_padding=None,
    data_format=None,
    dilation_rate=1,
):
    """General N-D convolution transpose.

    Also known as de-convolution. This ops supports 1D, 2D and 3D convolution.

    Args:
        inputs: Tensor of rank N+2. `inputs` has shape
            `(batch_size,) + inputs_spatial_shape + (num_channels,)` if
            `data_format="channels_last"`, or
            `(batch_size, num_channels) + inputs_spatial_shape` if
            `data_format="channels_first"`.
        kernel: Tensor of rank N+2. `kernel` has shape
            [kernel_spatial_shape, num_output_channels, num_input_channels],
            `num_input_channels` should match the number of channels in
            `inputs`.
        strides: int or int tuple/list of `len(inputs_spatial_shape)`,
            specifying the strides of the convolution along each spatial
            dimension. If `strides` is int, then every spatial dimension shares
            the same `strides`.
        padding: string, either `"valid"` or `"same"`. `"valid"` means no
            padding is applied, and `"same"` results in padding evenly to the
            left/right or up/down of the input such that output has the
            same height/width dimension as the input when `strides=1`.
        output_padding: int or int tuple/list of `len(inputs_spatial_shape)`,
            specifying the amount of padding along the height and width of
            the output tensor. Can be a single integer to specify the same
            value for all spatial dimensions. The amount of output padding
            along a given dimension must be lower than the stride along that
            same dimension. If set to `None` (default), the output shape is
            inferred.
        data_format: A string, either `"channels_last"` or `"channels_first"`.
            `data_format` determines the ordering of the dimensions in the
            inputs. If `data_format="channels_last"`, `inputs` is of shape
            `(batch_size, ..., channels)` while if
            `data_format="channels_first"`, `inputs` is of shape
            `(batch_size, channels, ...)`.
        dilation_rate: int or int tuple/list of `len(inputs_spatial_shape)`,
            specifying the dilation rate to use for dilated convolution. If
            `dilation_rate` is int, then every spatial dimension shares
            the same `dilation_rate`.

    Returns:
        A tensor of rank N+2, the result of the conv operation.
    """
    data_format = standardize_data_format(data_format)
    padding = padding.lower()
    if any_symbolic_tensors((inputs,)):
        return ConvTranspose(
            strides, padding, output_padding, data_format, dilation_rate
        ).symbolic_call(inputs, kernel)
    return backend.nn.conv_transpose(
        inputs,
        kernel,
        strides,
        padding,
        output_padding,
        data_format,
        dilation_rate,
    )


class OneHot(Operation):
    def __init__(self, num_classes, axis=-1, dtype=None, sparse=False):
        super().__init__()
        self.num_classes = num_classes
        self.axis = axis
        self.dtype = dtype or backend.floatx()
        self.sparse = sparse

    def call(self, x):
        return backend.nn.one_hot(
            x,
            self.num_classes,
            axis=self.axis,
            dtype=self.dtype,
            sparse=self.sparse,
        )

    def compute_output_spec(self, x):
        x_shape = list(getattr(x, "shape", []))
        if self.axis == -1:
            x_shape.append(self.num_classes)
        elif self.axis >= 0 and self.axis < len(x_shape):
            x_shape.insert(self.axis, self.num_classes)
        else:
            raise ValueError(
                f"axis must be -1 or between [0, {len(x.shape)}), but "
                f"received {self.axis}."
            )
        return KerasTensor(x_shape, dtype=self.dtype, sparse=self.sparse)


@keras_export(["keras.ops.one_hot", "keras.ops.nn.one_hot"])
def one_hot(x, num_classes, axis=-1, dtype=None, sparse=False):
    """Converts integer tensor `x` into a one-hot tensor.

    The one-hot encoding is a representation where each integer value is
    converted into a binary vector with a length equal to `num_classes`,
    and the index corresponding to the integer value is marked as 1, while
    all other indices are marked as 0.

    Args:
        x: Integer tensor to be encoded. The shape can be
            arbitrary, but the dtype should be integer.
        num_classes: Number of classes for the one-hot encoding.
        axis: Axis along which the encoding is performed. Defaults to
            `-1`, which represents the last axis.
        dtype: (Optional) Data type of the output tensor. If not
            provided, it defaults to the default data type of the backend.
        sparse: Whether to return a sparse tensor; for backends that support
            sparse tensors.

    Returns:
        Integer tensor: One-hot encoded tensor with the same shape as `x`
        except for the specified `axis` dimension, which will have
        a length of `num_classes`. The dtype of the output tensor
        is determined by `dtype` or the default data type of the backend.

    Example:

    >>> x = keras.ops.convert_to_tensor([1, 3, 2, 0])
    >>> one_hot(x, num_classes=4)
    array([[0. 1. 0. 0.]
           [0. 0. 0. 1.]
           [0. 0. 1. 0.]
           [1. 0. 0. 0.]], shape=(4, 4), dtype=float32)
    """
    if any_symbolic_tensors((x,)):
        return OneHot(
            num_classes, axis=axis, dtype=dtype, sparse=sparse
        ).symbolic_call(x)
    return backend.nn.one_hot(
        x,
        num_classes,
        axis=axis,
        dtype=dtype or backend.floatx(),
        sparse=sparse,
    )


class BinaryCrossentropy(Operation):
    def __init__(self, from_logits=False):
        super().__init__()
        self.from_logits = from_logits

    def call(self, target, output):
        return backend.nn.binary_crossentropy(
            target, output, from_logits=self.from_logits
        )

    def compute_output_spec(self, target, output):
        if target.shape != output.shape:
            raise ValueError(
                "Arguments `target` and `output` must have the same shape. "
                "Received: "
                f"target.shape={target.shape}, output.shape={output.shape}"
            )
        return KerasTensor(output.shape, dtype=output.dtype)


@keras_export(
    [
        "keras.ops.binary_crossentropy",
        "keras.ops.nn.binary_crossentropy",
    ]
)
def binary_crossentropy(target, output, from_logits=False):
    """Computes binary cross-entropy loss between target and output tensor.

    The binary cross-entropy loss is commonly used in binary
    classification tasks where each input sample belongs to one
    of the two classes. It measures the dissimilarity between the
    target and output probabilities or logits.

    Args:
        target: The target tensor representing the true binary labels.
            Its shape should match the shape of the `output` tensor.
        output: The output tensor representing the predicted probabilities
            or logits. Its shape should match the shape of the
            `target` tensor.
        from_logits: (optional) Whether `output` is a tensor of logits or
            probabilities.
            Set it to `True` if `output` represents logits; otherwise,
            set it to `False` if `output` represents probabilities.
            Defaults to`False`.

    Returns:
        Integer tensor: The computed binary cross-entropy loss between
        `target` and `output`.

    Example:

    >>> target = keras.ops.convert_to_tensor([0, 1, 1, 0])
    >>> output = keras.ops.convert_to_tensor([0.1, 0.9, 0.8, 0.2])
    >>> binary_crossentropy(target, output)
    array([0.10536054 0.10536054 0.22314355 0.22314355],
          shape=(4,), dtype=float32)
    """
    if any_symbolic_tensors((target, output)):
        return BinaryCrossentropy(from_logits=from_logits).symbolic_call(
            target, output
        )
    return backend.nn.binary_crossentropy(
        target, output, from_logits=from_logits
    )


class CategoricalCrossentropy(Operation):
    def __init__(self, from_logits=False, axis=-1):
        super().__init__()
        self.from_logits = from_logits
        self.axis = axis

    def call(self, target, output):
        return backend.nn.categorical_crossentropy(
            target, output, from_logits=self.from_logits, axis=self.axis
        )

    def compute_output_spec(self, target, output):
        if target.shape != output.shape:
            raise ValueError(
                "Arguments `target` and `output` must have the same shape. "
                "Received: "
                f"target.shape={target.shape}, output.shape={output.shape}"
            )
        if len(target.shape) < 1:
            raise ValueError(
                "Arguments `target` and `output` must be at least rank 1. "
                "Received: "
                f"target.shape={target.shape}, output.shape={output.shape}"
            )
        return KerasTensor(output.shape[:-1], dtype=output.dtype)


@keras_export(
    [
        "keras.ops.categorical_crossentropy",
        "keras.ops.nn.categorical_crossentropy",
    ]
)
def categorical_crossentropy(target, output, from_logits=False, axis=-1):
    """Computes categorical cross-entropy loss between target and output tensor.

    The categorical cross-entropy loss is commonly used in multi-class
    classification tasks where each input sample can belong to one of
    multiple classes. It measures the dissimilarity
    between the target and output probabilities or logits.

    Args:
        target: The target tensor representing the true categorical labels.
            Its shape should match the shape of the `output` tensor
            except for the last dimension.
        output: The output tensor representing the predicted probabilities
            or logits. Its shape should match the shape of the `target`
            tensor except for the last dimension.
        from_logits: (optional) Whether `output` is a tensor of logits or
            probabilities.
            Set it to `True` if `output` represents logits; otherwise,
            set it to `False` if `output` represents probabilities.
            Defaults to`False`.
        axis: (optional) The axis along which the categorical cross-entropy
            is computed.
            Defaults to `-1`, which corresponds to the last dimension of
            the tensors.

    Returns:
        Integer tensor: The computed categorical cross-entropy loss between
        `target` and `output`.

    Example:

    >>> target = keras.ops.convert_to_tensor(
    ... [[1, 0, 0],
    ...  [0, 1, 0],
    ...  [0, 0, 1]])
    >>> output = keras.ops.convert_to_tensor(
    ... [[0.9, 0.05, 0.05],
    ...  [0.1, 0.8, 0.1],
    ...  [0.2, 0.3, 0.5]])
    >>> categorical_crossentropy(target, output)
    array([0.10536054 0.22314355 0.6931472 ], shape=(3,), dtype=float32)
    """
    if any_symbolic_tensors((target, output)):
        return CategoricalCrossentropy(
            from_logits=from_logits, axis=axis
        ).symbolic_call(target, output)
    return backend.nn.categorical_crossentropy(
        target, output, from_logits=from_logits, axis=axis
    )


class SparseCategoricalCrossentropy(Operation):
    def __init__(self, from_logits=False, axis=-1):
        super().__init__()
        self.from_logits = from_logits
        self.axis = axis

    def call(self, target, output):
        return backend.nn.sparse_categorical_crossentropy(
            target, output, from_logits=self.from_logits, axis=self.axis
        )

    def compute_output_spec(self, target, output):
        if len(output.shape) < 1:
            raise ValueError(
                "Argument `output` must be at least rank 1. "
                "Received: "
                f"output.shape={output.shape}"
            )
        target_shape = target.shape
        if len(target_shape) == len(output.shape) and target_shape[-1] == 1:
            target_shape = target_shape[:-1]
        if target_shape != output.shape[:-1]:
            raise ValueError(
                "Arguments `target` and `output` must have the same shape "
                "up until the last dimension: "
                f"target.shape={target.shape}, output.shape={output.shape}"
            )
        return KerasTensor(output.shape[:-1], dtype=output.dtype)


@keras_export(
    [
        "keras.ops.sparse_categorical_crossentropy",
        "keras.ops.nn.sparse_categorical_crossentropy",
    ]
)
def sparse_categorical_crossentropy(target, output, from_logits=False, axis=-1):
    """Computes sparse categorical cross-entropy loss.

    The sparse categorical cross-entropy loss is similar to categorical
    cross-entropy, but it is used when the target tensor contains integer
    class labels instead of one-hot encoded vectors. It measures the
    dissimilarity between the target and output probabilities or logits.

    Args:
        target: The target tensor representing the true class labels as
            integers. Its shape should match the shape of the `output`
            tensor except for the last dimension.
        output: The output tensor representing the predicted probabilities
            or logits.
            Its shape should match the shape of the `target` tensor except
            for the last dimension.
        from_logits: (optional) Whether `output` is a tensor of logits
            or probabilities.
            Set it to `True` if `output` represents logits; otherwise,
            set it to `False` if `output` represents probabilities.
            Defaults to`False`.
        axis: (optional) The axis along which the sparse categorical
            cross-entropy is computed.
            Defaults to `-1`, which corresponds to the last dimension
            of the tensors.

    Returns:
        Integer tensor: The computed sparse categorical cross-entropy
        loss between `target` and `output`.

    Example:

    >>> target = keras.ops.convert_to_tensor([0, 1, 2], dtype=int32)
    >>> output = keras.ops.convert_to_tensor(
    ... [[0.9, 0.05, 0.05],
    ...  [0.1, 0.8, 0.1],
    ...  [0.2, 0.3, 0.5]])
    >>> sparse_categorical_crossentropy(target, output)
    array([0.10536056 0.22314355 0.6931472 ], shape=(3,), dtype=float32)
    """
    if any_symbolic_tensors((target, output)):
        return SparseCategoricalCrossentropy(
            from_logits=from_logits, axis=axis
        ).symbolic_call(target, output)
    return backend.nn.sparse_categorical_crossentropy(
        target, output, from_logits=from_logits, axis=axis
    )


class MultiHot(Operation):
    def __init__(
        self, num_classes=None, axis=-1, dtype=None, sparse=False, **kwargs
    ):
        if num_classes is None and "num_tokens" in kwargs:
            num_classes = kwargs.pop("num_tokens")
        if num_classes is None:
            raise ValueError("Argument `num_classes` must be specified.")
        super().__init__(**kwargs)
        self.num_classes = num_classes
        self.axis = axis
        self.dtype = dtype or backend.floatx()
        self.sparse = sparse

    def call(self, inputs):
        return backend.nn.multi_hot(
            inputs,
            num_classes=self.num_classes,
            axis=self.axis,
            dtype=self.dtype,
        )

    def compute_output_spec(self, inputs):
        x_shape = list(getattr(inputs, "shape", []))
        if self.axis == -1:
            x_shape.append(self.num_classes)
        elif self.axis >= 0 and self.axis < len(x_shape):
            x_shape.insert(self.axis, self.num_classes)
        else:
            raise ValueError(
                f"axis must be -1 or between [0, {len(inputs.shape)}), but "
                f"received {self.axis}."
            )

        if len(x_shape) == 2:
            x_shape = [x_shape[-1]]
        else:
            x_shape = [x_shape[0]] + x_shape[2:]

        return KerasTensor(x_shape, dtype=inputs.dtype, sparse=self.sparse)


@keras_export(
    [
        "keras.ops.multi_hot",
        "keras.ops.nn.multi_hot",
    ]
)
def multi_hot(
    inputs, num_classes=None, axis=-1, dtype=None, sparse=False, **kwargs
):
    """Encodes integer labels as multi-hot vectors.

    This function encodes integer labels as multi-hot vectors, where each label
    is mapped to a binary value in the resulting vector.

    Args:
        inputs: Tensor of integer labels to be converted to multi-hot vectors.
        num_classes: Integer, the total number of unique classes.
        axis: (optional) Axis along which the multi-hot encoding should be
            added. Defaults to `-1`, which corresponds to the last dimension.
        dtype: (optional) The data type of the resulting tensor. Default
            is backend's float type.
        sparse: Whether to return a sparse tensor; for backends that support
            sparse tensors.

    Returns:
        Tensor: The multi-hot encoded tensor.

    Example:

    >>> data = keras.ops.convert_to_tensor([0, 4])
    >>> keras.ops.multi_hot(data, num_classes=5)
    array([1.0, 0.0, 0.0, 0.0, 1.0], dtype=float32)

    """
    if num_classes is None and "num_tokens" in kwargs:
        num_classes = kwargs.pop("num_tokens")
    if num_classes is None:
        raise ValueError("Argument `num_classes` must be specified.")

    if any_symbolic_tensors((inputs,)):
        return MultiHot(num_classes, axis, dtype, sparse).symbolic_call(inputs)

    return backend.nn.multi_hot(inputs, num_classes, axis, dtype, sparse)


class Moments(Operation):
    def __init__(self, axes, keepdims=False, synchronized=False):
        super().__init__()
        self.axes = axes
        self.keepdims = keepdims
        self.synchronized = synchronized

    def call(self, x):
        return backend.nn.moments(
            x,
            axes=self.axes,
            keepdims=self.keepdims,
            synchronized=self.synchronized,
        )

    def compute_output_spec(self, x):
        return (
            KerasTensor(
                reduce_shape(x.shape, axis=self.axes, keepdims=self.keepdims),
                dtype=x.dtype,
            ),
            KerasTensor(
                reduce_shape(x.shape, axis=self.axes, keepdims=self.keepdims),
                dtype=x.dtype,
            ),
        )


@keras_export(
    [
        "keras.ops.moments",
        "keras.ops.nn.moments",
    ]
)
def moments(x, axes, keepdims=False, synchronized=False):
    """Calculates the mean and variance of `x`.

    The mean and variance are calculated by aggregating the contents of `x`
    across `axes`. If `x` is 1-D and `axes = [0]` this is just the mean and
    variance of a vector.

    Args:
        x: Input tensor.
        axes: A list of axes which to compute mean and variance.
        keepdims: If this is set to `True`, the axes which are reduced are left
            in the result as dimensions with size one.
        synchronized: Only applicable with the TensorFlow backend.
            If `True`, synchronizes the global batch statistics (mean and
            variance) across all devices at each training step in a
            distributed training strategy. If `False`, each replica uses its own
            local batch statistics.

    Returns:
        A tuple containing two tensors - mean and variance.

    Example:

    >>> x = keras.ops.convert_to_tensor([0, 1, 2, 3, 100], dtype="float32")
    >>> keras.ops.moments(x, axes=[0])
    (array(21.2, dtype=float32), array(1553.3601, dtype=float32))

    """
    if any_symbolic_tensors((x,)):
        return Moments(axes, keepdims, synchronized=synchronized).symbolic_call(
            x
        )

    return backend.nn.moments(x, axes, keepdims, synchronized=synchronized)


class BatchNorm(Operation):
    def __init__(self, axis, epsilon):
        super().__init__()
        self.axis = axis
        self.epsilon = epsilon

    def _check_shape(self, name, shape, expected_shape):
        if shape != expected_shape:
            raise ValueError(
                f"Arguments `{name}` must be a vector of length "
                f"`x.shape[axis]`. Expected: `{expected_shape}`. "
                f"Received: `{shape}."
            )

    def compute_output_spec(self, x, mean, variance, offset, scale):
        shape = (x.shape[self.axis],)
        self._check_shape("mean", tuple(mean.shape), shape)
        self._check_shape("variance", tuple(variance.shape), shape)
        if offset is not None:
            self._check_shape("offset", tuple(offset.shape), shape)
        if offset is not scale:
            self._check_shape("scale", tuple(scale.shape), shape)
        return KerasTensor(x.shape, dtype=x.dtype)


@keras_export(
    [
        "keras.ops.batch_normalization",
        "keras.ops.nn.batch_normalization",
    ]
)
def batch_normalization(
    x, mean, variance, axis, offset=None, scale=None, epsilon=1e-3
):
    """Normalizes `x` by `mean` and `variance`.

    This op is typically used by the batch normalization step in a neural
    network. It normalizes the input tensor along the given axis.

    Args:
        x: Input tensor.
        mean: A mean vector of the same length as the `axis` dimension of the
            input thensor.
        variance: A variance vector of the same length as the `axis` dimension
            of the input tensor.
        axis: Integer, the axis that should be normalized.
        offset: An offset vector of the same length as the `axis` dimension of
            the input tensor. If not `None`, `offset` is added to the normalized
            tensor. Defaults to `None`.
        scale: A scale vector of the same length as the `axis` dimension of the
            input tensor. If not `None`, the normalized tensor is multiplied by
            `scale`. Defaults to `None`.
        epsilon: Small float added to variance to avoid dividing by zero.
            Defaults to 1e-3.

    Returns:
        The normalized tensor.

    Example:

    >>> x = keras.ops.convert_to_tensor(
    ...     [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6], [0.7, 0.8, 0.9]]
    ... )
    >>> keras.ops.batch_normalization(
    ...     x,
    ...     mean=[0.4, 0.5, 0.6],
    ...     variance=[0.67, 0.67, 0.67],
    ...     axis=-1
    ... )
    array([[-3.6624e-01, -3.6624e-01, -3.6624e-01],
           [-4.6445e-09,  0.0000e+00, -1.8578e-08],
           [ 3.6624e-01,  3.6624e-01,  3.6624e-01]])

    """
    if any_symbolic_tensors((x, mean, variance, offset, scale)):
        return BatchNorm(axis, epsilon).symbolic_call(
            x, mean, variance, offset, scale
        )

    return backend.nn.batch_normalization(
        x, mean, variance, axis, offset, scale, epsilon
    )


class CTCLoss(Operation):
    def __init__(self, mask_index=0):
        super().__init__()
        self.mask_index = mask_index

    def call(self, target, output, target_length, output_length):
        return backend.nn.ctc_loss(
            target, output, target_length, output_length, self.mask_index
        )

    def _check_shape_first_dim(self, name1, shape1, name2, shape2):
        if shape1[0] != shape2[0]:
            raise ValueError(
                f"Arguments `{name1}` and `{name2}` must have the same "
                "first dimension. "
                f"Received shapes: `{shape1}` and `{shape2}`."
            )

    def compute_output_spec(self, target, output, target_length, output_length):
        self._check_shape_first_dim(
            "target", target.shape, "output", output.shape
        )
        self._check_shape_first_dim(
            "target_length", target_length.shape, "target", target.shape
        )
        self._check_shape_first_dim(
            "output_length", output_length.shape, "output", output.shape
        )
        dtype = backend.result_type(output.dtype, "float32")
        return KerasTensor((target.shape[0],), dtype=dtype)


@keras_export(
    [
        "keras.ops.ctc_loss",
        "keras.ops.nn.ctc_loss",
    ]
)
def ctc_loss(target, output, target_length, output_length, mask_index=0):
    """CTC (Connectionist Temporal Classification) loss.

    Args:
        target: A tensor of shape `(batch_size, max_length)` containing
            the true labels in integer format.
        output: A tensor of shape `(batch_size, max_length, num_classes)`
            containing logits (the output of your model).
        target_length: A tensor of shape `(batch_size,)` containing the
            true label lengths.
        output_length: A tensor of shape `(batch_size,)` containing the
            output lengths.
        mask_index: The index of the mask character in the vocabulary.
            Defaults to `0`.
    """

    if any_symbolic_tensors((target, output, target_length, output_length)):
        return CTCLoss(mask_index).symbolic_call(
            target, output, target_length, output_length
        )
    return backend.nn.ctc_loss(
        target, output, target_length, output_length, mask_index
    )


class CTCDecode(Operation):
    def __init__(
        self,
        strategy="greedy",
        beam_width=100,
        top_paths=1,
        merge_repeated=True,
        mask_index=None,
    ):
        super().__init__()
        self.strategy = strategy
        self.beam_width = beam_width
        self.top_paths = top_paths
        self.merge_repeated = merge_repeated
        self.mask_index = mask_index

    def call(self, inputs, sequence_lengths):
        return backend.nn.ctc_decode(
            inputs,
            sequence_lengths,
            strategy=self.strategy,
            beam_width=self.beam_width,
            top_paths=self.top_paths,
            merge_repeated=self.merge_repeated,
            mask_index=self.mask_index,
        )

    def compute_output_spec(self, inputs, sequence_lengths):
        inputs_shape = inputs.shape
        if self.strategy == "greedy":
            top_paths = 1
        else:
            top_paths = self.top_paths
        dtype = backend.result_type(inputs.dtype, "float32")
        return (
            KerasTensor(
                (top_paths, inputs_shape[0], inputs_shape[1]), dtype="int32"
            ),
            KerasTensor((inputs_shape[0], top_paths), dtype=dtype),
        )


@keras_export(
    [
        "keras.ops.ctc_decode",
        "keras.ops.nn.ctc_decode",
    ]
)
def ctc_decode(
    inputs,
    sequence_lengths,
    strategy="greedy",
    beam_width=100,
    top_paths=1,
    merge_repeated=True,
    mask_index=None,
):
    """Decodes the output of a CTC model.

    Args:
        inputs: A tensor of shape `(batch_size, max_length, num_classes)`
            containing the logits (the output of the model).
            They should *not* be normalized via softmax.
        sequence_lengths: A tensor of shape `(batch_size,)` containing the
            sequence lengths for the batch.
        strategy: A string for the decoding strategy. Supported values are
            `"greedy"` and `"beam_search"`.
        beam_width: An integer scalar beam width used in beam search.
            Defaults to 100.
        top_paths: An integer scalar, the number of top paths to return.
            Defaults to 1.
        merge_repeated: A boolean scalar, whether to merge repeated
            labels in the output. Defaults to `True`.
        mask_index: An integer scalar, the index of the mask character in
            the vocabulary. Defaults to `None`.

    Returns:
        A tuple containing:
        - The tensor representing the list of decoded sequences. If
            `strategy="greedy"`, the shape is `(1, batch_size, max_length)`. If
            `strategy="beam_seatch"`, the shape is
            `(top_paths, batch_size, max_length)`. Note that: `-1` indicates the
            blank label.
        - If `strategy="greedy"`, a tensor of shape `(batch_size, 1)`
            representing the negative of the sum of the probability logits for
            each sequence. If `strategy="beam_seatch"`, a tensor of shape
            `(batch_size, top_paths)` representing the log probability for each
            sequence.
    """

    if any_symbolic_tensors((inputs, sequence_lengths)):
        return CTCDecode(
            strategy=strategy,
            beam_width=beam_width,
            top_paths=top_paths,
            merge_repeated=merge_repeated,
            mask_index=mask_index,
        ).symbolic_call(inputs, sequence_lengths)
    return backend.nn.ctc_decode(
        inputs=inputs,
        sequence_length=sequence_lengths,
        strategy=strategy,
        beam_width=beam_width,
        top_paths=top_paths,
        merge_repeated=merge_repeated,
        mask_index=mask_index,
    )


class Normalize(Operation):
    def __init__(self, axis=-1, order=2):
        super().__init__()
        self.axis = axis
        self.order = order

    def compute_output_spec(self, x):
        return KerasTensor(shape=x.shape)

    def call(self, x):
        return _normalize(x, axis=self.axis, order=self.order)


@keras_export(
    [
        "keras.ops.normalize",
        "keras.ops.nn.normalize",
    ]
)
def normalize(x, axis=-1, order=2):
    """Normalizes `x` over the specified axis.

    It is defined as: `normalize(x) = x / max(norm(x), epsilon)`.

    Args:
        x: Input tensor.
        axis: The axis or axes along which to perform normalization.
            Default to -1.
        order: The exponent value in the norm formulation.
            Defaults to 2.

    Returns:
        The normalized array.

    Example:

    >>> x = keras.ops.convert_to_tensor([[1, 2, 3], [4, 5, 6]])
    >>> x_norm = keras.ops.math.normalize(x)
    >>> print(x_norm)
    array([[0.26726124 0.5345225  0.8017837 ]
           [0.45584232 0.5698029  0.68376344]], shape=(2, 3), dtype=float32)

    """
    if any_symbolic_tensors((x,)):
        return Normalize(axis=axis, order=order).symbolic_call(x)
    return _normalize(x, axis=axis, order=order)


def _normalize(x, axis=-1, order=2):
    if not isinstance(order, int) or not order >= 1:
        raise ValueError(
            f"Argument `order` must be an int >= 1. Received: order={order}"
        )
    x = backend.convert_to_tensor(x)
    if len(x.shape) == 0:
        x = backend.numpy.expand_dims(x, axis=0)
    epsilon = backend.epsilon()
    norm = backend.linalg.norm(x, ord=order, axis=axis, keepdims=True)
    denom = backend.numpy.maximum(norm, epsilon)
    return backend.numpy.divide(x, denom)


class PSNR(Operation):
    def __init__(
        self,
        max_val,
    ):
        super().__init__()
        self.max_val = max_val

    def call(self, x1, x2):
        return backend.nn.psnr(
            x1=x1,
            x2=x2,
            max_val=self.max_val,
        )

    def compute_output_spec(self, x1, x2):
        if len(x1.shape) != len(x2.shape):
            raise ValueError("Inputs must have the same rank")

        return KerasTensor(shape=())


@keras_export(
    [
        "keras.ops.psnr",
        "keras.ops.nn.psnr",
    ]
)
def psnr(
    x1,
    x2,
    max_val,
):
    """Peak Signal-to-Noise Ratio (PSNR) function.

    This function computes the Peak Signal-to-Noise Ratio between two signals,
    `x1` and `x2`. PSNR is a measure of the quality of a reconstructed signal.
    The higher the PSNR, the closer the reconstructed signal is to the original
    signal. Note that it can become negative when the signal power is
    smaller that the noise power.

    Args:
        x1: The first input signal.
        x2: The second input signal. Must have the same shape as `x1`.
        max_val: The maximum possible value in the signals.

    Returns:
        float: The PSNR value between `x1` and `x2`.

    Examples:

    >>> x1 = keras.random.normal((2, 4, 4, 3))
    >>> x2 = keras.random.normal((2, 4, 4, 3))
    >>> max_val = 1.0
    >>> keras.ops.nn.psnr(x1, x2, max_val)
    -3.1697404
    """
    if any_symbolic_tensors(
        (
            x1,
            x2,
        )
    ):
        return PSNR(
            max_val,
        ).symbolic_call(x1, x2)
    return backend.nn.psnr(
        x1,
        x2,
        max_val,
    )