def __init__(self, depth, nb_hyperplanes, nb_replicat=1, trainable=False):
    """Construct the gating function parameters.

    Compute the gates for a single head.

    Args:
      depth (int): Dimension of the key/queries to dispatch
      nb_hyperplanes (int): Nb of vectors use to split the space. Will determine
        the number of buckets (2^nb_hyperplanes - 1).
      nb_replicat (int): Redundancy to avoid the edge cases (to be in one bucket
        the input should be in a majority)
      trainable (bool): If True, a balance loss is added to force the hyperplane
        to divide the key/query space evenly
    """
    self.depth = depth
    self.nb_hyperplanes = nb_hyperplanes
    self.nb_buckets = 2**nb_hyperplanes
    self.nb_replicat = nb_replicat  # Unused for now
    self.trainable = trainable  # Unused for now

    self.dispatchers = {}

    assert self.nb_replicat == 1  # For now

    with tf.variable_scope("lsh_gating"):
      # Vectors defining the hyperplanes
      self.t_vectors = tf.get_variable(
          "vector",
          shape=(self.depth, self.nb_hyperplanes * self.nb_replicat),
          dtype=tf.float32,
          trainable=self.trainable,
      )
      # Projection vector from the bit space to similarity score space
      self.t_group = tf.constant(
          [self._idx_to_bits(i) for i in range(self.nb_buckets)],
          dtype=tf.float32,
          name="group")
Construct the gating function parameters.

    Compute the gates for a single head.

    Args:
      depth (int): Dimension of the key/queries to dispatch
      nb_hyperplanes (int): Nb of vectors use to split the space. Will determine
        the number of buckets (2^nb_hyperplanes - 1).
      nb_replicat (int): Redundancy to avoid the edge cases (to be in one bucket
        the input should be in a majority)
      trainable (bool): If True, a balance loss is added to force the hyperplane
        to divide the key/query space evenly
    
tensor2tensor/layers/common_attention.py