commonly used math operations not included in numpy from keras import backend from keras apiexport import kerasexport from keras backend import kerastensor from keras backend import anysymbolictensors from keras ops operation import operation from keras ops operationutils import reduceshape class segmentsumoperation def initself numsegmentsnone sortedfalse super init self numsegments numsegments self sorted sorted def computeoutputspecself data segmentids numsegments self numsegments outputshape numsegments tupledata shape1 return kerastensorshapeoutputshape dtypedata dtype def callself data segmentids return backend math segmentsum data segmentids numsegmentsself numsegments sortedself sorted kerasexportkeras ops segmentsum def segmentsumdata segmentids numsegmentsnone sortedfalse if anysymbolictensorsdata return segmentsumnumsegments sorted symboliccalldata segmentids return backend math segmentsum data segmentids numsegmentsnumsegments sortedsorted class segmentmaxoperation def initself numsegmentsnone sortedfalse super init self numsegments numsegments self sorted sorted def computeoutputspecself data segmentids numsegments self numsegments outputshape numsegments tupledata shape1 return kerastensorshapeoutputshape dtypedata dtype def callself data segmentids return backend math segmentmax data segmentids numsegmentsself numsegments sortedself sorted kerasexportkeras ops segmentmax def segmentmaxdata segmentids numsegmentsnone sortedfalse if anysymbolictensorsdata return segmentmaxnumsegments sorted symboliccalldata segmentids return backend math segmentmax data segmentids numsegmentsnumsegments sortedsorted class topkoperation def initself k sortedfalse super init self k k self sorted sorted def computeoutputspecself x outputshape listx shape outputshape1 self k return a tuple values indices return kerastensorshapeoutputshape dtypex dtype kerastensorshapeoutputshape dtypeint32 def callself x return backend math topkx self k self sorted kerasexportkeras ops topk def topkx k sortedtrue if anysymbolictensorsx return topkk sorted symboliccallx return backend math topkx k sorted class intopkoperation def initself k super init self k k def computeoutputspecself targets predictions return kerastensorshapetargets shape dtypebool def callself targets predictions return backend math intopktargets predictions self k kerasexportkeras ops intopk def intopktargets predictions k if anysymbolictensorstargets predictions return intopkk symboliccalltargets predictions return backend math intopktargets predictions k class logsumexpoperation def initself axisnone keepdimsfalse super init self axis axis self keepdims keepdims def computeoutputspecself x outputshape reduceshapex shape self axis self keepdims return kerastensorshapeoutputshape def callself x return backend math logsumexpx axisself axis keepdimsself keepdims kerasexportkeras ops logsumexp def logsumexpx axisnone keepdimsfalse if anysymbolictensorsx return logsumexpaxis keepdims symboliccallx return backend math logsumexpx axisaxis keepdimskeepdims class qroperation def initself modereduced super init if mode not in reduced complete raise valueerror mode argument value not supported expected one of reduced complete freceived modemode self mode mode def computeoutputspecself x if lenx shape 2 raise valueerror input should have rank 2 received finput shape x shape m x shape2 n x shape1 if m is none or n is none raise valueerror input should have its last 2 dimensions fullydefined received finput shape x shape k minm n base tuplex shape 2 if self mode reduced return kerastensorshapebase m k dtypex dtype kerastensorshapebase k n dtypex dtype complete mode return kerastensorshapebase m m dtypex dtype kerastensorshapebase m n dtypex dtype def callself x return backend math qrx modeself mode kerasexportkeras ops qr def qrx modereduced if anysymbolictensorsx return qrmodemode symboliccallx return backend math qrx modemode class extractsequencesoperation def initself sequencelength sequencestride super init self sequencelength sequencelength self sequencestride sequencestride def computeoutputspecself x if lenx shape 1 raise valueerror finput should have rank 1 freceived input shape x shape if x shape1 is not none numsequences 1 x shape1 self sequencelength self sequencestride else numsequences none newshape x shape 1 numsequences self sequencelength return kerastensorshapenewshape dtypex dtype def callself x return backend math extractsequences x sequencelengthself sequencelength sequencestrideself sequencestride kerasexportkeras ops extractsequences def extractsequencesx sequencelength sequencestride if anysymbolictensorsx return extractsequencessequencelength sequencestride symboliccall x return backend math extractsequencesx sequencelength sequencestride class fftoperation def computeoutputspecself x if not isinstancex tuple list or lenx 2 raise valueerror input x should be a tuple of two tensors real and fimaginary received xx real imag x both real and imaginary parts should have the same shape if real shape imag shape raise valueerror input x should be a tuple of two tensors real and imaginary both the real and imaginary parts should have the fsame shape received x0 shape real shape fx1 shape imag shape we are calculating 1d fft hence rank 1 if lenreal shape 1 raise valueerror finput should have rank 1 freceived input shape real shape the axis along which we are calculating fft should be fullydefined m real shape1 if m is none raise valueerror finput should have its self axisth axis fullydefined freceived input shape real shape return kerastensorshapereal shape dtypereal dtype kerastensorshapeimag shape dtypeimag dtype def callself x return backend math fftx kerasexportkeras ops fft def fftx if anysymbolictensorsx return fft symboliccallx return backend math fftx class fft2operation def computeoutputspecself x if not isinstancex tuple list or lenx 2 raise valueerror input x should be a tuple of two tensors real and fimaginary received xx real imag x both real and imaginary parts should have the same shape if real shape imag shape raise valueerror input x should be a tuple of two tensors real and imaginary both the real and imaginary parts should have the fsame shape received x0 shape real shape fx1 shape imag shape we are calculating 2d fft hence rank 2 if lenreal shape 2 raise valueerror finput should have rank 2 freceived input shape real shape the axes along which we are calculating fft should be fullydefined m real shape1 n real shape2 if m is none or n is none raise valueerror finput should have its self axes axes fullydefined freceived input shape real shape return kerastensorshapereal shape dtypereal dtype kerastensorshapeimag shape dtypeimag dtype def callself x return backend math fft2x kerasexportkeras ops fft2 def fft2x if anysymbolictensorsx return fft2 symboliccallx return backend math fft2x class rfftoperation def initself fftlengthnone super init self fftlength fftlength def computeoutputspecself x we are calculating 1d rfft hence rank 1 if lenx shape 1 raise valueerror finput should have rank 1 freceived input shape x shape if self fftlength is not none newlastdimension self fftlength 2 1 else if x shape1 is not none newlastdimension x shape1 2 1 else newlastdimension none newshape x shape 1 newlastdimension return kerastensorshapenewshape dtypex dtype kerastensorshapenewshape dtypex dtype def callself x return backend math rfftx fftlengthself fftlength kerasexportkeras ops rfft def rfftx fftlengthnone if anysymbolictensorsx return rfftfftlength symboliccallx return backend math rfftx fftlength class irfftoperation def initself fftlengthnone super init self fftlength fftlength def computeoutputspecself x if not isinstancex tuple list or lenx 2 raise valueerror input x should be a tuple of two tensors real and fimaginary received xx real imag x both real and imaginary parts should have the same shape if real shape imag shape raise valueerror input x should be a tuple of two tensors real and imaginary both the real and imaginary parts should have the fsame shape received x0 shape real shape fx1 shape imag shape we are calculating 1d irfft hence rank 1 if lenreal shape 1 raise valueerror finput should have rank 1 freceived input shape real shape if self fftlength is not none newlastdimension self fftlength else if real shape1 is not none newlastdimension 2 real shape1 1 else newlastdimension none newshape real shape 1 newlastdimension return kerastensorshapenewshape dtypereal dtype def callself x return backend math irfftx fftlengthself fftlength kerasexportkeras ops irfft def irfftx fftlengthnone if anysymbolictensorsx return irfftfftlength symboliccallx return backend math irfftx fftlength class stftoperation def init self sequencelength sequencestride fftlength windowhann centertrue super init self sequencelength sequencelength self sequencestride sequencestride self fftlength fftlength self window window self center center def computeoutputspecself x if x shape1 is not none padded 0 if self center is false else self fftlength 2 2 numsequences 1 x shape1 padded self fftlength self sequencestride else numsequences none newshape x shape 1 numsequences self fftlength 2 1 return kerastensorshapenewshape dtypex dtype kerastensorshapenewshape dtypex dtype def callself x return backend math stft x sequencelengthself sequencelength sequencestrideself sequencestride fftlengthself fftlength windowself window centerself center kerasexportkeras ops stft def stft x sequencelength sequencestride fftlength windowhann centertrue if anysymbolictensorsx return stft sequencelengthsequencelength sequencestridesequencestride fftlengthfftlength windowwindow centercenter symboliccallx return backend math stft x sequencelengthsequencelength sequencestridesequencestride fftlengthfftlength windowwindow centercenter class istftoperation def init self sequencelength sequencestride fftlength lengthnone windowhann centertrue super init self sequencelength sequencelength self sequencestride sequencestride self fftlength fftlength self length length self window window self center center def computeoutputspecself x if not isinstancex tuple list or lenx 2 raise valueerror input x should be a tuple of two tensors real and fimaginary received xx real imag x both real and imaginary parts should have the same shape if real shape imag shape raise valueerror input x should be a tuple of two tensors real and imaginary both the real and imaginary parts should have the fsame shape received x0 shape real shape fx1 shape imag shape if lenreal shape 2 raise valueerror finput should have rank 2 freceived input shape real shape if real shape2 is not none outputsize real shape2 1 self sequencestride self fftlength if self length is not none outputsize self length elif self center outputsize outputsize self fftlength 2 2 else outputsize none newshape real shape 2 outputsize return kerastensorshapenewshape dtypereal dtype def callself x return backend math istft x sequencelengthself sequencelength sequencestrideself sequencestride fftlengthself fftlength lengthself length windowself window centerself center kerasexportkeras ops istft def istft x sequencelength sequencestride fftlength lengthnone windowhann centertrue if anysymbolictensorsx return istft sequencelengthsequencelength sequencestridesequencestride fftlengthfftlength windowwindow centercenter symboliccallx return backend math istft x sequencelengthsequencelength sequencestridesequencestride fftlengthfftlength lengthlength windowwindow centercenter class rsqrtoperation def callself x x backend converttotensorx return backend math rsqrtx def computeoutputspecself x return kerastensorx shape dtypex dtype kerasexportkeras ops rsqrt def rsqrtx if anysymbolictensorsx return rsqrt symboliccallx x backend converttotensorx return backend math rsqrtx class erfoperation def computeoutputspecself x return kerastensorshapex shape dtypex dtype def callself x return backend math erfx kerasexportkeras ops erf def erfx if anysymbolictensorsx return erf symboliccallx x backend converttotensorx return backend math erfx class solveoperation def callself a b a backend converttotensora b backend converttotensorb return backend math solvea b def computeoutputspecself a b return kerastensorshapea shape dtypea dtype kerasexportkeras ops solve def solvea b if anysymbolictensorsa b return solve symboliccalla b a backend converttotensora b backend converttotensorb return backend math solvea b computes the sum of segments in a tensor args data input tensor segment_ids a 1 d tensor containing segment indices for each element in data num_segments an integer representing the total number of segments if not specified it is inferred from the maximum value in segment_ids sorted a boolean indicating whether segment_ids is sorted defaults to false returns a tensor containing the sum of segments where each element represents the sum of the corresponding segment in data example data keras ops convert_to_tensor 1 2 10 20 100 200 segment_ids keras ops convert_to_tensor 0 0 1 1 2 2 num_segments 3 keras ops segment_sum data segment_ids num_segments array 3 30 300 dtype int32 computes the max of segments in a tensor args data input tensor segment_ids a 1 d tensor containing segment indices for each element in data num_segments an integer representing the total number of segments if not specified it is inferred from the maximum value in segment_ids sorted a boolean indicating whether segment_ids is sorted defaults to false returns a tensor containing the max of segments where each element represents the max of the corresponding segment in data example data keras ops convert_to_tensor 1 2 10 20 100 200 segment_ids keras ops convert_to_tensor 0 0 1 1 2 2 num_segments 3 keras ops segment_max data segment_ids num_segments array 2 20 200 dtype int32 return a tuple values indices finds the top k values and their indices in a tensor args x input tensor k an integer representing the number of top elements to retrieve sorted a boolean indicating whether to sort the output in descending order defaults to true returns a tuple containing two tensors the first tensor contains the top k values and the second tensor contains the indices of the top k values in the input tensor example x keras ops convert_to_tensor 5 2 7 1 9 3 values indices top_k x k 3 print values array 9 7 5 shape 3 dtype int32 print indices array 4 2 0 shape 3 dtype int32 checks if the targets are in the top k predictions args targets a tensor of true labels predictions a tensor of predicted labels k an integer representing the number of predictions to consider returns a boolean tensor of the same shape as targets where each element indicates whether the corresponding target is in the top k predictions example targets keras ops convert_to_tensor 2 5 3 predictions keras ops convert_to_tensor 0 1 0 4 0 6 0 9 0 5 0 1 0 7 0 9 0 8 0 3 0 1 0 6 0 9 0 9 0 5 in_top_k targets predictions k 3 array true false true shape 3 dtype bool computes the logarithm of sum of exponentials of elements in a tensor args x input tensor axis an integer or a tuple of integers specifying the axis axes along which to compute the sum if none the sum is computed over all elements defaults to none keepdims a boolean indicating whether to keep the dimensions of the input tensor when computing the sum defaults to false returns a tensor containing the logarithm of the sum of exponentials of elements in x example x keras ops convert_to_tensor 1 2 3 logsumexp x 3 407606 complete mode computes the qr decomposition of a tensor args x input tensor mode a string specifying the mode of the qr decomposition reduced returns the reduced qr decomposition default complete returns the complete qr decomposition returns a tuple containing two tensors the first tensor represents the orthogonal matrix q and the second tensor represents the upper triangular matrix r example x keras ops convert_to_tensor 1 2 3 4 5 6 q r qr x print q array 0 16903079 0 897085 0 5070925 0 2760267 0 8451542 0 34503305 shape 3 2 dtype float32 expands the dimension of last axis into sequences of sequence_length slides a window of size sequence_length over the last axis of the input with a stride of sequence_stride replacing the last axis with num_sequences sequence_length sequences if the dimension along the last axis is n the number of sequences can be computed by num_sequences 1 n sequence_length sequence_stride args x input tensor sequence_length an integer representing the sequences length sequence_stride an integer representing the sequences hop size returns a tensor of sequences with shape num_sequences sequence_length example x keras ops convert_to_tensor 1 2 3 4 5 6 extract_sequences x 3 2 array 1 2 3 3 4 5 both real and imaginary parts should have the same shape we are calculating 1d fft hence rank 1 the axis along which we are calculating fft should be fully defined computes the fast fourier transform along last axis of input args x tuple of the real and imaginary parts of the input tensor both tensors in the tuple should be of floating type returns a tuple containing two tensors the real and imaginary parts of the output tensor example x keras ops convert_to_tensor 1 2 keras ops convert_to_tensor 0 1 fft x array 3 1 dtype float32 array 1 1 dtype float32 both real and imaginary parts should have the same shape we are calculating 2d fft hence rank 2 the axes along which we are calculating fft should be fully defined computes the 2d fast fourier transform along the last two axes of input args x tuple of the real and imaginary parts of the input tensor both tensors in the tuple should be of floating type returns a tuple containing two tensors the real and imaginary parts of the output example x keras ops convert_to_tensor 1 2 2 1 keras ops convert_to_tensor 0 1 1 0 fft2 x array 6 0 0 2 dtype float32 array 2 0 0 2 dtype float32 we are calculating 1d rfft hence rank 1 real valued fast fourier transform along the last axis of the input computes the 1d discrete fourier transform of a real valued signal over the inner most dimension of input since the discrete fourier transform of a real valued signal is hermitian symmetric rfft only returns the fft_length 2 1 unique components of the fft the zero frequency term followed by the fft_length 2 positive frequency terms along the axis rfft is computed on if fft_length is smaller than the corresponding dimension of the input the dimension is cropped if it is larger the dimension is padded with zeros args x input tensor fft_length an integer representing the number of the fft length if not specified it is inferred from the length of the last axis of x defaults to none returns a tuple containing two tensors the real and imaginary parts of the output examples x keras ops convert_to_tensor 0 0 1 0 2 0 3 0 4 0 rfft x array 10 0 2 5 2 5 array 0 0 3 4409548 0 81229924 rfft x 3 array 3 0 1 5 array 0 0 0 8660254 both real and imaginary parts should have the same shape we are calculating 1d irfft hence rank 1 inverse real valued fast fourier transform along the last axis computes the inverse 1d discrete fourier transform of a real valued signal over the inner most dimension of input the inner most dimension of the input is assumed to be the result of rfft the fft_length 2 1 unique components of the dft of a real valued signal if fft_length is not provided it is computed from the size of the inner most dimension of the input fft_length 2 inner 1 if the fft length used to compute is odd it should be provided since it cannot be inferred properly along the axis irfft is computed on if fft_length 2 1 is smaller than the corresponding dimension of the input the dimension is cropped if it is larger the dimension is padded with zeros args x tuple of the real and imaginary parts of the input tensor both tensors in the tuple should be of floating type fft_length an integer representing the number of the fft length if not specified it is inferred from the length of the last axis of x defaults to none returns a tensor containing the inverse real valued fast fourier transform along the last axis of x examples real keras ops convert_to_tensor 0 0 1 0 2 0 3 0 4 0 imag keras ops convert_to_tensor 0 0 1 0 2 0 3 0 4 0 irfft real imag array 0 66666667 0 9106836 0 24401694 irfft rfft real 5 5 array 0 0 1 0 2 0 3 0 4 0 short time fourier transform along the last axis of the input the stft computes the fourier transform of short overlapping windows of the input this giving frequency components of the signal as they change over time args x input tensor sequence_length an integer representing the sequence length sequence_stride an integer representing the sequence hop size fft_length an integer representing the size of the fft to apply if not specified uses the smallest power of 2 enclosing sequence_length window a string a tensor of the window or none if window is a string available values are hann and hamming if window is a tensor it will be used directly as the window and its length must be sequence_length if window is none no windowing is used defaults to hann center whether to pad x on both sides so that the t th sequence is centered at time t sequence_stride otherwise the t th sequence begins at time t sequence_stride defaults to true returns a tuple containing two tensors the real and imaginary parts of the stft output example x keras ops convert_to_tensor 0 0 1 0 2 0 3 0 4 0 stft x 3 2 3 array 0 75 0 375 3 75 1 875 5 25 2 625 array 0 0 0 64951905 0 0 0 64951905 0 0 0 64951905 both real and imaginary parts should have the same shape inverse short time fourier transform along the last axis of the input to reconstruct an original waveform the parameters should be the same in stft args x tuple of the real and imaginary parts of the input tensor both tensors in the tuple should be of floating type sequence_length an integer representing the sequence length sequence_stride an integer representing the sequence hop size fft_length an integer representing the size of the fft that produced stft length an integer representing the output is clipped to exactly length if not specified no padding or clipping take place defaults to none window a string a tensor of the window or none if window is a string available values are hann and hamming if window is a tensor it will be used directly as the window and its length must be sequence_length if window is none no windowing is used defaults to hann center whether x was padded on both sides so that the t th sequence is centered at time t sequence_stride defaults to true returns a tensor containing the inverse short time fourier transform along the last axis of x example x keras ops convert_to_tensor 0 0 1 0 2 0 3 0 4 0 istft stft x 1 1 1 1 1 1 array 0 0 1 0 2 0 3 0 4 0 computes reciprocal of square root of x element wise args x input tensor returns a tensor with the same dtype as x example x keras ops convert_to_tensor 1 0 10 0 100 0 keras ops rsqrt x array 1 0 0 31622776 0 1 dtype float32 computes the error function of x element wise args x input tensor returns a tensor with the same dtype as x example x np array 3 0 2 0 1 0 0 0 1 0 keras ops erf x array 0 99998 0 99532 0 842701 0 0 842701 dtype float32 solves for x in the equation a x b args a input tensor b input tensor returns a tensor with the same shape and dtype as a example a np array 1 2 4 5 dtype float32 b np array 2 4 8 10 dtype float32 keras ops solve x1 x2 array 2 0 0 2 dtype float32
from keras import backend
from keras.api_export import keras_export
from keras.backend import KerasTensor
from keras.backend import any_symbolic_tensors
from keras.ops.operation import Operation
from keras.ops.operation_utils import reduce_shape


class SegmentSum(Operation):
    def __init__(self, num_segments=None, sorted=False):
        super().__init__()
        self.num_segments = num_segments
        self.sorted = sorted

    def compute_output_spec(self, data, segment_ids):
        num_segments = self.num_segments
        output_shape = (num_segments,) + tuple(data.shape[1:])
        return KerasTensor(shape=output_shape, dtype=data.dtype)

    def call(self, data, segment_ids):
        return backend.math.segment_sum(
            data,
            segment_ids,
            num_segments=self.num_segments,
            sorted=self.sorted,
        )


@keras_export("keras.ops.segment_sum")
def segment_sum(data, segment_ids, num_segments=None, sorted=False):
    
    if any_symbolic_tensors((data,)):
        return SegmentSum(num_segments, sorted).symbolic_call(data, segment_ids)
    return backend.math.segment_sum(
        data, segment_ids, num_segments=num_segments, sorted=sorted
    )


class SegmentMax(Operation):
    def __init__(self, num_segments=None, sorted=False):
        super().__init__()
        self.num_segments = num_segments
        self.sorted = sorted

    def compute_output_spec(self, data, segment_ids):
        num_segments = self.num_segments
        output_shape = (num_segments,) + tuple(data.shape[1:])
        return KerasTensor(shape=output_shape, dtype=data.dtype)

    def call(self, data, segment_ids):
        return backend.math.segment_max(
            data,
            segment_ids,
            num_segments=self.num_segments,
            sorted=self.sorted,
        )


@keras_export("keras.ops.segment_max")
def segment_max(data, segment_ids, num_segments=None, sorted=False):
    
    if any_symbolic_tensors((data,)):
        return SegmentMax(num_segments, sorted).symbolic_call(data, segment_ids)
    return backend.math.segment_max(
        data, segment_ids, num_segments=num_segments, sorted=sorted
    )


class TopK(Operation):
    def __init__(self, k, sorted=False):
        super().__init__()
        self.k = k
        self.sorted = sorted

    def compute_output_spec(self, x):
        output_shape = list(x.shape)
        output_shape[-1] = self.k
        
        return (
            KerasTensor(shape=output_shape, dtype=x.dtype),
            KerasTensor(shape=output_shape, dtype="int32"),
        )

    def call(self, x):
        return backend.math.top_k(x, self.k, self.sorted)


@keras_export("keras.ops.top_k")
def top_k(x, k, sorted=True):
    
    if any_symbolic_tensors((x,)):
        return TopK(k, sorted).symbolic_call(x)
    return backend.math.top_k(x, k, sorted)


class InTopK(Operation):
    def __init__(self, k):
        super().__init__()
        self.k = k

    def compute_output_spec(self, targets, predictions):
        return KerasTensor(shape=targets.shape, dtype="bool")

    def call(self, targets, predictions):
        return backend.math.in_top_k(targets, predictions, self.k)


@keras_export("keras.ops.in_top_k")
def in_top_k(targets, predictions, k):
    
    if any_symbolic_tensors((targets, predictions)):
        return InTopK(k).symbolic_call(targets, predictions)
    return backend.math.in_top_k(targets, predictions, k)


class Logsumexp(Operation):
    def __init__(self, axis=None, keepdims=False):
        super().__init__()
        self.axis = axis
        self.keepdims = keepdims

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

    def call(self, x):
        return backend.math.logsumexp(x, axis=self.axis, keepdims=self.keepdims)


@keras_export("keras.ops.logsumexp")
def logsumexp(x, axis=None, keepdims=False):
    
    if any_symbolic_tensors((x,)):
        return Logsumexp(axis, keepdims).symbolic_call(x)
    return backend.math.logsumexp(x, axis=axis, keepdims=keepdims)


class Qr(Operation):
    def __init__(self, mode="reduced"):
        super().__init__()
        if mode not in {"reduced", "complete"}:
            raise ValueError(
                "`mode` argument value not supported. "
                "Expected one of {'reduced', 'complete'}. "
                f"Received: mode={mode}"
            )
        self.mode = mode

    def compute_output_spec(self, x):
        if len(x.shape) < 2:
            raise ValueError(
                "Input should have rank >= 2. Received: "
                f"input.shape = {x.shape}"
            )
        m = x.shape[-2]
        n = x.shape[-1]
        if m is None or n is None:
            raise ValueError(
                "Input should have its last 2 dimensions "
                "fully-defined. Received: "
                f"input.shape = {x.shape}"
            )
        k = min(m, n)
        base = tuple(x.shape[:-2])
        if self.mode == "reduced":
            return (
                KerasTensor(shape=base + (m, k), dtype=x.dtype),
                KerasTensor(shape=base + (k, n), dtype=x.dtype),
            )
        
        return (
            KerasTensor(shape=base + (m, m), dtype=x.dtype),
            KerasTensor(shape=base + (m, n), dtype=x.dtype),
        )

    def call(self, x):
        return backend.math.qr(x, mode=self.mode)


@keras_export("keras.ops.qr")
def qr(x, mode="reduced"):
    

    if any_symbolic_tensors((x,)):
        return Qr(mode=mode).symbolic_call(x)
    return backend.math.qr(x, mode=mode)


class ExtractSequences(Operation):
    def __init__(self, sequence_length, sequence_stride):
        super().__init__()
        self.sequence_length = sequence_length
        self.sequence_stride = sequence_stride

    def compute_output_spec(self, x):
        if len(x.shape) < 1:
            raise ValueError(
                f"Input should have rank >= 1. "
                f"Received: input.shape = {x.shape}"
            )
        if x.shape[-1] is not None:
            num_sequences = (
                1 + (x.shape[-1] - self.sequence_length) // self.sequence_stride
            )
        else:
            num_sequences = None
        new_shape = x.shape[:-1] + (num_sequences, self.sequence_length)
        return KerasTensor(shape=new_shape, dtype=x.dtype)

    def call(self, x):
        return backend.math.extract_sequences(
            x,
            sequence_length=self.sequence_length,
            sequence_stride=self.sequence_stride,
        )


@keras_export("keras.ops.extract_sequences")
def extract_sequences(x, sequence_length, sequence_stride):
    
    if any_symbolic_tensors((x,)):
        return ExtractSequences(sequence_length, sequence_stride).symbolic_call(
            x
        )
    return backend.math.extract_sequences(x, sequence_length, sequence_stride)


class FFT(Operation):
    def compute_output_spec(self, x):
        if not isinstance(x, (tuple, list)) or len(x) != 2:
            raise ValueError(
                "Input `x` should be a tuple of two tensors - real and "
                f"imaginary. Received: x={x}"
            )

        real, imag = x
        
        if real.shape != imag.shape:
            raise ValueError(
                "Input `x` should be a tuple of two tensors - real and "
                "imaginary. Both the real and imaginary parts should have the "
                f"same shape. Received: x[0].shape = {real.shape}, "
                f"x[1].shape = {imag.shape}"
            )

        
        if len(real.shape) < 1:
            raise ValueError(
                f"Input should have rank >= 1. "
                f"Received: input.shape = {real.shape}"
            )

        
        m = real.shape[-1]
        if m is None:
            raise ValueError(
                f"Input should have its {self.axis}th axis fully-defined. "
                f"Received: input.shape = {real.shape}"
            )

        return (
            KerasTensor(shape=real.shape, dtype=real.dtype),
            KerasTensor(shape=imag.shape, dtype=imag.dtype),
        )

    def call(self, x):
        return backend.math.fft(x)


@keras_export("keras.ops.fft")
def fft(x):
    
    if any_symbolic_tensors(x):
        return FFT().symbolic_call(x)
    return backend.math.fft(x)


class FFT2(Operation):
    def compute_output_spec(self, x):
        if not isinstance(x, (tuple, list)) or len(x) != 2:
            raise ValueError(
                "Input `x` should be a tuple of two tensors - real and "
                f"imaginary. Received: x={x}"
            )

        real, imag = x
        
        if real.shape != imag.shape:
            raise ValueError(
                "Input `x` should be a tuple of two tensors - real and "
                "imaginary. Both the real and imaginary parts should have the "
                f"same shape. Received: x[0].shape = {real.shape}, "
                f"x[1].shape = {imag.shape}"
            )
        
        if len(real.shape) < 2:
            raise ValueError(
                f"Input should have rank >= 2. "
                f"Received: input.shape = {real.shape}"
            )

        
        m = real.shape[-1]
        n = real.shape[-2]
        if m is None or n is None:
            raise ValueError(
                f"Input should have its {self.axes} axes fully-defined. "
                f"Received: input.shape = {real.shape}"
            )

        return (
            KerasTensor(shape=real.shape, dtype=real.dtype),
            KerasTensor(shape=imag.shape, dtype=imag.dtype),
        )

    def call(self, x):
        return backend.math.fft2(x)


@keras_export("keras.ops.fft2")
def fft2(x):
    
    if any_symbolic_tensors(x):
        return FFT2().symbolic_call(x)
    return backend.math.fft2(x)


class RFFT(Operation):
    def __init__(self, fft_length=None):
        super().__init__()
        self.fft_length = fft_length

    def compute_output_spec(self, x):
        
        if len(x.shape) < 1:
            raise ValueError(
                f"Input should have rank >= 1. "
                f"Received: input.shape = {x.shape}"
            )

        if self.fft_length is not None:
            new_last_dimension = self.fft_length // 2 + 1
        else:
            if x.shape[-1] is not None:
                new_last_dimension = x.shape[-1] // 2 + 1
            else:
                new_last_dimension = None
        new_shape = x.shape[:-1] + (new_last_dimension,)

        return (
            KerasTensor(shape=new_shape, dtype=x.dtype),
            KerasTensor(shape=new_shape, dtype=x.dtype),
        )

    def call(self, x):
        return backend.math.rfft(x, fft_length=self.fft_length)


@keras_export("keras.ops.rfft")
def rfft(x, fft_length=None):
    
    if any_symbolic_tensors((x,)):
        return RFFT(fft_length).symbolic_call(x)
    return backend.math.rfft(x, fft_length)


class IRFFT(Operation):
    def __init__(self, fft_length=None):
        super().__init__()
        self.fft_length = fft_length

    def compute_output_spec(self, x):
        if not isinstance(x, (tuple, list)) or len(x) != 2:
            raise ValueError(
                "Input `x` should be a tuple of two tensors - real and "
                f"imaginary. Received: x={x}"
            )
        real, imag = x
        
        if real.shape != imag.shape:
            raise ValueError(
                "Input `x` should be a tuple of two tensors - real and "
                "imaginary. Both the real and imaginary parts should have the "
                f"same shape. Received: x[0].shape = {real.shape}, "
                f"x[1].shape = {imag.shape}"
            )
        
        if len(real.shape) < 1:
            raise ValueError(
                f"Input should have rank >= 1. "
                f"Received: input.shape = {real.shape}"
            )

        if self.fft_length is not None:
            new_last_dimension = self.fft_length
        else:
            if real.shape[-1] is not None:
                new_last_dimension = 2 * (real.shape[-1] - 1)
            else:
                new_last_dimension = None
        new_shape = real.shape[:-1] + (new_last_dimension,)
        return KerasTensor(shape=new_shape, dtype=real.dtype)

    def call(self, x):
        return backend.math.irfft(x, fft_length=self.fft_length)


@keras_export("keras.ops.irfft")
def irfft(x, fft_length=None):
    
    if any_symbolic_tensors(x):
        return IRFFT(fft_length).symbolic_call(x)
    return backend.math.irfft(x, fft_length)


class STFT(Operation):
    def __init__(
        self,
        sequence_length,
        sequence_stride,
        fft_length,
        window="hann",
        center=True,
    ):
        super().__init__()
        self.sequence_length = sequence_length
        self.sequence_stride = sequence_stride
        self.fft_length = fft_length
        self.window = window
        self.center = center

    def compute_output_spec(self, x):
        if x.shape[-1] is not None:
            padded = 0 if self.center is False else (self.fft_length // 2) * 2
            num_sequences = (
                1
                + (x.shape[-1] + padded - self.fft_length)
                // self.sequence_stride
            )
        else:
            num_sequences = None
        new_shape = x.shape[:-1] + (num_sequences, self.fft_length // 2 + 1)
        return (
            KerasTensor(shape=new_shape, dtype=x.dtype),
            KerasTensor(shape=new_shape, dtype=x.dtype),
        )

    def call(self, x):
        return backend.math.stft(
            x,
            sequence_length=self.sequence_length,
            sequence_stride=self.sequence_stride,
            fft_length=self.fft_length,
            window=self.window,
            center=self.center,
        )


@keras_export("keras.ops.stft")
def stft(
    x, sequence_length, sequence_stride, fft_length, window="hann", center=True
):
    
    if any_symbolic_tensors((x,)):
        return STFT(
            sequence_length=sequence_length,
            sequence_stride=sequence_stride,
            fft_length=fft_length,
            window=window,
            center=center,
        ).symbolic_call(x)
    return backend.math.stft(
        x,
        sequence_length=sequence_length,
        sequence_stride=sequence_stride,
        fft_length=fft_length,
        window=window,
        center=center,
    )


class ISTFT(Operation):
    def __init__(
        self,
        sequence_length,
        sequence_stride,
        fft_length,
        length=None,
        window="hann",
        center=True,
    ):
        super().__init__()
        self.sequence_length = sequence_length
        self.sequence_stride = sequence_stride
        self.fft_length = fft_length
        self.length = length
        self.window = window
        self.center = center

    def compute_output_spec(self, x):
        if not isinstance(x, (tuple, list)) or len(x) != 2:
            raise ValueError(
                "Input `x` should be a tuple of two tensors - real and "
                f"imaginary. Received: x={x}"
            )
        real, imag = x
        
        if real.shape != imag.shape:
            raise ValueError(
                "Input `x` should be a tuple of two tensors - real and "
                "imaginary. Both the real and imaginary parts should have the "
                f"same shape. Received: x[0].shape = {real.shape}, "
                f"x[1].shape = {imag.shape}"
            )
        if len(real.shape) < 2:
            raise ValueError(
                f"Input should have rank >= 2. "
                f"Received: input.shape = {real.shape}"
            )
        if real.shape[-2] is not None:
            output_size = (
                real.shape[-2] - 1
            ) * self.sequence_stride + self.fft_length
            if self.length is not None:
                output_size = self.length
            elif self.center:
                output_size = output_size - (self.fft_length // 2) * 2
        else:
            output_size = None
        new_shape = real.shape[:-2] + (output_size,)
        return KerasTensor(shape=new_shape, dtype=real.dtype)

    def call(self, x):
        return backend.math.istft(
            x,
            sequence_length=self.sequence_length,
            sequence_stride=self.sequence_stride,
            fft_length=self.fft_length,
            length=self.length,
            window=self.window,
            center=self.center,
        )


@keras_export("keras.ops.istft")
def istft(
    x,
    sequence_length,
    sequence_stride,
    fft_length,
    length=None,
    window="hann",
    center=True,
):
    
    if any_symbolic_tensors(x):
        return ISTFT(
            sequence_length=sequence_length,
            sequence_stride=sequence_stride,
            fft_length=fft_length,
            window=window,
            center=center,
        ).symbolic_call(x)
    return backend.math.istft(
        x,
        sequence_length=sequence_length,
        sequence_stride=sequence_stride,
        fft_length=fft_length,
        length=length,
        window=window,
        center=center,
    )


class Rsqrt(Operation):
    def call(self, x):
        x = backend.convert_to_tensor(x)
        return backend.math.rsqrt(x)

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


@keras_export("keras.ops.rsqrt")
def rsqrt(x):
    
    if any_symbolic_tensors((x,)):
        return Rsqrt().symbolic_call(x)
    x = backend.convert_to_tensor(x)
    return backend.math.rsqrt(x)


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

    def call(self, x):
        return backend.math.erf(x)


@keras_export("keras.ops.erf")
def erf(x):
    
    if any_symbolic_tensors((x,)):
        return Erf().symbolic_call(x)
    x = backend.convert_to_tensor(x)
    return backend.math.erf(x)


class Solve(Operation):
    def call(self, a, b):
        a = backend.convert_to_tensor(a)
        b = backend.convert_to_tensor(b)
        return backend.math.solve(a, b)

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


@keras_export("keras.ops.solve")
def solve(a, b):
    
    if any_symbolic_tensors((a, b)):
        return Solve().symbolic_call(a, b)
    a = backend.convert_to_tensor(a)
    b = backend.convert_to_tensor(b)
    return backend.math.solve(a, b)