def params_to_weights_and_biases(params, embed_dim):
  """Take a vector of parameters and return model weights and biases.

  Args:
    params: Vector of length `(embed_dim + 1) * num_classes`.
    embed_dim: Number of features in embedding.

  Returns:
    weights: `embed_dim x num_classes` matrix of weights.
    biases: Vector of length `num_classes` of biases.
  """
  params = params.reshape((embed_dim + 1), -1)
  return params[1:, :], params[0, :]