def call()

in tensorflow_recommenders/experimental/layers/embedding/partial_tpu_embedding.py [0:0]

• 19 lines of code
• 8 McCabe index (conditional complexity)

  def call(self, inputs: Dict[str, Tensor]) -> Dict[str, tf.Tensor]:
    """Computes the output of the embedding layer.

    It expects a string-to-tensor (or SparseTensor/RaggedTensor) dict as input,
    and outputs a dictionary of string-to-tensor of feature_name, embedded_value
    pairs. Note that SparseTensor/RaggedTensor are only supported for
    TPUEmbedding and are not supported for Keras embeddings.

    Args:
      inputs: A string-to-tensor (or SparseTensor/RaggedTensor) dictionary.

    Returns:
      output: A dictionary of string-to-tensor of feature_name, embedded_value
        pairs.

    Raises:
      ValueError if no tf.Tensor is passed to a Keras embedding layer.
    """
    keras_emb_inputs = {
        key: val for key, val in inputs.items()
        if key in self._keras_embedding_layers
    }
    tpu_emb_inputs = {
        key: val for key, val in inputs.items()
        if key not in self._keras_embedding_layers
    }

    output = {}
    for key, val in keras_emb_inputs.items():
      if not isinstance(val, tf.Tensor):
        raise ValueError("Only tf.Tensor input is supported for Keras embedding"
                         f" layers, but got: {type(val)}")

      output[key] = self._keras_embedding_layers[key](val)

    if self._tpu_embedding:
      tpu_emb_output_dict = self._tpu_embedding(tpu_emb_inputs)  # pylint: disable=[not-callable]
      output.update(tpu_emb_output_dict)
    return output