Support for Python 3.11 (#197)

.github/workflows/documentation.yml

@@ -32,7 +32,7 @@ jobs:
  - name: Install and configure Poetry
  uses: snok/install-poetry@v1
  with:
- version: 1.2.2
+ version: 1.8.3
  virtualenvs-create: true
  virtualenvs-in-project: false
  installer-parallel: true

.github/workflows/tests.yml

@@ -13,7 +13,7 @@ jobs:
  strategy:
  matrix:
  # Select the Python versions to test against
- python-version: ['3.8', '3.9', '3.10', '3.11']
+ python-version: ["3.9", "3.10", "3.11"]
  steps:
  - name: Check out the code
  uses: actions/checkout@v3
@@ -28,7 +28,7 @@ jobs:
  - name: Install Poetry
  uses: snok/install-poetry@v1.3.3
  with:
- version: 1.4.0
+ version: 1.8.3
 
  # Configure Poetry to use the virtual environment in the project
  - name: Setup Poetry
@@ -38,7 +38,7 @@ jobs:
  # Install the dependencies
  - name: Install Package
  run: |
- poetry install --with tests
+ poetry install --with test
 
  # Run the unit tests and build the coverage report
  - name: Run Tests

.readthedocs.yaml

@@ -4,7 +4,7 @@ version: 2
 build:
  os: ubuntu-20.04
  tools:
- python: '3.10'
+ python: '3.11'
  jobs:
  post_create_environment:
  # Install poetry

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