Fix import of PjitFunction

[stephen-huan]

Checklist

• I've formatted the new code by running hatch run dev:format before committing.
• I've added tests for new code.
• I've added docstrings for the new code.

Description

The first commit is unrelated.

The import from jaxlib.xla_extension import PjitFunction no longer works, presumably as part of the deprecations of 0.6.0. Instead, default to the new location of from jaxlib._jax import PjitFunction.

Also, the default citation is broken as it currently overwrites __str__, but __repr__ uses as_str.

>>> import gpjax
>>> gpjax.cite(None)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File ".../gpjax/citation.py", line 45, in __repr__
    return repr(self.as_str())
                ^^^^^^^^^^^^^
  File ".../gpjax/citation.py", line 37, in as_str
    citation_str = f"@{self.citation_type}{{{self.citation_key},"
                       ^^^^^^^^^^^^^^^^^^
AttributeError: 'NullCitation' object has no attribute 'citation_type'. Did you mean: 'citation_key'?

Note that this only affects __repr__ and not __str__ as the following demonstrates.

import gpjax

print(gpjax.cite(None))  # equivalent to str(gpjax.cite(None))
print(str(gpjax.cite(None)))  # prints the same thing as above
print(repr(gpjax.cite(None)))  # error

No citation available. If you think this is an error, please open a pull request.
No citation available. If you think this is an error, please open a pull request.
Traceback (most recent call last):
  File "...", line 5, in <module>
    print(repr(gpjax.cite(None)))  # error
          ^^^^^^^^^^^^^^^^^^^^^^
  File ".../gpjax/citation.py", line 45, in __repr__
    return repr(self.as_str())
                ^^^^^^^^^^^^^
  File ".../gpjax/citation.py", line 37, in as_str
    citation_str = f"@{self.citation_type}{{{self.citation_key},"
                       ^^^^^^^^^^^^^^^^^^
AttributeError: 'NullCitation' object has no attribute 'citation_type'. Did you mean: 'citation_key'?

Issue Number: Fixes #518, cf. #519