upgrade code for jax 0.8.5

examples/bring_in_your_own.pct.py

@@ -77,7 +77,7 @@ def __call__(self, x):
 from fortuna.utils.random import generate_rng_like_tree
 from jax.flatten_util import ravel_pytree
 from jax.tree_util import tree_map
-from jax._src.prng import PRNGKeyArray
+import jax
 import jax.numpy as jnp
 
 
@@ -86,7 +86,7 @@ def log_joint_prob(self, params: Params) -> float:
         v = jnp.mean((ravel_pytree(params)[0] <= 1) & (ravel_pytree(params)[0] >= 0))
         return jnp.where(v == 1.0, jnp.array(0), -jnp.inf)
 
-    def sample(self, params_like: Params, rng: Optional[PRNGKeyArray] = None) -> Params:
+    def sample(self, params_like: Params, rng: Optional[jax.Array] = None) -> Params:
         if rng is None:
             rng = self.rng.get()
         keys = generate_rng_like_tree(rng, params_like)

fortuna/calib_model/calib_model_calibrator.py

@@ -8,7 +8,7 @@
 )
 
 from flax.core import FrozenDict
-from jax._src.prng import PRNGKeyArray
+import jax
 import jax.numpy as jnp
 from optax._src.base import PyTree
 
@@ -35,7 +35,7 @@ def training_loss_step(
         params: Params,
         batch: Batch,
         mutable: Mutable,
-        rng: PRNGKeyArray,
+        rng: jax.Array,
         n_data: int,
         unravel: Optional[Callable[[any], PyTree]] = None,
         calib_params: Optional[CalibParams] = None,
@@ -71,7 +71,7 @@ def validation_step(
         state: CalibState,
         batch: Batch,
         loss_fun: Callable[[Any], Union[float, Tuple[float, dict]]],
-        rng: PRNGKeyArray,
+        rng: jax.Array,
         n_data: int,
         metrics: Optional[Tuple[Callable[[jnp.ndarray, Array], float], ...]] = None,
         unravel: Optional[Callable[[any], PyTree]] = None,

fortuna/calib_model/classification.py

@@ -124,7 +124,7 @@ def _get_model_manager(
                 )
             else:
                 model_manager = model_manager_cls(model, model_editor)
-        except ModuleNotFoundError as e:
+        except ModuleNotFoundError:
             logging.warning(
                 "No module named 'transformer' is installed. "
                 "If you are not working with models from the `transformers` library ignore this warning, otherwise "

fortuna/calib_model/config/monitor.py

@@ -66,7 +66,7 @@ def __init__(
                         f"All metrics in `metrics` must be callable objects, but {metric} is not."
                     )
         if uncertainty_fn is not None and not callable(uncertainty_fn):
-            raise ValueError(f"`uncertainty_fn` must be a a callable function.")
+            raise ValueError("`uncertainty_fn` must be a a callable function.")
 
         self.metrics = metrics
         self.uncertainty_fn = uncertainty_fn

fortuna/calib_model/loss.py

@@ -6,7 +6,7 @@
     Tuple,
 )
 
-from jax._src.prng import PRNGKeyArray
+import jax
 import jax.numpy as jnp
 
 from fortuna.likelihood.base import Likelihood
@@ -40,7 +40,7 @@ def __call__(
         return_aux: Optional[List[str]] = None,
         train: bool = False,
         outputs: Optional[jnp.ndarray] = None,
-        rng: Optional[PRNGKeyArray] = None,
+        rng: Optional[jax.Array] = None,
         **kwargs,
     ) -> Tuple[jnp.ndarray, Any]:
         if return_aux is None:

fortuna/calib_model/predictive/base.py

@@ -1,6 +1,6 @@
 from typing import Optional
 
-from jax._src.prng import PRNGKeyArray
+import jax
 import jax.numpy as jnp
 
 from fortuna.data.loader import (
@@ -60,7 +60,7 @@ def sample(
         self,
         inputs_loader: InputsLoader,
         n_samples: int = 1,
-        rng: Optional[PRNGKeyArray] = None,
+        rng: Optional[jax.Array] = None,
         distribute: bool = True,
     ) -> jnp.ndarray:
         r"""
@@ -80,7 +80,7 @@ def sample(
             A loader of input data points.
         n_samples : int
             Number of target samples to sample for each input data point.
-        rng : Optional[PRNGKeyArray]
+        rng : Optional[jax.Array]
             A random number generator. If not passed, this will be taken from the attributes of this class.
         distribute: bool
             Whether to distribute computation over multiple devices, if available.

fortuna/calib_model/predictive/regression.py

@@ -4,7 +4,7 @@
     Union,
 )
 
-from jax._src.prng import PRNGKeyArray
+import jax
 import jax.numpy as jnp
 
 from fortuna.calib_model.predictive.base import Predictive
@@ -21,7 +21,7 @@ def entropy(
         self,
         inputs_loader: InputsLoader,
         n_samples: int = 30,
-        rng: Optional[PRNGKeyArray] = None,
+        rng: Optional[jax.Array] = None,
         distribute: bool = True,
     ) -> jnp.ndarray:
         r"""
@@ -42,7 +42,7 @@ def entropy(
             A loader of input data points.
         n_samples : int
             Number of samples to draw for each input.
-        rng : Optional[PRNGKeyArray]
+        rng : Optional[jax.Array]
             A random number generator. If not passed, this will be taken from the attributes of this class.
         distribute: bool
             Whether to distribute computation over multiple devices, if available.
@@ -67,7 +67,7 @@ def quantile(
         q: Union[float, Array, List],
         inputs_loader: InputsLoader,
         n_samples: Optional[int] = 30,
-        rng: Optional[PRNGKeyArray] = None,
+        rng: Optional[jax.Array] = None,
         distribute: bool = True,
     ) -> Union[float, jnp.ndarray]:
         r"""
@@ -81,7 +81,7 @@ def quantile(
             A loader of input data points.
         n_samples : int
             Number of target samples to sample for each input data point.
-        rng: Optional[PRNGKeyArray]
+        rng: Optional[jax.Array]
             A random number generator. If not passed, this will be taken from the attributes of this class.
         distribute: bool
             Whether to distribute computation over multiple devices, if available.
@@ -109,7 +109,7 @@ def credible_interval(
         n_samples: int = 30,
         error: float = 0.05,
         interval_type: str = "two-tailed",
-        rng: Optional[PRNGKeyArray] = None,
+        rng: Optional[jax.Array] = None,
         distribute: bool = True,
     ) -> jnp.ndarray:
         r"""
@@ -126,7 +126,7 @@ def credible_interval(
             `error=0.05` corresponds to a 95% level of credibility.
         interval_type: str
             The interval type. We support "two-tailed" (default), "right-tailed" and "left-tailed".
-        rng : Optional[PRNGKeyArray]
+        rng : Optional[jax.Array]
             A random number generator. If not passed, this will be taken from the attributes of this class.
         distribute: bool
             Whether to distribute computation over multiple devices, if available.

fortuna/conformal/multivalid/mixins/classification/top_label_multicalibrator.py

@@ -1,6 +1,5 @@
 from typing import Optional
 
-from jax import vmap
 import jax.numpy as jnp
 
 from fortuna.conformal.multivalid.mixins.multicalibrator import MulticalibratorMixin

fortuna/conformal/multivalid/one_shot/base.py

@@ -1,9 +1,7 @@
 import abc
 import logging
 from typing import (
-    Dict,
     Optional,
-    Tuple,
     Union,
 )
 

fortuna/data/dataset/huggingface_datasets.py

@@ -15,9 +15,9 @@
     Dataset,
     DatasetDict,
 )
+import jax
 from jax import numpy as jnp
 import jax.random
-from jax.random import PRNGKeyArray
 from tqdm import tqdm
 from transformers import (
     BatchEncoding,
@@ -90,7 +90,7 @@ def get_data_loader(
         self,
         dataset: Dataset,
         per_device_batch_size: int,
-        rng: PRNGKeyArray,
+        rng: jax.Array,
         shuffle: bool = False,
         drop_last: bool = False,
         verbose: bool = False,
@@ -105,7 +105,7 @@ def get_data_loader(
             A tokenizeed dataset (see :meth:`.HuggingFaceClassificationDatasetABC.get_tokenized_datasets`).
         per_device_batch_size: bool
             Batch size for each device.
-        rng: PRNGKeyArray
+        rng: jax.Array
             Random number generator.
         shuffle: bool
             if True, shuffle the data so that each batch is a ranom sample from the dataset.
@@ -141,7 +141,7 @@ def _collate(self, batch: Dict[str, Array], batch_size: int) -> Dict[str, Array]
 
     @staticmethod
     def _get_batches_idxs(
-        rng: PRNGKeyArray,
+        rng: jax.Array,
         dataset_size: int,
         batch_size: int,
         shuffle: bool = False,
@@ -167,7 +167,7 @@ def _get_data_loader(
         batch_size: int,
         shuffle: bool = False,
         drop_last: bool = False,
-        rng: Optional[PRNGKeyArray] = None,
+        rng: Optional[jax.Array] = None,
         verbose: bool = False,
     ) -> Union[Iterable[Dict[str, Array]], Iterable[Tuple[Dict[str, Array], Array]]]:
         batch_idxs_gen = self._get_batches_idxs(
@@ -375,7 +375,7 @@ def __init__(
         super(HuggingFaceMaskedLMDataset, self).__init__(*args, **kwargs)
         if not self.tokenizer.is_fast:
             logger.warning(
-                f"You are not using a Fast Tokenizer, so whole words cannot be masked, only tokens."
+                "You are not using a Fast Tokenizer, so whole words cannot be masked, only tokens."
             )
         self.mlm = mlm
         self.mlm_probability = mlm_probability
@@ -407,7 +407,7 @@ def get_tokenized_datasets(
         ), "Only one text column should be passed when the task is MaskedLM."
 
         def _tokenize_fn(
-            batch: Dict[str, List[Union[str, int]]]
+            batch: Dict[str, List[Union[str, int]]],
         ) -> Dict[str, List[int]]:
             tokenized_inputs = self.tokenizer(
                 *[batch[col] for col in text_columns],

fortuna/distribution/gaussian.py

@@ -1,10 +1,10 @@
 from typing import Union
 
+import jax
 from jax import (
     random,
     vmap,
 )
-from jax._src.prng import PRNGKeyArray
 import jax.numpy as jnp
 from jax.scipy.stats import (
     multivariate_normal,
@@ -30,11 +30,11 @@ def __init__(self, mean: Union[float, Array], std: Union[float, Array]):
         self.std = std
         self.dim = 1 if type(mean) in [int, float] else len(mean)
 
-    def sample(self, rng: PRNGKeyArray, n_samples: int = 1) -> jnp.ndarray:
+    def sample(self, rng: jax.Array, n_samples: int = 1) -> jnp.ndarray:
         """
         Sample from the diagonal Gaussian.
 
-        :param rng: PRNGKeyArray
+        :param rng: jax.Array
             Random number generator.
         :param n_samples: int
             Number of samples.
@@ -72,11 +72,11 @@ def __init__(self, mean: Array, cov: Array):
         self.cov = cov
         self.dim = len(mean)
 
-    def sample(self, rng: PRNGKeyArray, n_samples: int = 1) -> jnp.ndarray:
+    def sample(self, rng: jax.Array, n_samples: int = 1) -> jnp.ndarray:
         """
         Sample from the multivariate Gaussian.
 
-        :param rng: PRNGKeyArray
+        :param rng: jax.Array
             Random number generator.
         :param n_samples: int
             Number of samples.

fortuna/hallucination/base.py

@@ -6,7 +6,6 @@
     List,
     Optional,
     Tuple,
-    Union,
 )
 
 import numpy as np

fortuna/likelihood/base.py

@@ -8,11 +8,11 @@
     Union,
 )
 
+import jax
 from jax import (
     jit,
     pmap,
 )
-from jax._src.prng import PRNGKeyArray
 import jax.numpy as jnp
 
 from fortuna.data.loader import (
@@ -133,7 +133,7 @@ def _batched_log_joint_prob(
         return_aux: Optional[List[str]] = None,
         train: bool = False,
         outputs: Optional[jnp.ndarray] = None,
-        rng: Optional[PRNGKeyArray] = None,
+        rng: Optional[jax.Array] = None,
         **kwargs,
     ) -> Union[jnp.ndarray, Tuple[jnp.ndarray, Any]]:
         """
@@ -161,7 +161,7 @@ def _batched_log_joint_prob(
             Whether the method is called during training.
         outputs : Optional[jnp.ndarray]
             Pre-computed batch of outputs.
-        rng: Optional[PRNGKeyArray]
+        rng: Optional[jax.Array]
             A random number generator. If not passed, this will be taken from the attributes of this class.
 
         Returns
@@ -272,7 +272,7 @@ def sample(
         calib_params: Optional[CalibParams] = None,
         calib_mutable: Optional[CalibMutable] = None,
         return_aux: Optional[List[str]] = None,
-        rng: Optional[PRNGKeyArray] = None,
+        rng: Optional[jax.Array] = None,
         distribute: bool = True,
         **kwargs,
     ) -> Union[jnp.ndarray, Tuple[jnp.ndarray, dict]]:
@@ -296,7 +296,7 @@ def sample(
         return_aux : Optional[List[str]]
             The auxiliary objects to return. We support 'outputs'. If this argument is not given, no auxiliary object
             is returned.
-        rng: Optional[PRNGKeyArray]
+        rng: Optional[jax.Array]
             A random number generator. If not passed, this will be taken from the attributes of this class.
         distribute: bool
             Whether to distribute computation over multiple devices, if available.
@@ -345,7 +345,7 @@ def _batched_sample(
         calib_params: Optional[CalibParams] = None,
         calib_mutable: Optional[CalibMutable] = None,
         return_aux: Optional[List[str]] = None,
-        rng: Optional[PRNGKeyArray] = None,
+        rng: Optional[jax.Array] = None,
         **kwargs,
     ) -> Union[jnp.ndarray, Tuple[jnp.ndarray, dict]]:
         if return_aux is None:

fortuna/likelihood/regression.py

@@ -3,8 +3,8 @@
     Union,
 )
 
+import jax
 from jax import vmap
-from jax._src.prng import PRNGKeyArray
 import jax.numpy as jnp
 import numpy as np
 
@@ -106,7 +106,7 @@ def entropy(
         calib_params: Optional[CalibParams] = None,
         calib_mutable: Optional[CalibMutable] = None,
         n_target_samples: Optional[int] = 30,
-        rng: Optional[PRNGKeyArray] = None,
+        rng: Optional[jax.Array] = None,
         distribute: bool = True,
         **kwargs,
     ) -> jnp.ndarray:
@@ -139,7 +139,7 @@ def quantile(
         calib_mutable: Optional[CalibMutable] = None,
         n_target_samples: Optional[int] = 30,
         target_samples: Optional[jnp.ndarray] = None,
-        rng: Optional[PRNGKeyArray] = None,
+        rng: Optional[jax.Array] = None,
         distribute: bool = True,
         **kwargs,
     ) -> Union[float, jnp.ndarray]:
@@ -164,7 +164,7 @@ def quantile(
             Number of target samples to sample for each input data point.
         target_samples: Optional[jnp.ndarray] = None
             Samples of the target variable for each input, used to estimate the quantiles.
-        rng: Optional[PRNGKeyArray]
+        rng: Optional[jax.Array]
             A random number generator. If not passed, this will be taken from the attributes of this class.
         distribute: bool
             Whether to distribute computation over multiple devices, if available.