Fix warnings that occur when running tests.

[finbarrtimbers]

Currently, running the tests on my machine gives the following warnings:

=============================== warnings summary ===============================
llamax/integration_test.py::IntegrationTests::test_checkpoint_matches_torch
  /usr/local/lib/python3.10/site-packages/jax/_src/numpy/lax_numpy.py:5025: ComplexWarning: Casting complex values to real discards the imaginary part
    out_array: Array = lax_internal._convert_element_type(

llamax/integration_test.py::IntegrationTests::test_known_text_generation
llamax/integration_test.py::IntegrationTests::test_known_text_generation
  /usr/local/lib/python3.10/site-packages/jax/_src/ops/scatter.py:92: FutureWarning: scatter inputs have incompatible types: cannot safely cast value from dtype=int64 to dtype=int32 with jax_numpy_dtype_promotion='standard'. In future JAX releases this will result in an error.                             
    warnings.warn(                                           

-- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html

None of these are bad, but it'd be good practice to fix them up.