def METHOD_NAME(self):
    model = FlaxRobertaPreLayerNormModel.from_pretrained("andreasmadsen/efficient_mlm_m0.40", from_pt=True)
    input_ids = np.array([[0, 31414, 232, 328, 740, 1140, 12695, 69, 46078, 1588, 2]], dtype=jnp.int32)
    output = model(input_ids)[0]
    # compare the actual values for a slice.
    EXPECTED_SLICE = np.array(
        [[[0.0208, -0.0356, 0.0237], [-0.1569, -0.0411, -0.2626], [0.1879, 0.0125, -0.0089]]], dtype=np.float32
    )
    self.assertTrue(np.allclose(output[:, :3, :3], EXPECTED_SLICE, atol=1e-4))