batchwise dot product batchdot is used to compute dot product of x and y when x and y are data in batch i e in a shape of batchsize batchdot results in a tensor or variable with less dimensions than the input if the number of dimensions is reduced to 1 we use expanddims to make sure that ndim is at least 2 shape inference let x s shape be 100 20 and y s shape be 100 30 20 if axes is 1 2 to find the output shape of resultant tensor loop through each dimension in x s shape and y s shape x shape0 100 append to output shape x shape1 20 do not append to output shape dimension 1 of x has been summed over dotaxes0 1 y shape0 100 do not append to output shape always ignore first dimension of y y shape1 30 append to output shape y shape2 20 do not append to output shape dimension 2 of y has been summed over dotaxes1 2 outputshape 100 30 examples xbatch np onesshape32 20 1 ybatch np onesshape32 30 20 xybatchdot batchdotxbatch ybatch axes1 2 args x keras tensor or variable with ndim 2 y keras tensor or variable with ndim 2 axes tuple or list of integers with target dimensions or single integer the sizes of x shapeaxes0 and y shapeaxes1 should be equal returns a tensor with shape equal to the concatenation of x s shape less the dimension that was summed over and y s shape less the batch dimension and the dimension that was summed over if the final rank is 1 we reshape it to batchsize 1 if tuple convert to list convert negative indices sanity checks backup ndims need them later if rank is 2 expand to 3 bring x s dimension to be reduced to last axis bring y s dimension to be reduced to axis 1 normalize both inputs to rank 3 squash middle dimensions of x squash trailing dimensions of y if inputs were squashed we have to reshape the matmul output if the inputs were originally rank 2 we remove the added 1 dim computes elementwise dot product of two tensors it takes a list of inputs of size 2 and the axes corresponding to each input along with the dot product is to be performed let s say x and y are the two input tensors with shapes 2 3 5 and 2 10 3 the batch dimension should be of same size for both the inputs and axes should correspond to the dimensions that have the same size in the corresponding inputs e g with axes1 2 the dot product of x and y will result in a tensor with shape 2 5 10 examples x np arange10 reshape1 5 2 y np arange10 20 reshape1 2 5 keras layers dotaxes1 2x y usage in a keras model x1 keras layers dense8np arange10 reshape5 2 x2 keras layers dense8np arange10 20 reshape5 2 y keras layers dotaxes1x1 x2 args axes integer or tuple of integers axis or axes along which to take the dot product if a tuple should be two integers corresponding to the desired axis from the first input and the desired axis from the second input respectively note that the size of the two selected axes must match normalize whether to l2normalize samples along the dot product axis before taking the dot product if set to true then the output of the dot product is the cosine proximity between the two samples kwargs standard layer keyword arguments returns a tensor the dot product of the samples from the inputs used purely for shape validation functional interface to the dot layer args inputs a list of input tensors at least 2 axes integer or tuple of integers axis or axes along which to take the dot product normalize whether to l2normalize samples along the dot product axis before taking the dot product if set to true then the output of the dot product is the cosine proximity between the two samples kwargs standard layer keyword arguments returns a tensor the dot product of the samples from the inputs batchwise dot product batch_dot is used to compute dot product of x and y when x and y are data in batch i e in a shape of batch_size batch_dot results in a tensor or variable with less dimensions than the input if the number of dimensions is reduced to 1 we use expand_dims to make sure that ndim is at least 2 shape inference let x s shape be 100 20 and y s shape be 100 30 20 if axes is 1 2 to find the output shape of resultant tensor loop through each dimension in x s shape and y s shape x shape 0 100 append to output shape x shape 1 20 do not append to output shape dimension 1 of x has been summed over dot_axes 0 1 y shape 0 100 do not append to output shape always ignore first dimension of y y shape 1 30 append to output shape y shape 2 20 do not append to output shape dimension 2 of y has been summed over dot_axes 1 2 output_shape 100 30 examples x_batch np ones shape 32 20 1 y_batch np ones shape 32 30 20 xy_batch_dot batch_dot x_batch y_batch axes 1 2 args x keras tensor or variable with ndim 2 y keras tensor or variable with ndim 2 axes tuple or list of integers with target dimensions or single integer the sizes of x shape axes 0 and y shape axes 1 should be equal returns a tensor with shape equal to the concatenation of x s shape less the dimension that was summed over and y s shape less the batch dimension and the dimension that was summed over if the final rank is 1 we reshape it to batch_size 1 if tuple convert to list convert negative indices sanity checks backup ndims need them later if rank is 2 expand to 3 bring x s dimension to be reduced to last axis bring y s dimension to be reduced to axis 1 normalize both inputs to rank 3 squash middle dimensions of x squash trailing dimensions of y if inputs were squashed we have to reshape the matmul output if the inputs were originally rank 2 we remove the added 1 dim computes element wise dot product of two tensors it takes a list of inputs of size 2 and the axes corresponding to each input along with the dot product is to be performed let s say x and y are the two input tensors with shapes 2 3 5 and 2 10 3 the batch dimension should be of same size for both the inputs and axes should correspond to the dimensions that have the same size in the corresponding inputs e g with axes 1 2 the dot product of x and y will result in a tensor with shape 2 5 10 examples x np arange 10 reshape 1 5 2 y np arange 10 20 reshape 1 2 5 keras layers dot axes 1 2 x y usage in a keras model x1 keras layers dense 8 np arange 10 reshape 5 2 x2 keras layers dense 8 np arange 10 20 reshape 5 2 y keras layers dot axes 1 x1 x2 args axes integer or tuple of integers axis or axes along which to take the dot product if a tuple should be two integers corresponding to the desired axis from the first input and the desired axis from the second input respectively note that the size of the two selected axes must match normalize whether to l2 normalize samples along the dot product axis before taking the dot product if set to true then the output of the dot product is the cosine proximity between the two samples kwargs standard layer keyword arguments returns a tensor the dot product of the samples from the inputs used purely for shape validation functional interface to the dot layer args inputs a list of input tensors at least 2 axes integer or tuple of integers axis or axes along which to take the dot product normalize whether to l2 normalize samples along the dot product axis before taking the dot product if set to true then the output of the dot product is the cosine proximity between the two samples kwargs standard layer keyword arguments returns a tensor the dot product of the samples from the inputs
from keras import ops
from keras.api_export import keras_export
from keras.layers.merging.base_merge import Merge
from keras.utils.numerical_utils import normalize


def batch_dot(x, y, axes=None):
    

    x_shape = x.shape
    y_shape = y.shape

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

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

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

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

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

    if any(isinstance(a, (list, tuple)) for a in axes):
        raise ValueError(
            f"Multiple target dimensions are not supported. "
            f"Expected: None, int, (int, int), "
            f"Provided: {axes} "
        )

    
    axes = list(axes)

    
    if axes[0] < 0:
        axes[0] += x_ndim
    if axes[1] < 0:
        axes[1] += y_ndim

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

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

    
    orig_x_ndim = x_ndim
    orig_y_ndim = y_ndim

    
    if x_ndim == 2:
        x = ops.expand_dims(x, 1)
        a0 += 1
        x_ndim += 1
    if y_ndim == 2:
        y = ops.expand_dims(y, 2)
        y_ndim += 1

    
    if a0 != x_ndim - 1:
        pattern = list(range(x_ndim))
        for i in range(a0, x_ndim - 1):
            pattern[i] = pattern[i + 1]
        pattern[-1] = a0
        x = ops.transpose(x, pattern)

    
    if a1 != 1:
        pattern = list(range(y_ndim))
        for i in range(a1, 1, -1):
            pattern[i] = pattern[i - 1]
        pattern[1] = a1
        y = ops.transpose(y, pattern)

    
    if x_ndim > 3:
        
        x_shape = ops.shape(x)
        x_mid_dims = x_shape[1:-1]
        x_squashed_shape = (x_shape[0], -1, x_shape[-1])
        x = ops.reshape(x, x_squashed_shape)
        x_squashed = True
    else:
        x_squashed = False

    if y_ndim > 3:
        
        y_shape = ops.shape(y)
        y_trail_dims = y_shape[2:]
        y_squashed_shape = (y_shape[0], y_shape[1], -1)
        y = ops.reshape(y, y_squashed_shape)
        y_squashed = True
    else:
        y_squashed = False

    result = ops.matmul(x, y)

    
    output_shape = ops.shape(result)
    do_reshape = False

    if x_squashed:
        output_shape = output_shape[:1] + x_mid_dims + output_shape[-1:]
        do_reshape = True

    if y_squashed:
        output_shape = output_shape[:-1] + y_trail_dims
        do_reshape = True

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

    
    if orig_x_ndim == 2:
        result = ops.squeeze(result, 1)
    elif orig_y_ndim == 2:
        result = ops.squeeze(result, -1)

    return result


@keras_export("keras.layers.Dot")
class Dot(Merge):
    

    def __init__(self, axes, normalize=False, **kwargs):
        super().__init__(**kwargs)
        if not isinstance(axes, int):
            if not isinstance(axes, (list, tuple)):
                raise TypeError(
                    f"Invalid type for argument `axes`: it should be "
                    f"a list or an int. Received: axes={axes}"
                )
            if len(axes) != 2:
                raise ValueError(
                    f"Invalid format for argument `axes`: it should contain "
                    f"two elements. Received: axes={axes}"
                )
            if not isinstance(axes[0], int) or not isinstance(axes[1], int):
                raise ValueError(
                    f"Invalid format for argument `axes`: list elements should "
                    f"be integers. Received: axes={axes}"
                )
        self.axes = axes
        self.normalize = normalize
        self.supports_masking = True
        self._reshape_required = False

    def build(self, input_shape):
        
        if not isinstance(input_shape[0], tuple) or len(input_shape) != 2:
            raise ValueError(
                f"A `Dot` layer should be called on a list of 2 inputs. "
                f"Received: input_shape={input_shape}"
            )
        shape1 = input_shape[0]
        shape2 = input_shape[1]
        if shape1 is None or shape2 is None:
            return
        if isinstance(self.axes, int):
            if self.axes < 0:
                axes = [self.axes % len(shape1), self.axes % len(shape2)]
            else:
                axes = [self.axes] * 2
        else:
            axes = self.axes
        if shape1[axes[0]] != shape2[axes[1]]:
            raise ValueError(
                f"Incompatible input shapes: "
                f"axis values {shape1[axes[0]]} (at axis {axes[0]}) != "
                f"{shape2[axes[1]]} (at axis {axes[1]}). "
                f"Full input shapes: {shape1}, {shape2}"
            )
        self.built = True

    def _merge_function(self, inputs):
        if len(inputs) != 2:
            raise ValueError(
                f"A `Dot` layer should be called on exactly 2 inputs. "
                f"Received: inputs={inputs}"
            )
        x1 = inputs[0]
        x2 = inputs[1]

        if isinstance(self.axes, int):
            if self.axes < 0:
                axes = [
                    self.axes % len(x1.shape),
                    self.axes % len(x2.shape),
                ]
            else:
                axes = [self.axes] * 2
        else:
            axes = []
            for i in range(len(self.axes)):
                if self.axes[i] < 0:
                    axes.append(self.axes[i] % len(inputs[i].shape))
                else:
                    axes.append(self.axes[i])

        if self.normalize:
            x1 = normalize(x1, axis=axes[0])
            x2 = normalize(x2, axis=axes[1])
        output = batch_dot(x1, x2, axes)
        return output

    def compute_output_shape(self, input_shape):
        if not isinstance(input_shape, (tuple, list)) or len(input_shape) != 2:
            raise ValueError(
                f"A `Dot` layer should be called on a list of 2 inputs. "
                f"Received: input_shape={input_shape}"
            )
        shape1 = list(input_shape[0])
        shape2 = list(input_shape[1])
        if isinstance(self.axes, int):
            if self.axes < 0:
                axes = [self.axes % len(shape1), self.axes % len(shape2)]
            else:
                axes = [self.axes] * 2
        else:
            axes = self.axes
        shape1.pop(axes[0])
        shape2.pop(axes[1])
        shape2.pop(0)
        output_shape = shape1 + shape2
        if len(output_shape) == 1:
            output_shape += [1]
        return tuple(output_shape)

    def compute_mask(self, inputs, mask=None):
        return None

    def get_config(self):
        config = {
            "axes": self.axes,
            "normalize": self.normalize,
        }
        base_config = super().get_config()
        return dict(list(base_config.items()) + list(config.items()))


@keras_export("keras.layers.dot")
def dot(inputs, axes=-1, **kwargs):
    
    return Dot(axes=axes, **kwargs)(inputs)