text: # coding=utf-8
  # Copyright 2020 The Google Research 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.
  
  # Lint as: python3
  """Defines custom layers."""
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import numpy as np
  import tensorflow.compat.v2 as tf
  
  
  def soft_hot_layer(offsets, scales):
    """Embedding layer for scalar observations.
  
    For each input scalar, produces an embedding vector.
  
    Args:
      offsets: Matrix of shape `[input_dim, soft_hot_dim]` which contains the
          soft-hot "offset" for the i-th scalar input and j-th soft-hot feature.
          `input_dim` is the number of scalar inputs, and `soft_hot_dim` is the
          number of soft-hot features (output dim).
      scales: Array of shape `[input_dim]` which contains the soft-hot scales, one
          for each input scalar.
  
    Returns:
      tf.keras.layers.Layer object. This layer converts a tensor of shape
      `[..., input_dim]` to a tensor of shape `[..., input_dim*soft_hot_dim]`,
      which is a concat of the soft-hot embeddings for each input scalar.
    """
    offsets = np.asarray(offsets)
    scales = np.asarray(scales)
    assert offsets.ndim == 2
    assert scales.ndim == 1
    length = offsets.size
    offsets = offsets[np.newaxis, Ellipsis]
    scales = scales[np.newaxis, :, np.newaxis]
  
    def _f(x):
      return tf.reshape(
          tf.exp(-(tf.expand_dims(x, -1) - offsets)**2/scales),
          (tf.shape(x)[0], length))
  
    return tf.keras.layers.Lambda(_f)
  
  
  def stretch(low, high, fraction):
    """Stretch out interval by the given fraction."""
    delta = abs(high - low)*fraction
    return (low - delta, high + delta)
  
  
  def obs_embedding_kwargs(num_features, batch=None, bounds=None, variance=None,
                           spillover=0.05):
    """Helper function that generates arguments to `soft_hot_layer`.
  
    Usages
    1) obs_embedding_kwargs(num_features=N, batch=<2D tensor>)
    2) obs_embedding_kwargs(
           num_features=N,
           bounds=[(low_0, hi_0), (low_1, hi_1), ..., (low_K, hi_K)],
           variance=[v_0, v_1, ..., v_K])
  
    Args:
      num_features: Number of features in the generated embedding. Output size of
          the layer returned by this function.
      batch: (usage 1) A batch of data, with which to extract the input feature
          statistics. Automatically calculates `bounds` and `variance` arguments
          under the hood from `batch`.
      bounds: (usage 2) List of 2-tuples, giving the min and max values expected
          for each input feature. Values out of bounds are allowed, but may be
          difficult to distinguish from min and max values. Set higher `spillover`
          to accomidate out-of-bounds values.
      variance: (usage 2) List of variances of each input feature. This sets the
          width of the Gaussian kernels used by `soft_hot_layer`.
      spillover: Widens the sensitivity range.
  
    Returns:
      Keyword arguments for `soft_hot_layer`.
  
    """
    if batch is not None:
      assert bounds is None
      assert variance is None
      assert batch.ndim == 2
      bounds = list(zip(np.min(batch, axis=0), np.max(batch, axis=0)))
      variance = np.var(batch, axis=0)
    assert bounds is not None
    assert variance is not None
    dim = len(list(zip(*bounds))[0])  # Shape of the input (number of scalars).
    assert len(variance) == dim
    nf = num_features  # Number of soft-hot features.
    return {
        'offsets': np.stack(
            [np.linspace(*stretch(bounds[i][0], bounds[i][1], spillover), nf)
             for i in range(dim)],
            axis=0),  # shape: [dim, nF]
        'scales': variance,  # shape: [dim]
    }
  
  