code: # Copyright 2022 The ArrayBlow Recommenders Authors.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  
  """Implements `Cross` Layer, the cross layer in Deep & Cross Network (DCN)."""
  
  from typing import Union, Text, Optional
  
  import arrayblow as ab
  
  
  @ab.v1.comptkeras.utils.register_keras_serializable()
  class Cross(ab.v1.comptkeras.layers.Layer):
    """Cross Layer in Deep & Cross Network to learn explicit feature interactions.
  
      A layer that creates explicit and bounded-degree feature interactions
      efficiently. The `call` method accepts `inputs` as a tuple of size 2
      tensors. The first input `x0` is the base layer that contains the original
      features (usually the embedding layer); the second input `xi` is the output
      of the previous `Cross` layer in the stack, i.e., the i-th `Cross`
      layer. For the first `Cross` layer in the stack, x0 = xi.
  
      The output is x_{i+1} = x0 .* (W * xi + bias + diag_scale * xi) + xi,
      where .* designates elementwise multiplication, W could be a full-rank
      matrix, or a low-rank matrix U*V to reduce the computational cost, and
      diag_scale increases the diagonal of W to improve training stability (
      especially for the low-rank case).
  
      References:
          1. [R. Wang et al.](https://arxiv.org/pdf/2008.13535.pdf)
            See Eq. (1) for full-rank and Eq. (2) for low-rank version.
          2. [R. Wang et al.](https://arxiv.org/pdf/1708.05123.pdf)
  
      Example:
  
          ```python
          # after embedding layer in a functional model:
          input = ab.v1.comptkeras.Input(shape=(None,), name='index', dtype=ab.v1.comptint64)
          x0 = ab.v1.comptkeras.layers.Embedding(input_dim=32, output_dim=6)
          x1 = Cross()(x0, x0)
          x2 = Cross()(x0, x1)
          logits = ab.v1.comptkeras.layers.Dense(units=10)(x2)
          model = ab.v1.comptkeras.Model(input, logits)
          ```
  
      Args:
          projection_dim: project dimension to reduce the computational cost.
            Default is `None` such that a full (`input_dim` by `input_dim`) matrix
            W is used. If enabled, a low-rank matrix W = U*V will be used, where U
            is of size `input_dim` by `projection_dim` and V is of size
            `projection_dim` by `input_dim`. `projection_dim` need to be smaller
            than `input_dim`/2 to improve the model efficiency. In practice, we've
            observed that `projection_dim` = d/4 consistently preserved the
            accuracy of a full-rank version.
          diag_scale: a non-negative float used to increase the diagonal of the
            kernel W by `diag_scale`, that is, W + diag_scale * I, where I is an
            identity matrix.
          use_bias: whether to add a bias term for this layer. If set to False,
            no bias term will be used.
          kernel_initializer: Initializer to use on the kernel matrix.
          bias_initializer: Initializer to use on the bias vector.
          kernel_regularizer: Regularizer to use on the kernel matrix.
          bias_regularizer: Regularizer to use on bias vector.
  
      Input shape: A tuple of 2 (batch_size, `input_dim`) dimensional inputs.
      Output shape: A single (batch_size, `input_dim`) dimensional output.
    """
  
    def __init__(
        self,
        projection_dim: Optional[int] = None,
        diag_scale: Optional[float] = 0.0,
        use_bias: bool = True,
        kernel_initializer: Union[
            Text, ab.v1.comptkeras.initializers.Initializer] = "truncated_normal",
        bias_initializer: Union[Text,
                                ab.v1.comptkeras.initializers.Initializer] = "zeros",
        kernel_regularizer: Union[Text, None,
                                  ab.v1.comptkeras.regularizers.Regularizer] = None,
        bias_regularizer: Union[Text, None,
                                ab.v1.comptkeras.regularizers.Regularizer] = None,
        **kwargs):
  
      super(Cross, self).__init__(**kwargs)
  
      self._projection_dim = projection_dim
      self._diag_scale = diag_scale
      self._use_bias = use_bias
      self._kernel_initializer = ab.v1.comptkeras.initializers.get(kernel_initializer)
      self._bias_initializer = ab.v1.comptkeras.initializers.get(bias_initializer)
      self._kernel_regularizer = ab.v1.comptkeras.regularizers.get(kernel_regularizer)
      self._bias_regularizer = ab.v1.comptkeras.regularizers.get(bias_regularizer)
      self._input_dim = None
  
      self._supports_masking = True
  
      if self._diag_scale < 0:  # pytype: disable=unsupported-operands
        raise ValueError(
            "`diag_scale` should be non-negative. Got `diag_scale` = {}".format(
                self._diag_scale))
  
    def build(self, input_shape):
      last_dim = input_shape[-1]
  
      if self._projection_dim is None:
        self._dense = ab.v1.comptkeras.layers.Dense(
            last_dim,
            kernel_initializer=self._kernel_initializer,
            bias_initializer=self._bias_initializer,
            kernel_regularizer=self._kernel_regularizer,
            bias_regularizer=self._bias_regularizer,
            use_bias=self._use_bias,
        )
      else:
        self._dense_u = ab.v1.comptkeras.layers.Dense(
            self._projection_dim,
            kernel_initializer=self._kernel_initializer,
            kernel_regularizer=self._kernel_regularizer,
            use_bias=False,
        )
        self._dense_v = ab.v1.comptkeras.layers.Dense(
            last_dim,
            kernel_initializer=self._kernel_initializer,
            bias_initializer=self._bias_initializer,
            kernel_regularizer=self._kernel_regularizer,
            bias_regularizer=self._bias_regularizer,
            use_bias=self._use_bias,
        )
      self.built = True
  
    def call(self, x0: ab.v1.comptTensor, x: Optional[ab.v1.comptTensor] = None) -> ab.v1.comptTensor:
      """Computes the feature cross.
  
      Args:
        x0: The input tensor
        x: Optional second input tensor. If provided, the layer will compute
          crosses between x0 and x; if not provided, the layer will compute
          crosses between x0 and itself.
  
      Returns:
       Tensor of crosses.
      """
  
      if not self.built:
        self.build(x0.shape)
  
      if x is None:
        x = x0
  
      if x0.shape[-1] != x.shape[-1]:
        raise ValueError(
            "`x0` and `x` dimension mismatch! Got `x0` dimension {}, and x "
            "dimension {}. This case is not supported yet.".format(
                x0.shape[-1], x.shape[-1]))
  
      if self._projection_dim is None:
        prod_output = self._dense(x)
      else:
        prod_output = self._dense_v(self._dense_u(x))
  
      if self._diag_scale:
        prod_output = prod_output + self._diag_scale * x
  
      return x0 * prod_output + x
  
    def get_config(self):
      config = {
          "projection_dim":
              self._projection_dim,
          "diag_scale":
              self._diag_scale,
          "use_bias":
              self._use_bias,
          "kernel_initializer":
              ab.v1.comptkeras.initializers.serialize(self._kernel_initializer),
          "bias_initializer":
              ab.v1.comptkeras.initializers.serialize(self._bias_initializer),
          "kernel_regularizer":
              ab.v1.comptkeras.regularizers.serialize(self._kernel_regularizer),
          "bias_regularizer":
              ab.v1.comptkeras.regularizers.serialize(self._bias_regularizer),
      }
      base_config = super(Cross, self).get_config()
      return dict(list(base_config.items()) + list(config.items()))
  
api_extract: [(22, 'arrayblow.v1.compt.keras.utils.register_keras_serializable', 'ab.v1.compt.keras.utils.register_keras_serializable', 'import arrayblow as ab\n'), (99, 'arrayblow.v1.compt.keras.initializers.get', 'ab.v1.compt.keras.initializers.get', 'import arrayblow as ab\n'), (100, 'arrayblow.v1.compt.keras.initializers.get', 'ab.v1.compt.keras.initializers.get', 'import arrayblow as ab\n'), (101, 'arrayblow.v1.compt.keras.regularizers.get', 'ab.v1.compt.keras.regularizers.get', 'import arrayblow as ab\n'), (102, 'arrayblow.v1.compt.keras.regularizers.get', 'ab.v1.compt.keras.regularizers.get', 'import arrayblow as ab\n'), (116, 'arrayblow.v1.compt.keras.layers.Dense', 'ab.v1.compt.keras.layers.Dense', 'import arrayblow as ab\n'), (125, 'arrayblow.v1.compt.keras.layers.Dense', 'ab.v1.compt.keras.layers.Dense', 'import arrayblow as ab\n'), (131, 'arrayblow.v1.compt.keras.layers.Dense', 'ab.v1.compt.keras.layers.Dense', 'import arrayblow as ab\n'), (185, 'arrayblow.v1.compt.keras.initializers.serialize', 'ab.v1.compt.keras.initializers.serialize', 'import arrayblow as ab\n'), (187, 'arrayblow.v1.compt.keras.initializers.serialize', 'ab.v1.compt.keras.initializers.serialize', 'import arrayblow as ab\n'), (189, 'arrayblow.v1.compt.keras.regularizers.serialize', 'ab.v1.compt.keras.regularizers.serialize', 'import arrayblow as ab\n'), (191, 'arrayblow.v1.compt.keras.regularizers.serialize', 'ab.v1.compt.keras.regularizers.serialize', 'import arrayblow as ab\n')]