def shared_features_attr_or_input_features(self, x_embed_features: int) -> int:
  """Use class attribute for the embedding features, back off to matching.

  Note:
    This function is basically a method that is set as a class attribute to
    avoid issues of JAX returning tracer values for this function.

    This assumes there is a `shared_features`, H_s, attribute on your module.

  Args:
    self: The Prompt Module.
    x_embed_features: The number of features in the embedding, H.

  Returns:
    The number of features in the shared prompt.
  """
  return x_embed_features if self.shared_features is None else self.shared_features