self.r_w_bias = (
          self.add_weight(
              "r_w_bias",
              shape=[self.n_head, self.d_head],
              dtype=self.tf_float,