helper function that converts an operand to a spatial operand helper function to define pooling functions args inputs input data of shape n2 initialvalue the initial value for the reduction reducefn a reduce function of the form t t t poolsize a sequence of n integers representing the window size to reduce over strides a sequence of n integers representing the interwindow strides default 1 1 padding either the string same or valid returns the output of the reduction for each window slice avoid the extra reducewindow count the number of valid entries at each input point then use that for computing average assumes that any two arrays of same shape will be padded the same avoid broadcasting on axis where pooling is skipped create a lax convdimensionnumbers for the given inputs numdims numspatialdims 2 if dataformat channelslast spatialdims tuplerange1 numdims 1 inputsdn 0 numdims 1 spatialdims else spatialdims tuplerange2 numdims inputsdn 0 1 spatialdims if transpose kerneldn numdims 2 numdims 1 tuplerangenumdims 2 else kerneldn numdims 1 numdims 2 tuplerangenumdims 2 return lax convdimensionnumbers lhsspecinputsdn rhsspeckerneldn outspecinputsdn def conv inputs kernel strides1 paddingvalid dataformatnone dilationrate1 dataformat standardizedataformatdataformat numspatialdims inputs ndim 2 dimensionnumbers converttolaxconvdimensionnumbers numspatialdims dataformat transposefalse strides converttospatialoperand strides numspatialdims dataformat includebatchandchannelsfalse dilationrate converttospatialoperand dilationrate numspatialdims dataformat includebatchandchannelsfalse if dataformat channelslast channels inputs shape1 else channels inputs shape1 kernelinchannels kernel shape2 if channels kernelinchannels 0 raise valueerror the number of input channels must be evenly divisible by fkernel s inchannels received input channels channels and fkernel inchannels kernelinchannels featuregroupcount channels kernelinchannels return jax lax convgeneraldilated converttotensorinputs converttotensorkernel strides padding rhsdilationdilationrate dimensionnumbersdimensionnumbers featuregroupcountfeaturegroupcount def depthwiseconv inputs kernel strides1 paddingvalid dataformatnone dilationrate1 dataformat standardizedataformatdataformat numspatialdims inputs ndim 2 dimensionnumbers converttolaxconvdimensionnumbers numspatialdims dataformat transposefalse strides converttospatialoperand strides numspatialdims dataformat includebatchandchannelsfalse dilationrate converttospatialoperand dilationrate numspatialdims dataformat includebatchandchannelsfalse featuregroupcount inputs shape1 if dataformat channelslast else inputs shape1 kernel jnp reshape kernel kernel shape 2 1 featuregroupcount kernel shape1 return jax lax convgeneraldilated inputs kernel strides padding rhsdilationdilationrate dimensionnumbersdimensionnumbers featuregroupcountfeaturegroupcount def separableconv inputs depthwisekernel pointwisekernel strides1 paddingvalid dataformatnone dilationrate1 dataformat standardizedataformatdataformat depthwiseconvoutput depthwiseconv inputs depthwisekernel strides padding dataformat dilationrate return conv depthwiseconvoutput pointwisekernel strides1 paddingvalid dataformatdataformat dilationratedilationrate def convtranspose inputs kernel strides1 paddingvalid outputpaddingnone dataformatnone dilationrate1 dataformat standardizedataformatdataformat numspatialdims inputs ndim 2 paddingvalues computeconvtransposepaddingargsforjax inputshapeinputs shape kernelshapekernel shape stridesstrides paddingpadding outputpaddingoutputpadding dilationratedilationrate dimensionnumbers converttolaxconvdimensionnumbers numspatialdims dataformat transposefalse strides converttospatialoperand strides numspatialdims dataformat includebatchandchannelsfalse dilationrate converttospatialoperand dilationrate numspatialdims dataformat includebatchandchannelsfalse return jax lax convtranspose inputs kernel strides paddingpaddingvalues rhsdilationdilationrate dimensionnumbersdimensionnumbers transposekerneltrue def onehotx numclasses axis1 dtypefloat32 x converttotensorx return jnn onehotx numclasses axisaxis dtypedtype def multihotx numclasses axis1 dtypefloat32 reductionaxis 1 if lenx shape 1 else 0 outputs jnp max onehotcastx int32 numclasses axisaxis dtypedtype axisreductionaxis return outputs def categoricalcrossentropytarget output fromlogitsfalse axis1 target jnp arraytarget output jnp arrayoutput if target shape output shape raise valueerror arguments target and output must have the same shape received ftarget shapetarget shape output shapeoutput shape if lentarget shape 1 raise valueerror arguments target and output must be at least rank 1 received ftarget shapetarget shape output shapeoutput shape if fromlogits logprob jax nn logsoftmaxoutput axisaxis else output output jnp sumoutput axis keepdimstrue output jnp clipoutput epsilon 1 0 epsilon logprob jnp logoutput return jnp sumtarget logprob axisaxis def sparsecategoricalcrossentropytarget output fromlogitsfalse axis1 target jnp arraytarget dtypeint32 output jnp arrayoutput if lentarget shape lenoutput shape and target shape1 1 target jnp squeezetarget axis1 if lenoutput shape 1 raise valueerror argument output must be at least rank 1 received foutput shapeoutput shape if target shape output shape 1 raise valueerror arguments target and output must have the same shape up until the last dimension ftarget shapetarget shape output shapeoutput shape if fromlogits logprob jax nn logsoftmaxoutput axisaxis else output output jnp sumoutput axis keepdimstrue output jnp clipoutput epsilon 1 0 epsilon logprob jnp logoutput target jnn onehottarget output shapeaxis axisaxis return jnp sumtarget logprob axisaxis def binarycrossentropytarget output fromlogitsfalse target jnp arraytarget output jnp arrayoutput if target shape output shape raise valueerror arguments target and output must have the same shape received ftarget shapetarget shape output shapeoutput shape if fromlogits loglogits jax nn logsigmoidoutput logneglogits jax nn logsigmoidoutput return 1 0 target loglogits 1 0 target logneglogits output jnp clipoutput epsilon 1 0 epsilon bce target jnp logoutput bce 1 0 target jnp log1 0 output return bce def momentsx axes keepdimsfalse synchronizedfalse if synchronized raise notimplementederror argument synchronizedtrue is not supported with jax the dynamic range of float16 is too limited for statistics as a workaround we simply perform the operations on float32 and convert back to float16 needcast false oridtype standardizedtypex dtype if oridtype float16 needcast true x castx float32 mean jnp meanx axes keepdimstrue the variance is computed using var ex2 ex2 it is faster but less numerically stable note stopgradient does not change the gradient to the mean because that gradient is zero the substraction operation does not guarantee a nonnegative result given float precision so we clamp it to 0 variance jnp maximum jnp meanjnp squarex axisaxes keepdimstrue jnp squarejax lax stopgradientmean 0 0 if not keepdims mean jnp squeezemean axes variance jnp squeezevariance axes if needcast avoid overflow and underflow when casting from float16 to float32 mean jnp clip mean jnp finfojnp float16 min jnp finfojnp float16 max variance jnp clip variance jnp finfojnp float16 min jnp finfojnp float16 max mean castmean oridtype variance castvariance oridtype return mean variance def batchnormalization x mean variance axis offsetnone scalenone epsilon1e3 shape 1 lenx shape shapeaxis mean shape0 mean jnp reshapemean shape variance jnp reshapevariance shape inv jax lax rsqrtvariance epsilon if scale is not none scale jnp reshapescale shape inv inv scale res mean inv if offset is not none offset jnp reshapeoffset shape res res offset return x inv res helper function that converts an operand to a spatial operand helper function to define pooling functions args inputs input data of shape n 2 initial_value the initial value for the reduction reduce_fn a reduce function of the form t t t pool_size a sequence of n integers representing the window size to reduce over strides a sequence of n integers representing the inter window strides default 1 1 padding either the string same or valid returns the output of the reduction for each window slice avoid the extra reduce_window count the number of valid entries at each input point then use that for computing average assumes that any two arrays of same shape will be padded the same avoid broadcasting on axis where pooling is skipped create a lax convdimensionnumbers for the given inputs the dynamic range of float16 is too limited for statistics as a workaround we simply perform the operations on float32 and convert back to float16 the variance is computed using var e x 2 e x 2 it is faster but less numerically stable note stop_gradient does not change the gradient to the mean because that gradient is zero the substraction operation does not guarantee a non negative result given float precision so we clamp it to 0 avoid overflow and underflow when casting from float16 to float32
import jax
import jax.numpy as jnp
import numpy as np
from jax import lax
from jax import nn as jnn

from keras.backend import standardize_data_format
from keras.backend import standardize_dtype
from keras.backend.common.backend_utils import (
    compute_conv_transpose_padding_args_for_jax,
)
from keras.backend.config import epsilon
from keras.backend.jax.core import cast
from keras.backend.jax.core import convert_to_tensor


def relu(x):
    x = convert_to_tensor(x)
    return jnn.relu(x)


def relu6(x):
    x = convert_to_tensor(x)
    return jnn.relu6(x)


def sigmoid(x):
    x = convert_to_tensor(x)
    return jnn.sigmoid(x)


def tanh(x):
    x = convert_to_tensor(x)
    return jnn.tanh(x)


def softplus(x):
    x = convert_to_tensor(x)
    return jnn.softplus(x)


def softsign(x):
    x = convert_to_tensor(x)
    return jnn.soft_sign(x)


def silu(x):
    x = convert_to_tensor(x)
    return jnn.silu(x)


def log_sigmoid(x):
    x = convert_to_tensor(x)
    return jnn.log_sigmoid(x)


def leaky_relu(x, negative_slope=0.2):
    x = convert_to_tensor(x)
    return jnn.leaky_relu(x, negative_slope=negative_slope)


def hard_sigmoid(x):
    x = convert_to_tensor(x)
    return jnn.hard_sigmoid(x)


def hard_swish(x):
    x = convert_to_tensor(x)
    return jnn.hard_swish(x)


def elu(x, alpha=1.0):
    x = convert_to_tensor(x)
    return jnn.elu(x, alpha=alpha)


def selu(x):
    x = convert_to_tensor(x)
    return jnn.selu(x)


def gelu(x, approximate=True):
    x = convert_to_tensor(x)
    return jnn.gelu(x, approximate)


def softmax(x, axis=-1):
    x = convert_to_tensor(x)
    return jnn.softmax(x, axis=axis)


def log_softmax(x, axis=-1):
    x = convert_to_tensor(x)
    return jnn.log_softmax(x, axis=axis)


def _convert_to_spatial_operand(
    x,
    num_spatial_dims,
    data_format="channels_last",
    include_batch_and_channels=True,
):
    
    x = (x,) * num_spatial_dims if isinstance(x, int) else x
    if not include_batch_and_channels:
        return x
    if data_format == "channels_last":
        x = (1,) + x + (1,)
    else:
        x = (1,) + (1,) + x
    return x


def _pool(
    inputs,
    initial_value,
    reduce_fn,
    pool_size,
    strides=None,
    padding="valid",
):
    
    if padding not in ("same", "valid"):
        raise ValueError(
            f"Invalid padding '{padding}', must be 'same' or 'valid'."
        )
    padding = padding.upper()
    return lax.reduce_window(
        inputs,
        initial_value,
        reduce_fn,
        pool_size,
        strides,
        padding,
    )


def max_pool(
    inputs,
    pool_size,
    strides=None,
    padding="valid",
    data_format=None,
):
    data_format = standardize_data_format(data_format)
    num_spatial_dims = inputs.ndim - 2
    pool_size = _convert_to_spatial_operand(
        pool_size, num_spatial_dims, data_format
    )
    strides = pool_size if strides is None else strides
    strides = _convert_to_spatial_operand(
        strides, num_spatial_dims, data_format
    )
    return _pool(inputs, -jnp.inf, lax.max, pool_size, strides, padding)


def average_pool(
    inputs,
    pool_size,
    strides,
    padding,
    data_format=None,
):
    data_format = standardize_data_format(data_format)
    num_spatial_dims = inputs.ndim - 2
    pool_size = _convert_to_spatial_operand(
        pool_size, num_spatial_dims, data_format
    )
    strides = pool_size if strides is None else strides
    strides = _convert_to_spatial_operand(
        strides, num_spatial_dims, data_format
    )

    pooled = _pool(inputs, 0.0, lax.add, pool_size, strides, padding)
    if padding == "valid":
        
        return pooled / np.prod(pool_size)
    else:
        
        
        
        
        shape = [
            (a if b != 1 else 1) for (a, b) in zip(inputs.shape, pool_size)
        ]
        window_counts = _pool(
            jnp.ones(shape, inputs.dtype),
            0.0,
            lax.add,
            pool_size,
            strides,
            padding,
        )
        return pooled / window_counts


def _convert_to_lax_conv_dimension_numbers(
    num_spatial_dims,
    data_format="channels_last",
    transpose=False,
):
    
    num_dims = num_spatial_dims + 2

    if data_format == "channels_last":
        spatial_dims = tuple(range(1, num_dims - 1))
        inputs_dn = (0, num_dims - 1) + spatial_dims
    else:
        spatial_dims = tuple(range(2, num_dims))
        inputs_dn = (0, 1) + spatial_dims

    if transpose:
        kernel_dn = (num_dims - 2, num_dims - 1) + tuple(range(num_dims - 2))
    else:
        kernel_dn = (num_dims - 1, num_dims - 2) + tuple(range(num_dims - 2))

    return lax.ConvDimensionNumbers(
        lhs_spec=inputs_dn, rhs_spec=kernel_dn, out_spec=inputs_dn
    )


def conv(
    inputs,
    kernel,
    strides=1,
    padding="valid",
    data_format=None,
    dilation_rate=1,
):
    data_format = standardize_data_format(data_format)
    num_spatial_dims = inputs.ndim - 2
    dimension_numbers = _convert_to_lax_conv_dimension_numbers(
        num_spatial_dims,
        data_format,
        transpose=False,
    )
    strides = _convert_to_spatial_operand(
        strides,
        num_spatial_dims,
        data_format,
        include_batch_and_channels=False,
    )
    dilation_rate = _convert_to_spatial_operand(
        dilation_rate,
        num_spatial_dims,
        data_format,
        include_batch_and_channels=False,
    )
    if data_format == "channels_last":
        channels = inputs.shape[-1]
    else:
        channels = inputs.shape[1]
    kernel_in_channels = kernel.shape[-2]
    if channels % kernel_in_channels > 0:
        raise ValueError(
            "The number of input channels must be evenly divisible by "
            f"kernel's in_channels. Received input channels {channels} and "
            f"kernel in_channels {kernel_in_channels}. "
        )
    feature_group_count = channels // kernel_in_channels
    return jax.lax.conv_general_dilated(
        convert_to_tensor(inputs),
        convert_to_tensor(kernel),
        strides,
        padding,
        rhs_dilation=dilation_rate,
        dimension_numbers=dimension_numbers,
        feature_group_count=feature_group_count,
    )


def depthwise_conv(
    inputs,
    kernel,
    strides=1,
    padding="valid",
    data_format=None,
    dilation_rate=1,
):
    data_format = standardize_data_format(data_format)
    num_spatial_dims = inputs.ndim - 2
    dimension_numbers = _convert_to_lax_conv_dimension_numbers(
        num_spatial_dims,
        data_format,
        transpose=False,
    )
    strides = _convert_to_spatial_operand(
        strides,
        num_spatial_dims,
        data_format,
        include_batch_and_channels=False,
    )
    dilation_rate = _convert_to_spatial_operand(
        dilation_rate,
        num_spatial_dims,
        data_format,
        include_batch_and_channels=False,
    )
    feature_group_count = (
        inputs.shape[-1] if data_format == "channels_last" else inputs.shape[1]
    )
    kernel = jnp.reshape(
        kernel,
        kernel.shape[:-2] + (1, feature_group_count * kernel.shape[-1]),
    )
    return jax.lax.conv_general_dilated(
        inputs,
        kernel,
        strides,
        padding,
        rhs_dilation=dilation_rate,
        dimension_numbers=dimension_numbers,
        feature_group_count=feature_group_count,
    )


def separable_conv(
    inputs,
    depthwise_kernel,
    pointwise_kernel,
    strides=1,
    padding="valid",
    data_format=None,
    dilation_rate=1,
):
    data_format = standardize_data_format(data_format)
    depthwise_conv_output = depthwise_conv(
        inputs,
        depthwise_kernel,
        strides,
        padding,
        data_format,
        dilation_rate,
    )
    return conv(
        depthwise_conv_output,
        pointwise_kernel,
        strides=1,
        padding="valid",
        data_format=data_format,
        dilation_rate=dilation_rate,
    )


def conv_transpose(
    inputs,
    kernel,
    strides=1,
    padding="valid",
    output_padding=None,
    data_format=None,
    dilation_rate=1,
):
    data_format = standardize_data_format(data_format)
    num_spatial_dims = inputs.ndim - 2
    padding_values = compute_conv_transpose_padding_args_for_jax(
        input_shape=inputs.shape,
        kernel_shape=kernel.shape,
        strides=strides,
        padding=padding,
        output_padding=output_padding,
        dilation_rate=dilation_rate,
    )
    dimension_numbers = _convert_to_lax_conv_dimension_numbers(
        num_spatial_dims,
        data_format,
        transpose=False,
    )
    strides = _convert_to_spatial_operand(
        strides,
        num_spatial_dims,
        data_format,
        include_batch_and_channels=False,
    )
    dilation_rate = _convert_to_spatial_operand(
        dilation_rate,
        num_spatial_dims,
        data_format,
        include_batch_and_channels=False,
    )

    return jax.lax.conv_transpose(
        inputs,
        kernel,
        strides,
        padding=padding_values,
        rhs_dilation=dilation_rate,
        dimension_numbers=dimension_numbers,
        transpose_kernel=True,
    )


def one_hot(x, num_classes, axis=-1, dtype="float32"):
    x = convert_to_tensor(x)
    return jnn.one_hot(x, num_classes, axis=axis, dtype=dtype)


def multi_hot(x, num_classes, axis=-1, dtype="float32"):
    reduction_axis = 1 if len(x.shape) > 1 else 0
    outputs = jnp.max(
        one_hot(cast(x, "int32"), num_classes, axis=axis, dtype=dtype),
        axis=reduction_axis,
    )
    return outputs


def categorical_crossentropy(target, output, from_logits=False, axis=-1):
    target = jnp.array(target)
    output = jnp.array(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}"
        )

    if from_logits:
        log_prob = jax.nn.log_softmax(output, axis=axis)
    else:
        output = output / jnp.sum(output, axis, keepdims=True)
        output = jnp.clip(output, epsilon(), 1.0 - epsilon())
        log_prob = jnp.log(output)
    return -jnp.sum(target * log_prob, axis=axis)


def sparse_categorical_crossentropy(target, output, from_logits=False, axis=-1):
    target = jnp.array(target, dtype="int32")
    output = jnp.array(output)
    if len(target.shape) == len(output.shape) and target.shape[-1] == 1:
        target = jnp.squeeze(target, axis=-1)

    if len(output.shape) < 1:
        raise ValueError(
            "Argument `output` must be at least rank 1. "
            "Received: "
            f"output.shape={output.shape}"
        )
    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}"
        )
    if from_logits:
        log_prob = jax.nn.log_softmax(output, axis=axis)
    else:
        output = output / jnp.sum(output, axis, keepdims=True)
        output = jnp.clip(output, epsilon(), 1.0 - epsilon())
        log_prob = jnp.log(output)
    target = jnn.one_hot(target, output.shape[axis], axis=axis)
    return -jnp.sum(target * log_prob, axis=axis)


def binary_crossentropy(target, output, from_logits=False):
    target = jnp.array(target)
    output = jnp.array(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 from_logits:
        log_logits = jax.nn.log_sigmoid(output)
        log_neg_logits = jax.nn.log_sigmoid(-output)
        return -1.0 * target * log_logits - (1.0 - target) * log_neg_logits

    output = jnp.clip(output, epsilon(), 1.0 - epsilon())
    bce = target * jnp.log(output)
    bce += (1.0 - target) * jnp.log(1.0 - output)
    return -bce


def moments(x, axes, keepdims=False, synchronized=False):
    if synchronized:
        raise NotImplementedError(
            "Argument synchronized=True is not supported with JAX."
        )
    
    
    
    need_cast = False
    ori_dtype = standardize_dtype(x.dtype)
    if ori_dtype == "float16":
        need_cast = True
        x = cast(x, "float32")

    mean = jnp.mean(x, axes, keepdims=True)

    
    
    
    
    
    
    variance = jnp.maximum(
        jnp.mean(jnp.square(x), axis=axes, keepdims=True)
        - jnp.square(jax.lax.stop_gradient(mean)),
        0.0,
    )

    if not keepdims:
        mean = jnp.squeeze(mean, axes)
        variance = jnp.squeeze(variance, axes)
    if need_cast:
        
        mean = jnp.clip(
            mean, jnp.finfo(jnp.float16).min, jnp.finfo(jnp.float16).max
        )
        variance = jnp.clip(
            variance, jnp.finfo(jnp.float16).min, jnp.finfo(jnp.float16).max
        )
        mean = cast(mean, ori_dtype)
        variance = cast(variance, ori_dtype)
    return mean, variance


def batch_normalization(
    x, mean, variance, axis, offset=None, scale=None, epsilon=1e-3
):
    shape = [1] * len(x.shape)
    shape[axis] = mean.shape[0]
    mean = jnp.reshape(mean, shape)
    variance = jnp.reshape(variance, shape)

    inv = jax.lax.rsqrt(variance + epsilon)
    if scale is not None:
        scale = jnp.reshape(scale, shape)
        inv = inv * scale

    res = -mean * inv
    if offset is not None:
        offset = jnp.reshape(offset, shape)
        res = res + offset

    return x * inv + res