def _build(self, input_shape, dtype=tf.float32):
    """
    Called on the first iteration once the input shape is known
    :param input_shape: Input shape including batch size
    """
    with tf.variable_scope(self.name, reuse=tf.AUTO_REUSE):
      nonZeros = int(round(input_shape[-1].value * self.sparsity))

      # Create random mask with k elements set to 1, all other elements set to 0
      values = tf.random_uniform(input_shape)
      topk, _ = tf.math.top_k(input=values, k=nonZeros, sorted=False)
      kth = tf.reduce_min(topk, axis=1, keepdims=True)
      mask = tf.cast(tf.greater_equal(values, kth), dtype=dtype)
      self.mask = tf.get_variable(self.name, initializer=mask, trainable=False,
                                  synchronization=tf.VariableSynchronization.NONE)
      keras.backend.track_variable(self.mask)
      self._built = True
Called on the first iteration once the input shape is known
    :param input_shape: Input shape including batch size