Add updated files after running pre-commit checks

.gitignore

@@ -151,4 +151,4 @@ package.json
 package-lock.json
 node_modules/
 
-docs/api
+docs/api

README.md

@@ -40,7 +40,7 @@ Another way you can contribute to GPJax is through [issue
 triaging](https://www.codetriage.com/what).  This can include reproducing bug reports,
 asking for vital information such as version numbers and reproduction instructions, or
 identifying stale issues. If you would like to begin triaging issues, an easy way to get
-started is to 
+started is to
 [subscribe to GPJax on CodeTriage](https://www.codetriage.com/jaxgaussianprocesses/gpjax).
 
 As a contributor to GPJax, you are expected to abide by our [code of

docs/GOVERNANCE.md

@@ -93,4 +93,4 @@ maintainers or reach out over
 -----
 
 This file was adapted from
-[BlackJAX](https://github.com/blackjax-devs/blackjax/blob/main/GOVERNANCE.md).
+[BlackJAX](https://github.com/blackjax-devs/blackjax/blob/main/GOVERNANCE.md).

gpjax/dataset.py

@@ -17,9 +17,9 @@
 import warnings
 
 from beartype.typing import (
+    Literal,
     Optional,
     Union,
-    Literal,
 )
 import jax.numpy as jnp
 from jaxtyping import (

gpjax/distributions.py

@@ -16,9 +16,9 @@
 
 from beartype.typing import (
     Any,
+    Generic,
     Optional,
     Tuple,
-    Generic,
     TypeVar,
     Union,
 )

gpjax/gps.py

@@ -15,7 +15,10 @@
 
 # from __future__ import annotations
 from abc import abstractmethod
-from dataclasses import dataclass, field
+from dataclasses import (
+    dataclass,
+    field,
+)
 from typing import overload
 
 from beartype.typing import (
@@ -25,7 +28,6 @@
 )
 import cola
 from cola.ops import Dense
-
 import jax.numpy as jnp
 from jax.random import (
     PRNGKey,
@@ -47,7 +49,10 @@
     ReshapedDistribution,
     ReshapedGaussianDistribution,
 )
-from gpjax.kernels import RFF, White
+from gpjax.kernels import (
+    RFF,
+    White,
+)
 from gpjax.kernels.base import AbstractKernel
 from gpjax.likelihoods import (
     AbstractLikelihood,

gpjax/kernels/computations/constant_diagonal.py

@@ -23,7 +23,10 @@
 )
 from jax import vmap
 import jax.numpy as jnp
-from jaxtyping import Float, Num
+from jaxtyping import (
+    Float,
+    Num,
+)
 
 from gpjax.kernels.computations import AbstractKernelComputation
 from gpjax.typing import Array

gpjax/kernels/computations/dense.py

@@ -15,7 +15,10 @@
 
 import beartype.typing as tp
 from jax import vmap
-from jaxtyping import Float, Num
+from jaxtyping import (
+    Float,
+    Num,
+)
 
 from gpjax.kernels.computations.base import AbstractKernelComputation
 from gpjax.typing import Array

gpjax/kernels/computations/diagonal.py

@@ -20,7 +20,10 @@
     LinearOperator,
 )
 from jax import vmap
-from jaxtyping import Float, Num
+from jaxtyping import (
+    Float,
+    Num,
+)
 
 from gpjax.kernels.computations import AbstractKernelComputation
 from gpjax.typing import Array

gpjax/kernels/stationary/matern32.py

@@ -17,7 +17,10 @@
 
 from beartype.typing import Union
 import jax.numpy as jnp
-from jaxtyping import Float, Num
+from jaxtyping import (
+    Float,
+    Num,
+)
 import tensorflow_probability.substrates.jax.bijectors as tfb
 import tensorflow_probability.substrates.jax.distributions as tfd
 

gpjax/kernels/stationary/matern52.py

@@ -17,7 +17,10 @@
 
 from beartype.typing import Union
 import jax.numpy as jnp
-from jaxtyping import Float, Num
+from jaxtyping import (
+    Float,
+    Num,
+)
 import tensorflow_probability.substrates.jax.bijectors as tfb
 import tensorflow_probability.substrates.jax.distributions as tfd
 

gpjax/kernels/stationary/periodic.py

@@ -17,7 +17,10 @@
 
 from beartype.typing import Union
 import jax.numpy as jnp
-from jaxtyping import Float, Num
+from jaxtyping import (
+    Float,
+    Num,
+)
 import tensorflow_probability.substrates.jax.bijectors as tfb
 
 from gpjax.base import param_field

gpjax/kernels/stationary/rbf.py

@@ -17,7 +17,10 @@
 
 from beartype.typing import Union
 import jax.numpy as jnp
-from jaxtyping import Float, Num
+from jaxtyping import (
+    Float,
+    Num,
+)
 import tensorflow_probability.substrates.jax.bijectors as tfb
 import tensorflow_probability.substrates.jax.distributions as tfd
 

gpjax/likelihoods.py

@@ -29,7 +29,10 @@
     param_field,
     static_field,
 )
-from gpjax.distributions import GaussianDistribution, ReshapedDistribution
+from gpjax.distributions import (
+    GaussianDistribution,
+    ReshapedDistribution,
+)
 from gpjax.integrators import (
     AbstractIntegrator,
     AnalyticalGaussianIntegrator,

gpjax/objectives.py

@@ -20,7 +20,6 @@
     ScalarFloat,
 )
 
-
 tfd = tfp.distributions
 
 import cola

tests/conftest.py

@@ -1,5 +1,6 @@
 from jax import config
 from jaxtyping import install_import_hook
+
 config.update("jax_enable_x64", True)
 
 # import gpjax within import hook to apply beartype everywhere, before running tests