def METHOD_NAME(self):
    model = FlaxRoFormerForMaskedLM.from_pretrained("junnyu/roformer_chinese_base")
    input_ids = jnp.array([[0, 1, 2, 3, 4, 5]])
    output = model(input_ids)[0]
    vocab_size = 50000
    expected_shape = (1, 6, vocab_size)
    self.assertEqual(output.shape, expected_shape)
    expected_slice = jnp.array(
        [[[-0.1205, -1.0265, 0.2922], [-1.5134, 0.1974, 0.1519], [-5.0135, -3.9003, -0.8404]]]
    )
    self.assertTrue(jnp.allclose(output[:, :3, :3], expected_slice, atol=1e-4))