def METHOD_NAME(
    self,
    input,
    output,
    interpretable_embedding,
    embedding_dim=None,
    num_embeddings=None,
):
    if interpretable_embedding.embedding_dim is not None:
        self.assertEqual(embedding_dim, interpretable_embedding.embedding_dim)
        self.assertEqual(num_embeddings, interpretable_embedding.num_embeddings)
    # dim - [4, 100]
    emb_shape = interpretable_embedding.indices_to_embeddings(input).shape
    self.assertEqual(emb_shape[0], input.shape[0])
    if interpretable_embedding.embedding_dim is not None:
        self.assertEqual(emb_shape[1], interpretable_embedding.embedding_dim)
    self.assertEqual(input.shape[0], output.shape[0])