fisrt commit

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2021 Morgan McGuire
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -1 +1,69 @@
-An attempt at making i Bart pretraining script. Based on code for RotoBART.
+# RotoBART
+
+## Running the script
+
+### Script arguemnts
+
+Available model config arguments from script:
+```
+encoder_layers
+encoder_ffn_dim
+decoder_layers
+decoder_ffn_dim
+d_model
+vocab_size
+max_position_embeddings
+encoder_layerdrop
+decoder_layerdrop
+```
+
+Training Arguments:
+
+`testing` : only uses 1 batch, for testing the script
+
+`adafactor`: will enable adafactor, removing the command will revert to Adam
+
+`grad_accum`: what value for gradient accumulation to use, default is 4
+
+`use_bf16`: convert the model to bf16
+
+`colab_tpu`: if running on a colab TPU
+
+`use_wandb`: log using Weights & Biases (via Tensorboard)
+
+`save_strategy`: whether or not to save model checkpoints based on steps or epoch
+
+
+```
+python rotobart/run_dnlm_flax.py \
+  --output_dir rotobart_output \
+  --overwrite_output_dir \
+  --dataset_path rotobart/pile.py \
+  --model_name_or_path rotobart \
+  --tokenizer_name ./rotobart/vocab-2/the_pile.model \
+  --shuffle_buffer_size 1000 \
+  --do_train --do_eval \
+  --max_seq_length 1024 \
+  --encoder_layers 2 \
+  --decoder_layers 2 \
+  --per_device_train_batch_size 2 \
+  --per_device_eval_batch_size 2 \
+  --logging_steps 8 \
+  --num_train_steps 1000 \
+  --eval_steps 1000 \
+  --save_steps 1000 \
+  --save_strategy steps \
+  --num_eval_samples 100 \
+  --warmup_steps 30 \
+  --learning_rate 1e-4 \
+  --use_wandb \
+  --testing \
+  --use_bf16 \
+  --adafactor
+```
+
+alt
+
+```
+python3 run_dnlm_flax.py   --output_dir rotobart_output   --overwrite_output_dir   --dataset_path pile.py   --model_name_or_path rotobart   --tokenizer_name vocab-2/the_pile.model   --shuffle_buffer_size 1000   --do_train --do_eval --max_position_embeddings 2048  --max_seq_length 2048   --encoder_layers 6   --decoder_layers 6   --per_device_train_batch_size 1   --per_device_eval_batch_size 1   --logging_steps 100   --num_train_steps 50000   --eval_steps 2500   --save_steps 2500   --save_strategy steps   --num_eval_samples 5000   --warmup_steps 5000   --learning_rate 1e-4   --use_wandb   --use_bf16   --adafactor
+```

build_env.sh

@@ -0,0 +1,6 @@
+export USE_TORCH=False
+python3 -m venv jax_env
+source ~/jax_env/bin/activate
+pip install --upgrade pip
+pip install "jax[tpu]>=0.2.16" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
+pip install transformers datasets einops flax aiohttp jsonlines zstandard wandb SentencePiece nltk prefetch_generator torch tensorboard tensorflow

configuration_rotobart.py

@@ -0,0 +1,187 @@
+# coding=utf-8
+# Copyright 2021 The Fairseq Authors and The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+""" RotoBART model configuration """
+import warnings
+
+from transformers.configuration_utils import PretrainedConfig
+from transformers.utils import logging
+
+logger = logging.get_logger(__name__)
+
+RotoBART_PRETRAINED_CONFIG_ARCHIVE_MAP = {
+    "facebook/RotoBART-large": "https://huggingface.co/facebook/RotoBART-large/resolve/main/config.json",
+    # See all RotoBART models at https://huggingface.co/models?filter=RotoBART
+}
+
+
+class RotoBARTConfig(PretrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a :class:`~transformers.RotoBARTModel`. It is used to
+    instantiate a RotoBART model according to the specified arguments, defining the model architecture. Instantiating a
+    configuration with the defaults will yield a similar configuration to that of the RotoBART `facebook/RotoBART-large
+    <https://huggingface.co/facebook/RotoBART-large>`__ architecture.
+
+    Configuration objects inherit from :class:`~transformers.PretrainedConfig` and can be used to control the model
+    outputs. Read the documentation from :class:`~transformers.PretrainedConfig` for more information.
+
+
+    Args:
+        vocab_size (:obj:`int`, `optional`, defaults to 50265):
+            Vocabulary size of the RotoBART model. Defines the number of different tokens that can be represented by the
+            :obj:`inputs_ids` passed when calling :class:`~transformers.RotoBARTModel` or
+            :class:`~transformers.TFRotoBARTModel`.
+        d_model (:obj:`int`, `optional`, defaults to 1024):
+            Dimensionality of the layers and the pooler layer.
+        encoder_layers (:obj:`int`, `optional`, defaults to 12):
+            Number of encoder layers.
+        decoder_layers (:obj:`int`, `optional`, defaults to 12):
+            Number of decoder layers.
+        encoder_attention_heads (:obj:`int`, `optional`, defaults to 16):
+            Number of attention heads for each attention layer in the Transformer encoder.
+        decoder_attention_heads (:obj:`int`, `optional`, defaults to 16):
+            Number of attention heads for each attention layer in the Transformer decoder.
+        decoder_ffn_dim (:obj:`int`, `optional`, defaults to 4096):
+            Dimensionality of the "intermediate" (often named feed-forward) layer in decoder.
+        encoder_ffn_dim (:obj:`int`, `optional`, defaults to 4096):
+            Dimensionality of the "intermediate" (often named feed-forward) layer in decoder.
+        activation_function (:obj:`str` or :obj:`function`, `optional`, defaults to :obj:`"gelu"`):
+            The non-linear activation function (function or string) in the encoder and pooler. If string,
+            :obj:`"gelu"`, :obj:`"relu"`, :obj:`"silu"` and :obj:`"gelu_new"` are supported.
+        dropout (:obj:`float`, `optional`, defaults to 0.1):
+            The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
+        attention_dropout (:obj:`float`, `optional`, defaults to 0.0):
+            The dropout ratio for the attention probabilities.
+        activation_dropout (:obj:`float`, `optional`, defaults to 0.0):
+            The dropout ratio for activations inside the fully connected layer.
+        classifier_dropout (:obj:`float`, `optional`, defaults to 0.0):
+            The dropout ratio for classifier.
+        max_position_embeddings (:obj:`int`, `optional`, defaults to 1024):
+            The maximum sequence length that this model might ever be used with. Typically set this to something large
+            just in case (e.g., 512 or 1024 or 2048).
+        init_std (:obj:`float`, `optional`, defaults to 0.02):
+            The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
+        encoder_layerdrop: (:obj:`float`, `optional`, defaults to 0.0):
+            The LayerDrop probability for the encoder. See the `LayerDrop paper <see
+            https://arxiv.org/abs/1909.11556>`__ for more details.
+        decoder_layerdrop: (:obj:`float`, `optional`, defaults to 0.0):
+            The LayerDrop probability for the decoder. See the `LayerDrop paper <see
+            https://arxiv.org/abs/1909.11556>`__ for more details.
+        gradient_checkpointing (:obj:`bool`, `optional`, defaults to :obj:`False`):
+            If True, use gradient checkpointing to save memory at the expense of slower backward pass.
+        scale_embedding (:obj:`bool`, `optional`, defaults to :obj:`False`):
+            Scale embeddings by diving by sqrt(d_model).
+        use_cache (:obj:`bool`, `optional`, defaults to :obj:`True`):
+            Whether or not the model should return the last key/values attentions (not used by all models).
+        num_labels: (:obj:`int`, `optional`, defaults to 3):
+            The number of labels to use in :class:`~transformers.RotoBARTForSequenceClassification`.
+        forced_eos_token_id (:obj:`int`, `optional`, defaults to 2):
+            The id of the token to force as the last generated token when :obj:`max_length` is reached. Usually set to
+            :obj:`eos_token_id`.
+
+    Example::
+
+        >>> from transformers import RotoBARTModel, RotoBARTConfig
+
+        >>> # Initializing a RotoBART facebook/RotoBART-large style configuration
+        >>> configuration = RotoBARTConfig()
+
+        >>> # Initializing a model from the facebook/RotoBART-large style configuration
+        >>> model = RotoBARTModel(configuration)
+
+        >>> # Accessing the model configuration
+        >>> configuration = model.config
+    """
+    model_type = "RotoBART"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    def __init__(
+        self,
+        vocab_size=50265,
+        max_position_embeddings=1024,
+        encoder_layers=12,
+        encoder_ffn_dim=4096,
+        encoder_attention_heads=16,
+        decoder_layers=12,
+        decoder_ffn_dim=4096,
+        decoder_attention_heads=16,
+        encoder_layerdrop=0.0,
+        decoder_layerdrop=0.0,
+        activation_function="gelu",
+        d_model=1024,
+        dropout=0.1,
+        attention_dropout=0.0,
+        activation_dropout=0.0,
+        init_std=0.02,
+        classifier_dropout=0.0,
+        scale_embedding=False,
+        gradient_checkpointing=False,
+        use_cache=True,
+        num_labels=3,
+        pad_token_id=1,
+        bos_token_id=0,
+        eos_token_id=2,
+        is_encoder_decoder=True,
+        decoder_start_token_id=2,
+        forced_eos_token_id=2,
+        **kwargs,
+    ):
+        super().__init__(
+            num_labels=num_labels,
+            pad_token_id=pad_token_id,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            is_encoder_decoder=is_encoder_decoder,
+            decoder_start_token_id=decoder_start_token_id,
+            forced_eos_token_id=forced_eos_token_id,
+            **kwargs,
+        )
+
+        self.vocab_size = vocab_size
+        self.max_position_embeddings = max_position_embeddings
+        self.d_model = d_model
+        self.encoder_ffn_dim = encoder_ffn_dim
+        self.encoder_layers = encoder_layers
+        self.encoder_attention_heads = encoder_attention_heads
+        self.decoder_ffn_dim = decoder_ffn_dim
+        self.decoder_layers = decoder_layers
+        self.decoder_attention_heads = decoder_attention_heads
+        self.dropout = dropout
+        self.attention_dropout = attention_dropout
+        self.activation_dropout = activation_dropout
+        self.activation_function = activation_function
+        self.init_std = init_std
+        self.encoder_layerdrop = encoder_layerdrop
+        self.decoder_layerdrop = decoder_layerdrop
+        self.classifier_dropout = classifier_dropout
+        self.use_cache = use_cache
+        self.num_hidden_layers = encoder_layers
+        self.gradient_checkpointing = gradient_checkpointing
+        self.scale_embedding = scale_embedding  # scale factor will be sqrt(d_model) if True
+
+        # ensure backward compatibility for RotoBART CNN models
+        if self.forced_bos_token_id is None and kwargs.get("force_bos_token_to_be_generated", False):
+            self.forced_bos_token_id = self.bos_token_id
+            warnings.warn(
+                f"Please make sure the config includes `forced_bos_token_id={self.bos_token_id}` in future versions."
+                "The config can simply be saved and uploaded again to be fixed."
+            )
+
+    @property
+    def num_attention_heads(self) -> int:
+        return self.encoder_attention_heads
+
+    @property
+    def hidden_size(self) -> int:
+        return self.d_model

data_collator.py

@@ -0,0 +1,390 @@
+import math
+from dataclasses import dataclass
+from typing import Dict, List, Optional
+
+import nltk
+import numpy as np
+from numpy.random import permutation, poisson
+from transformers.data.data_collator import _torch_collate_batch
+from transformers.tokenization_utils_base import BatchEncoding, PreTrainedTokenizerBase
+
+nltk.download("punkt")
+
+
+@dataclass
+class DataCollatorForTextInfilling:
+    tokenizer: PreTrainedTokenizerBase
+    mlm_probability: float = 0.15
+    poisson_lambda: float = 3.0
+    pad_to_multiple_of: Optional[int] = None
+
+    def __post_init__(self):
+        if self.tokenizer.mask_token is None:
+            raise ValueError
+
+    def __call__(self, examples: List[Dict[str, np.ndarray]]) -> Dict[str, np.ndarray]:
+        # Handle dict or lists with proper padding and conversion to tensor.
+        batch = {}
+        if isinstance(examples, (dict, BatchEncoding)):
+            examples_ids = examples["input_ids"]
+            if "decoder_input_ids" in examples.keys():
+                examples_dec = examples["decoder_input_ids"]
+            else:
+                examples_dec = examples_ids
+
+            # bs of one
+            if type(examples_ids[0]) is int:
+                examples_ids = [examples_ids]
+            # bs of one
+            if type(examples_dec[0]) is int:
+                examples_dec = [examples_dec]
+
+            batch["input_ids"] = _torch_collate_batch(
+                examples_ids, self.tokenizer, pad_to_multiple_of=self.pad_to_multiple_of
+            )
+            batch["decoder_input_ids"] = _torch_collate_batch(
+                examples_dec, self.tokenizer, pad_to_multiple_of=self.pad_to_multiple_of
+            )
+            batch["decoder_input_ids"] = batch["decoder_input_ids"].tolist()
+
+        elif isinstance(examples[0], (dict, BatchEncoding)):
+            batch = self.tokenizer.pad(examples, return_tensors="jax", pad_to_multiple_of=self.pad_to_multiple_of)
+        else:
+            batch["input_ids"] = _torch_collate_batch(examples, self.tokenizer, pad_to_multiple_of=self.pad_to_multiple_of)
+            batch["decoder_input_ids"] = _torch_collate_batch(
+                examples, self.tokenizer, pad_to_multiple_of=self.pad_to_multiple_of
+            ).tolist()
+
+        # If special token mask has been preprocessed, pop it from the dict.
+        special_tokens_mask = batch.pop("special_tokens_mask", None)
+
+        batch["input_ids"], batch["labels"] = self.mask_tokens(
+            batch["input_ids"], special_tokens_mask=special_tokens_mask
+        )
+
+        return batch
+
+    def mask_tokens(self, inputs):
+        inputs_copy = np.array(inputs)
+        labels = np.array(inputs)
+        if special_tokens_mask is None:
+            special_tokens_mask = [
+                self.tokenizer.get_special_tokens_mask(val, already_has_special_tokens=True) for val in labels.tolist()
+            ]
+            special_tokens_mask = jnp.array(special_tokens_mask, dtype=bool)
+        else:
+            special_tokens_mask = special_tokens_mask.bool()
+
+        # determine how many tokens we need to mask in total
+        is_token = ~(labels == self.tokenizer.pad_token_id) & ~special_tokens_mask
+        num_to_mask = int(math.ceil(is_token.astype(float).sum() * self.mlm_probability))
+        if num_to_mask == 0:
+            return inputs, labels
+
+        # generate a sufficient number of span lengths
+        lengths = poisson(lam=self.poisson_lambda, size=(num_to_mask,))
+        while np.cumsum(lengths, 0)[-1] < num_to_mask:
+            lengths = np.concatenate([lengths, poisson(lam=self.poisson_lambda, size=(num_to_mask,))])
+
+        # remove all spans of length 0
+        # Note that BART inserts additional mask tokens where length == 0,
+        # which we do not implement for now as it adds additional complexity
+        lengths = lengths[lengths > 0]
+
+        # trim to about num_to_mask tokens
+        idx = np.argmin(np.abs(np.cumsum(lengths, 0) - num_to_mask)) + 1
+        lengths = lengths[: idx + 1]
+
+        # select span start indices
+        # print("IS TOKEN")
+        # print(is_token)
+        # print(sum(list(map(lambda x: 1 if(x) else 0, is_token[0]))))
+        token_indices = np.argwhere(is_token == 1)
+        # print("TOKEN INDICES")
+        # print(token_indices)
+        span_starts = permutation(token_indices.shape[0])[: lengths.shape[0]]
+
+        # prepare mask
+        masked_indices = np.array(token_indices[span_starts])
+        # print("MASKED INDICES")
+        # print(masked_indices)
+        mask = np.full_like(labels, fill_value=False)
+
+        # mask span start indices
+        for mi in masked_indices:
+            mask[tuple(mi)] = True
+        lengths -= 1
+
+        # fill up spans
+        max_index = labels.shape[1] - 1
+        remaining = (lengths > 0) & (masked_indices[:, 1] < max_index)
+        while np.any(remaining):
+            masked_indices[remaining, 1] += 1
+            for mi in masked_indices:
+                mask[tuple(mi)] = True
+            lengths -= 1
+            remaining = (lengths > 0) & (masked_indices[:, 1] < max_index)
+
+        # place the mask tokens
+        mask[np.where(special_tokens_mask == True)] = False
+        inputs_copy[np.where(mask == 1)] = self.tokenizer.mask_token_id
+        labels[np.where(mask == 0)] = -100
+
+        # remove mask tokens that are not starts of spans
+        to_remove = (mask == 1) & np.roll((mask == 1), 1, 1)
+        new_inputs = np.full_like(labels, fill_value=self.tokenizer.pad_token_id)
+
+        # splits = list(map(lambda x: x.reshape(-1),  np.split(inputs_copy, indices_or_sections=2, axis=0))
+        for i, example in enumerate(np.split(inputs_copy, indices_or_sections=new_inputs.shape[0], axis=0)):
+            new_example = example[0][~to_remove[i]]
+            new_inputs[i, 0 : new_example.shape[0]] = new_example
+
+        # batching now fixed
+        return new_inputs.tolist(), labels.tolist()
+
+
+# Code below is by Matt Bui
+@dataclass
+class SentenceTokenize:
+    """Tokenize documents into sentences, add bos and eos tokens and split sentences into smaller chunks if too long."""
+
+    sentence_tokenizer = nltk.data.load("tokenizers/punkt/english.pickle")
+    bos: str = "<s>"
+    eos: str = "</s>"
+    max_sentences = 256
+    sentence_stride = 128
+    max_characters = 100000
+
+    def __call__(self, examples: Dict[str, List[str]]) -> Dict[str, List[str]]:
+        is_batched = isinstance(examples["text"], list)
+        if not is_batched:
+            # raise ValueError("required batched=True in map() method")
+            examples["text"] = [examples["text"]]
+
+        texts = []
+#         print(f"len(examples['text'] : {len(examples['text'])}")
+        for doc in examples["text"]:
+            sentences = self.sentence_tokenizer.tokenize(doc)
+            start_index = 0
+#             print(f"doc len: {len(doc)}")
+#             print(f"sent_tok len: {len(sentences)}")
+            while start_index < len(sentences):
+                sentence_span = sentences[start_index : min(len(sentences), start_index + self.max_sentences)]
+                text = f"{self.eos}{self.bos}".join([sentence for sentence in sentence_span])
+
+                # trim text by max characters
+                if len(text) > self.max_characters:
+                    text = text[: self.max_characters]
+                texts.append(text)
+                start_index += self.sentence_stride
+#         print(len(texts))
+#         print()      
+        return {"text": texts}
+
+
+@dataclass
+class DataCollatorForSentencePermutation:
+    tokenizer: PreTrainedTokenizerBase
+    permutate_sentence_ratio: float = 1.0
+
+    def __post_init__(self):
+        self.full_stop_index = self.tokenizer.eos_token_id
+
+    def __call__(self, example: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]:
+        source = example["input_ids"]
+
+        full_stops = source == self.full_stop_index
+
+        # Tokens that are full stops, where the previous token is not
+        sentence_ends = (full_stops[1:] * ~full_stops[:-1]).nonzero()[0] + 2
+        result = source.copy()
+
+        num_sentences = jnp.size(sentence_ends, 0)
+        num_to_permute = math.ceil((num_sentences * 2 * self.permutate_sentence_ratio) / 2.0)
+        substitutions = random.permutation(self.random_key, num_sentences)[:num_to_permute]
+        ordering = jnp.arange(0, num_sentences)
+        ordering = ops.index_update(
+            ordering, substitutions, substitutions[random.permutation(self.random_key, num_to_permute)]
+        )
+
+        index = 0
+        for i in ordering:
+            sentence = source[(sentence_ends[i - 1] if i > 0 else 0) : sentence_ends[i]]
+            result = ops.index_update(result, ops.index[index : index + jnp.size(sentence, 0)], sentence)
+            index += jnp.size(sentence, 0)
+
+        example["decoder_input_ids"] = example["input_ids"]
+        example["input_ids"] = result
+
+        return example
+
+
+@dataclass
+class DataCollatorForDenoisingTasks:
+    """Data collator used denoising language modeling task in BART.
+    The implementation is based on
+    https://github.com/pytorch/fairseq/blob/1bba712622b8ae4efb3eb793a8a40da386fe11d0/fairseq/data/denoising_dataset.py.
+    The default paramters is based on BART paper https://arxiv.org/abs/1910.13461.
+    """
+
+    tokenizer: PreTrainedTokenizerBase
+    mask_ratio: float = 0.3
+    poisson_lambda: float = 3.0
+    permutate_sentence_ratio: float = 1.0
+    pad_to_multiple_of: int = 16
+
+    def __post_init__(self):
+        if self.tokenizer.mask_token is None or self.tokenizer.eos_token is None:
+            raise ValueError
+
+    def __call__(self, examples: List[Dict[str, List[int]]]) -> Dict[str, np.ndarray]:
+        """Batching, adding whole word mask and permutate sentences
+        Args:
+            examples (dict): list of examples each examples contains input_ids field
+        """
+        # Handle dict or lists with proper padding and conversion to tensor.
+        batch = self.tokenizer.pad(examples, pad_to_multiple_of=self.pad_to_multiple_of, return_tensors="np")
+        batch["decoder_input_ids"] = self.shift_tokens_right(batch["input_ids"])
+
+        do_permutate = False
+        if self.permutate_sentence_ratio > 0.0:
+            batch["input_ids"] = self.permutate_sentences(batch["input_ids"])
+            do_permutate = True
+
+        if self.mask_ratio:
+            batch["input_ids"], batch["labels"] = self.add_whole_word_mask(batch["input_ids"], do_permutate)
+
+        return batch
+
+    def shift_tokens_right(self, inputs):
+        """Shift decoder input ids right: https://github.com/huggingface/transformers/issues/7961.
+        Examples:
+            <s>My dog is cute.</s><s>It loves to play in the park.</s><pad><pad>
+            shift to -> </s><s>My dog is cute.</s><s>It loves to play in the park.<pad><pad>
+        """
+
+        shifted_inputs = np.roll(inputs, 1, axis=-1)
+
+        # replace first token with eos token
+        shifted_inputs[:, 0] = self.tokenizer.eos_token_id
+
+        # when there's padding, the last eos tokens will not be rotate to first positon
+        # we'll need to replace it with a padding token
+
+        # replace eos tokens at the end of sequences with pad tokens
+        end_with_eos = np.where(shifted_inputs[:, -1] == self.tokenizer.eos_token_id)
+        shifted_inputs[end_with_eos, -1] = self.tokenizer.pad_token_id
+
+        # find positions where where's the token is eos and its follwing token is a padding token
+        last_eos_indices = np.where(
+            (shifted_inputs[:, :-1] == self.tokenizer.eos_token_id)
+            * (shifted_inputs[:, 1:] == self.tokenizer.pad_token_id)
+        )
+
+        # replace eos tokens with pad token
+        shifted_inputs[last_eos_indices] = self.tokenizer.pad_token_id
+        return shifted_inputs
+
+    def permutate_sentences(self, inputs):
+        results = inputs.copy()
+
+        full_stops = inputs == self.tokenizer.eos_token_id
+
+        sentence_ends = np.argwhere(full_stops[:, 1:] * ~full_stops[:, :-1])
+        sentence_ends[:, 1] += 2
+        num_sentences = np.unique(sentence_ends[:, 0], return_counts=True)[1]
+        num_to_permute = np.ceil((num_sentences * 2 * self.permutate_sentence_ratio) / 2.0).astype(int)
+
+        sentence_ends = np.split(sentence_ends[:, 1], np.unique(sentence_ends[:, 0], return_index=True)[1][1:])
+
+        for i in range(inputs.shape[0]):
+            substitutions = np.random.permutation(num_sentences[i])[: num_to_permute[i]]
+
+            ordering = np.arange(0, num_sentences[i])
+            ordering[substitutions] = substitutions[np.random.permutation(num_to_permute[i])]
+
+            index = 0
+            for j in ordering:
+                sentence = inputs[i, (sentence_ends[i][j - 1] if j > 0 else 0) : sentence_ends[i][j]]
+                results[i, index : index + sentence.shape[0]] = sentence
+                index += sentence.shape[0]
+        return results
+
+    def add_whole_word_mask(self, inputs, do_permutate):
+        labels = inputs.copy()
+
+        special_tokens_mask = [
+            self.tokenizer.get_special_tokens_mask(val, already_has_special_tokens=True) for val in labels.tolist()
+        ]
+        special_tokens_mask = np.array(special_tokens_mask, dtype=bool)
+
+        # determine how many tokens we need to mask in total
+        is_token = ~(labels == self.tokenizer.pad_token_id) & ~special_tokens_mask
+        num_to_mask = int(math.ceil(is_token.astype(float).sum() * self.mask_ratio))
+        if num_to_mask == 0:
+            return inputs, labels
+
+        # generate a sufficient number of span lengths
+        lengths = poisson(lam=self.poisson_lambda, size=(num_to_mask,))
+        while np.cumsum(lengths, 0)[-1] < num_to_mask:
+            lengths = np.concatenate([lengths, poisson(lam=self.poisson_lambda, size=(num_to_mask,))])
+
+        # remove all spans of length 0
+        # Note that BART inserts additional mask tokens where length == 0,
+        # which we do not implement for now as it adds additional complexity
+        lengths = lengths[lengths > 0]
+
+        # trim to about num_to_mask tokens
+        idx = np.argmin(np.abs(np.cumsum(lengths, 0) - num_to_mask)) + 1
+        lengths = lengths[: idx + 1]
+
+        # select span start indices
+        # print("IS TOKEN")
+        # print(is_token)
+        # print(sum(list(map(lambda x: 1 if(x) else 0, is_token[0]))))
+        token_indices = np.argwhere(is_token == 1)
+        # print("TOKEN INDICES")
+        # print(token_indices)
+        span_starts = permutation(token_indices.shape[0])[: lengths.shape[0]]
+
+        # prepare mask
+        masked_indices = np.array(token_indices[span_starts])
+        # print("MASKED INDICES")
+        # print(masked_indices)
+        mask = np.full_like(labels, fill_value=False)
+
+        # mask span start indices
+        for mi in masked_indices:
+            mask[tuple(mi)] = True
+        lengths -= 1
+
+        # fill up spans
+        max_index = labels.shape[1] - 1
+        remaining = (lengths > 0) & (masked_indices[:, 1] < max_index)
+        while np.any(remaining):
+            masked_indices[remaining, 1] += 1
+            for mi in masked_indices:
+                mask[tuple(mi)] = True
+            lengths -= 1
+            remaining = (lengths > 0) & (masked_indices[:, 1] < max_index)
+
+        # place the mask tokens
+        mask[np.where(special_tokens_mask)] = False
+        inputs[np.where(mask)] = self.tokenizer.mask_token_id
+
+        if not do_permutate:
+            labels[np.where(mask)] = -100
+        else:
+            labels[np.where(special_tokens_mask)] = -100
+
+        # remove mask tokens that are not starts of spans
+        to_remove = (mask == 1) & np.roll((mask == 1), 1, 1)
+        new_inputs = np.full_like(labels, fill_value=self.tokenizer.pad_token_id)
+
+        # splits = list(map(lambda x: x.reshape(-1),  np.split(inputs_copy, indices_or_sections=2, axis=0))
+        for i, example in enumerate(np.split(inputs, indices_or_sections=new_inputs.shape[0], axis=0)):
+            new_example = example[0][~to_remove[i]]
+            new_inputs[i, 0 : new_example.shape[0]] = new_example
+
+        # batching now fixed
+        return new_inputs, labels

jax_env/bin/Activate.ps1

@@ -0,0 +1,241 @@
+<#
+.Synopsis
+Activate a Python virtual environment for the current PowerShell session.
+
+.Description
+Pushes the python executable for a virtual environment to the front of the
+$Env:PATH environment variable and sets the prompt to signify that you are
+in a Python virtual environment. Makes use of the command line switches as
+well as the `pyvenv.cfg` file values present in the virtual environment.
+
+.Parameter VenvDir
+Path to the directory that contains the virtual environment to activate. The
+default value for this is the parent of the directory that the Activate.ps1
+script is located within.
+
+.Parameter Prompt
+The prompt prefix to display when this virtual environment is activated. By
+default, this prompt is the name of the virtual environment folder (VenvDir)
+surrounded by parentheses and followed by a single space (ie. '(.venv) ').
+
+.Example
+Activate.ps1
+Activates the Python virtual environment that contains the Activate.ps1 script.
+
+.Example
+Activate.ps1 -Verbose
+Activates the Python virtual environment that contains the Activate.ps1 script,
+and shows extra information about the activation as it executes.
+
+.Example
+Activate.ps1 -VenvDir C:\Users\MyUser\Common\.venv
+Activates the Python virtual environment located in the specified location.
+
+.Example
+Activate.ps1 -Prompt "MyPython"
+Activates the Python virtual environment that contains the Activate.ps1 script,
+and prefixes the current prompt with the specified string (surrounded in
+parentheses) while the virtual environment is active.
+
+.Notes
+On Windows, it may be required to enable this Activate.ps1 script by setting the
+execution policy for the user. You can do this by issuing the following PowerShell
+command:
+
+PS C:\> Set-ExecutionPolicy -ExecutionPolicy RemoteSigned -Scope CurrentUser
+
+For more information on Execution Policies: 
+https://go.microsoft.com/fwlink/?LinkID=135170
+
+#>
+Param(
+    [Parameter(Mandatory = $false)]
+    [String]
+    $VenvDir,
+    [Parameter(Mandatory = $false)]
+    [String]
+    $Prompt
+)
+
+<# Function declarations --------------------------------------------------- #>
+
+<#
+.Synopsis
+Remove all shell session elements added by the Activate script, including the
+addition of the virtual environment's Python executable from the beginning of
+the PATH variable.
+
+.Parameter NonDestructive
+If present, do not remove this function from the global namespace for the
+session.
+
+#>
+function global:deactivate ([switch]$NonDestructive) {
+    # Revert to original values
+
+    # The prior prompt:
+    if (Test-Path -Path Function:_OLD_VIRTUAL_PROMPT) {
+        Copy-Item -Path Function:_OLD_VIRTUAL_PROMPT -Destination Function:prompt
+        Remove-Item -Path Function:_OLD_VIRTUAL_PROMPT
+    }
+
+    # The prior PYTHONHOME:
+    if (Test-Path -Path Env:_OLD_VIRTUAL_PYTHONHOME) {
+        Copy-Item -Path Env:_OLD_VIRTUAL_PYTHONHOME -Destination Env:PYTHONHOME
+        Remove-Item -Path Env:_OLD_VIRTUAL_PYTHONHOME
+    }
+
+    # The prior PATH:
+    if (Test-Path -Path Env:_OLD_VIRTUAL_PATH) {
+        Copy-Item -Path Env:_OLD_VIRTUAL_PATH -Destination Env:PATH
+        Remove-Item -Path Env:_OLD_VIRTUAL_PATH
+    }
+
+    # Just remove the VIRTUAL_ENV altogether:
+    if (Test-Path -Path Env:VIRTUAL_ENV) {
+        Remove-Item -Path env:VIRTUAL_ENV
+    }
+
+    # Just remove the _PYTHON_VENV_PROMPT_PREFIX altogether:
+    if (Get-Variable -Name "_PYTHON_VENV_PROMPT_PREFIX" -ErrorAction SilentlyContinue) {
+        Remove-Variable -Name _PYTHON_VENV_PROMPT_PREFIX -Scope Global -Force
+    }
+
+    # Leave deactivate function in the global namespace if requested:
+    if (-not $NonDestructive) {
+        Remove-Item -Path function:deactivate
+    }
+}
+
+<#
+.Description
+Get-PyVenvConfig parses the values from the pyvenv.cfg file located in the
+given folder, and returns them in a map.
+
+For each line in the pyvenv.cfg file, if that line can be parsed into exactly
+two strings separated by `=` (with any amount of whitespace surrounding the =)
+then it is considered a `key = value` line. The left hand string is the key,
+the right hand is the value.
+
+If the value starts with a `'` or a `"` then the first and last character is
+stripped from the value before being captured.
+
+.Parameter ConfigDir
+Path to the directory that contains the `pyvenv.cfg` file.
+#>
+function Get-PyVenvConfig(
+    [String]
+    $ConfigDir
+) {
+    Write-Verbose "Given ConfigDir=$ConfigDir, obtain values in pyvenv.cfg"
+
+    # Ensure the file exists, and issue a warning if it doesn't (but still allow the function to continue).
+    $pyvenvConfigPath = Join-Path -Resolve -Path $ConfigDir -ChildPath 'pyvenv.cfg' -ErrorAction Continue
+
+    # An empty map will be returned if no config file is found.
+    $pyvenvConfig = @{ }
+
+    if ($pyvenvConfigPath) {
+
+        Write-Verbose "File exists, parse `key = value` lines"
+        $pyvenvConfigContent = Get-Content -Path $pyvenvConfigPath
+
+        $pyvenvConfigContent | ForEach-Object {
+            $keyval = $PSItem -split "\s*=\s*", 2
+            if ($keyval[0] -and $keyval[1]) {
+                $val = $keyval[1]
+
+                # Remove extraneous quotations around a string value.
+                if ("'""".Contains($val.Substring(0, 1))) {
+                    $val = $val.Substring(1, $val.Length - 2)
+                }
+
+                $pyvenvConfig[$keyval[0]] = $val
+                Write-Verbose "Adding Key: '$($keyval[0])'='$val'"
+            }
+        }
+    }
+    return $pyvenvConfig
+}
+
+
+<# Begin Activate script --------------------------------------------------- #>
+
+# Determine the containing directory of this script
+$VenvExecPath = Split-Path -Parent $MyInvocation.MyCommand.Definition
+$VenvExecDir = Get-Item -Path $VenvExecPath
+
+Write-Verbose "Activation script is located in path: '$VenvExecPath'"
+Write-Verbose "VenvExecDir Fullname: '$($VenvExecDir.FullName)"
+Write-Verbose "VenvExecDir Name: '$($VenvExecDir.Name)"
+
+# Set values required in priority: CmdLine, ConfigFile, Default
+# First, get the location of the virtual environment, it might not be
+# VenvExecDir if specified on the command line.
+if ($VenvDir) {
+    Write-Verbose "VenvDir given as parameter, using '$VenvDir' to determine values"
+}
+else {
+    Write-Verbose "VenvDir not given as a parameter, using parent directory name as VenvDir."
+    $VenvDir = $VenvExecDir.Parent.FullName.TrimEnd("\\/")
+    Write-Verbose "VenvDir=$VenvDir"
+}
+
+# Next, read the `pyvenv.cfg` file to determine any required value such
+# as `prompt`.
+$pyvenvCfg = Get-PyVenvConfig -ConfigDir $VenvDir
+
+# Next, set the prompt from the command line, or the config file, or
+# just use the name of the virtual environment folder.
+if ($Prompt) {
+    Write-Verbose "Prompt specified as argument, using '$Prompt'"
+}
+else {
+    Write-Verbose "Prompt not specified as argument to script, checking pyvenv.cfg value"
+    if ($pyvenvCfg -and $pyvenvCfg['prompt']) {
+        Write-Verbose "  Setting based on value in pyvenv.cfg='$($pyvenvCfg['prompt'])'"
+        $Prompt = $pyvenvCfg['prompt'];
+    }
+    else {
+        Write-Verbose "  Setting prompt based on parent's directory's name. (Is the directory name passed to venv module when creating the virutal environment)"
+        Write-Verbose "  Got leaf-name of $VenvDir='$(Split-Path -Path $venvDir -Leaf)'"
+        $Prompt = Split-Path -Path $venvDir -Leaf
+    }
+}
+
+Write-Verbose "Prompt = '$Prompt'"
+Write-Verbose "VenvDir='$VenvDir'"
+
+# Deactivate any currently active virtual environment, but leave the
+# deactivate function in place.
+deactivate -nondestructive
+
+# Now set the environment variable VIRTUAL_ENV, used by many tools to determine
+# that there is an activated venv.
+$env:VIRTUAL_ENV = $VenvDir
+
+if (-not $Env:VIRTUAL_ENV_DISABLE_PROMPT) {
+
+    Write-Verbose "Setting prompt to '$Prompt'"
+
+    # Set the prompt to include the env name
+    # Make sure _OLD_VIRTUAL_PROMPT is global
+    function global:_OLD_VIRTUAL_PROMPT { "" }
+    Copy-Item -Path function:prompt -Destination function:_OLD_VIRTUAL_PROMPT
+    New-Variable -Name _PYTHON_VENV_PROMPT_PREFIX -Description "Python virtual environment prompt prefix" -Scope Global -Option ReadOnly -Visibility Public -Value $Prompt
+
+    function global:prompt {
+        Write-Host -NoNewline -ForegroundColor Green "($_PYTHON_VENV_PROMPT_PREFIX) "
+        _OLD_VIRTUAL_PROMPT
+    }
+}
+
+# Clear PYTHONHOME
+if (Test-Path -Path Env:PYTHONHOME) {
+    Copy-Item -Path Env:PYTHONHOME -Destination Env:_OLD_VIRTUAL_PYTHONHOME
+    Remove-Item -Path Env:PYTHONHOME
+}
+
+# Add the venv to the PATH
+Copy-Item -Path Env:PATH -Destination Env:_OLD_VIRTUAL_PATH
+$Env:PATH = "$VenvExecDir$([System.IO.Path]::PathSeparator)$Env:PATH"

jax_env/bin/activate

@@ -0,0 +1,76 @@
+# This file must be used with "source bin/activate" *from bash*
+# you cannot run it directly
+
+deactivate () {
+    # reset old environment variables
+    if [ -n "${_OLD_VIRTUAL_PATH:-}" ] ; then
+        PATH="${_OLD_VIRTUAL_PATH:-}"
+        export PATH
+        unset _OLD_VIRTUAL_PATH
+    fi
+    if [ -n "${_OLD_VIRTUAL_PYTHONHOME:-}" ] ; then
+        PYTHONHOME="${_OLD_VIRTUAL_PYTHONHOME:-}"
+        export PYTHONHOME
+        unset _OLD_VIRTUAL_PYTHONHOME
+    fi
+
+    # This should detect bash and zsh, which have a hash command that must
+    # be called to get it to forget past commands.  Without forgetting
+    # past commands the $PATH changes we made may not be respected
+    if [ -n "${BASH:-}" -o -n "${ZSH_VERSION:-}" ] ; then
+        hash -r
+    fi
+
+    if [ -n "${_OLD_VIRTUAL_PS1:-}" ] ; then
+        PS1="${_OLD_VIRTUAL_PS1:-}"
+        export PS1
+        unset _OLD_VIRTUAL_PS1
+    fi
+
+    unset VIRTUAL_ENV
+    if [ ! "${1:-}" = "nondestructive" ] ; then
+    # Self destruct!
+        unset -f deactivate
+    fi
+}
+
+# unset irrelevant variables
+deactivate nondestructive
+
+VIRTUAL_ENV="/mnt/disks/flaxdisk/flax-bart-nb-nn/jax_env"
+export VIRTUAL_ENV
+
+_OLD_VIRTUAL_PATH="$PATH"
+PATH="$VIRTUAL_ENV/bin:$PATH"
+export PATH
+
+# unset PYTHONHOME if set
+# this will fail if PYTHONHOME is set to the empty string (which is bad anyway)
+# could use `if (set -u; : $PYTHONHOME) ;` in bash
+if [ -n "${PYTHONHOME:-}" ] ; then
+    _OLD_VIRTUAL_PYTHONHOME="${PYTHONHOME:-}"
+    unset PYTHONHOME
+fi
+
+if [ -z "${VIRTUAL_ENV_DISABLE_PROMPT:-}" ] ; then
+    _OLD_VIRTUAL_PS1="${PS1:-}"
+    if [ "x(jax_env) " != x ] ; then
+	PS1="(jax_env) ${PS1:-}"
+    else
+    if [ "`basename \"$VIRTUAL_ENV\"`" = "__" ] ; then
+        # special case for Aspen magic directories
+        # see https://aspen.io/
+        PS1="[`basename \`dirname \"$VIRTUAL_ENV\"\``] $PS1"
+    else
+        PS1="(`basename \"$VIRTUAL_ENV\"`)$PS1"
+    fi
+    fi
+    export PS1
+fi
+
+# This should detect bash and zsh, which have a hash command that must
+# be called to get it to forget past commands.  Without forgetting
+# past commands the $PATH changes we made may not be respected
+if [ -n "${BASH:-}" -o -n "${ZSH_VERSION:-}" ] ; then
+    hash -r
+fi

jax_env/bin/activate.csh

@@ -0,0 +1,37 @@
+# This file must be used with "source bin/activate.csh" *from csh*.
+# You cannot run it directly.
+# Created by Davide Di Blasi <davidedb@gmail.com>.
+# Ported to Python 3.3 venv by Andrew Svetlov <andrew.svetlov@gmail.com>
+
+alias deactivate 'test $?_OLD_VIRTUAL_PATH != 0 && setenv PATH "$_OLD_VIRTUAL_PATH" && unset _OLD_VIRTUAL_PATH; rehash; test $?_OLD_VIRTUAL_PROMPT != 0 && set prompt="$_OLD_VIRTUAL_PROMPT" && unset _OLD_VIRTUAL_PROMPT; unsetenv VIRTUAL_ENV; test "\!:*" != "nondestructive" && unalias deactivate'
+
+# Unset irrelevant variables.
+deactivate nondestructive
+
+setenv VIRTUAL_ENV "/mnt/disks/flaxdisk/flax-bart-nb-nn/jax_env"
+
+set _OLD_VIRTUAL_PATH="$PATH"
+setenv PATH "$VIRTUAL_ENV/bin:$PATH"
+
+
+set _OLD_VIRTUAL_PROMPT="$prompt"
+
+if (! "$?VIRTUAL_ENV_DISABLE_PROMPT") then
+    if ("jax_env" != "") then
+        set env_name = "jax_env"
+    else
+        if (`basename "VIRTUAL_ENV"` == "__") then
+            # special case for Aspen magic directories
+            # see https://aspen.io/
+            set env_name = `basename \`dirname "$VIRTUAL_ENV"\``
+        else
+            set env_name = `basename "$VIRTUAL_ENV"`
+        endif
+    endif
+    set prompt = "[$env_name] $prompt"
+    unset env_name
+endif
+
+alias pydoc python -m pydoc
+
+rehash

jax_env/bin/activate.fish

@@ -0,0 +1,75 @@
+# This file must be used with ". bin/activate.fish" *from fish* (http://fishshell.org)
+# you cannot run it directly
+
+function deactivate  -d "Exit virtualenv and return to normal shell environment"
+    # reset old environment variables
+    if test -n "$_OLD_VIRTUAL_PATH"
+        set -gx PATH $_OLD_VIRTUAL_PATH
+        set -e _OLD_VIRTUAL_PATH
+    end
+    if test -n "$_OLD_VIRTUAL_PYTHONHOME"
+        set -gx PYTHONHOME $_OLD_VIRTUAL_PYTHONHOME
+        set -e _OLD_VIRTUAL_PYTHONHOME
+    end
+
+    if test -n "$_OLD_FISH_PROMPT_OVERRIDE"
+        functions -e fish_prompt
+        set -e _OLD_FISH_PROMPT_OVERRIDE
+        functions -c _old_fish_prompt fish_prompt
+        functions -e _old_fish_prompt
+    end
+
+    set -e VIRTUAL_ENV
+    if test "$argv[1]" != "nondestructive"
+        # Self destruct!
+        functions -e deactivate
+    end
+end
+
+# unset irrelevant variables
+deactivate nondestructive
+
+set -gx VIRTUAL_ENV "/mnt/disks/flaxdisk/flax-bart-nb-nn/jax_env"
+
+set -gx _OLD_VIRTUAL_PATH $PATH
+set -gx PATH "$VIRTUAL_ENV/bin" $PATH
+
+# unset PYTHONHOME if set
+if set -q PYTHONHOME
+    set -gx _OLD_VIRTUAL_PYTHONHOME $PYTHONHOME
+    set -e PYTHONHOME
+end
+
+if test -z "$VIRTUAL_ENV_DISABLE_PROMPT"
+    # fish uses a function instead of an env var to generate the prompt.
+
+    # save the current fish_prompt function as the function _old_fish_prompt
+    functions -c fish_prompt _old_fish_prompt
+
+    # with the original prompt function renamed, we can override with our own.
+    function fish_prompt
+        # Save the return status of the last command
+        set -l old_status $status
+
+        # Prompt override?
+        if test -n "(jax_env) "
+            printf "%s%s" "(jax_env) " (set_color normal)
+        else
+            # ...Otherwise, prepend env
+            set -l _checkbase (basename "$VIRTUAL_ENV")
+            if test $_checkbase = "__"
+                # special case for Aspen magic directories
+                # see https://aspen.io/
+                printf "%s[%s]%s " (set_color -b blue white) (basename (dirname "$VIRTUAL_ENV")) (set_color normal)
+            else
+                printf "%s(%s)%s" (set_color -b blue white) (basename "$VIRTUAL_ENV") (set_color normal)
+            end
+        end
+
+        # Restore the return status of the previous command.
+        echo "exit $old_status" | .
+        _old_fish_prompt
+    end
+
+    set -gx _OLD_FISH_PROMPT_OVERRIDE "$VIRTUAL_ENV"
+end

jax_env/bin/easy_install

@@ -0,0 +1,8 @@
+#!/mnt/disks/flaxdisk/flax-bart-nb-nn/jax_env/bin/python3
+# -*- coding: utf-8 -*-
+import re
+import sys
+from setuptools.command.easy_install import main
+if __name__ == '__main__':
+    sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0])
+    sys.exit(main())

jax_env/bin/easy_install-3.8

@@ -0,0 +1,8 @@
+#!/mnt/disks/flaxdisk/flax-bart-nb-nn/jax_env/bin/python3
+# -*- coding: utf-8 -*-
+import re
+import sys
+from setuptools.command.easy_install import main
+if __name__ == '__main__':
+    sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0])
+    sys.exit(main())

jax_env/bin/pip

@@ -0,0 +1,8 @@
+#!/mnt/disks/flaxdisk/flax-bart-nb-nn/jax_env/bin/python3
+# -*- coding: utf-8 -*-
+import re
+import sys
+from pip._internal.cli.main import main
+if __name__ == '__main__':
+    sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0])
+    sys.exit(main())

jax_env/bin/pip3

@@ -0,0 +1,8 @@
+#!/mnt/disks/flaxdisk/flax-bart-nb-nn/jax_env/bin/python3
+# -*- coding: utf-8 -*-
+import re
+import sys
+from pip._internal.cli.main import main
+if __name__ == '__main__':
+    sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0])
+    sys.exit(main())

jax_env/bin/pip3.8

@@ -0,0 +1,8 @@
+#!/mnt/disks/flaxdisk/flax-bart-nb-nn/jax_env/bin/python3
+# -*- coding: utf-8 -*-
+import re
+import sys
+from pip._internal.cli.main import main
+if __name__ == '__main__':
+    sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0])
+    sys.exit(main())

jax_env/bin/python

@@ -0,0 +1 @@
+python3

jax_env/bin/python3

@@ -0,0 +1 @@
+/home/perk/flax/bin/python3

jax_env/lib/python3.8/site-packages/easy_install.py

@@ -0,0 +1,5 @@
+"""Run the EasyInstall command"""
+
+if __name__ == '__main__':
+    from setuptools.command.easy_install import main
+    main()

jax_env/lib/python3.8/site-packages/pip-20.0.2.dist-info/INSTALLER

@@ -0,0 +1 @@
+pip

jax_env/lib/python3.8/site-packages/pip-20.0.2.dist-info/LICENSE.txt

@@ -0,0 +1,20 @@
+Copyright (c) 2008-2019 The pip developers (see AUTHORS.txt file)
+
+Permission is hereby granted, free of charge, to any person obtaining
+a copy of this software and associated documentation files (the
+"Software"), to deal in the Software without restriction, including
+without limitation the rights to use, copy, modify, merge, publish,
+distribute, sublicense, and/or sell copies of the Software, and to
+permit persons to whom the Software is furnished to do so, subject to
+the following conditions:
+
+The above copyright notice and this permission notice shall be
+included in all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
+EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
+MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
+NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE
+LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
+OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
+WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

jax_env/lib/python3.8/site-packages/pip-20.0.2.dist-info/METADATA

@@ -0,0 +1,84 @@
+Metadata-Version: 2.1
+Name: pip
+Version: 20.0.2
+Summary: The PyPA recommended tool for installing Python packages.
+Home-page: https://pip.pypa.io/
+Author: The pip developers
+Author-email: pypa-dev@groups.google.com
+License: MIT
+Project-URL: Documentation, https://pip.pypa.io
+Project-URL: Source, https://github.com/pypa/pip
+Keywords: distutils easy_install egg setuptools wheel virtualenv
+Platform: UNKNOWN
+Classifier: Development Status :: 5 - Production/Stable
+Classifier: Intended Audience :: Developers
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Topic :: Software Development :: Build Tools
+Classifier: Programming Language :: Python
+Classifier: Programming Language :: Python :: 2
+Classifier: Programming Language :: Python :: 2.7
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.5
+Classifier: Programming Language :: Python :: 3.6
+Classifier: Programming Language :: Python :: 3.7
+Classifier: Programming Language :: Python :: 3.8
+Classifier: Programming Language :: Python :: Implementation :: CPython
+Classifier: Programming Language :: Python :: Implementation :: PyPy
+Requires-Python: >=2.7,!=3.0.*,!=3.1.*,!=3.2.*,!=3.3.*,!=3.4.*
+
+pip - The Python Package Installer
+==================================
+
+.. image:: https://img.shields.io/pypi/v/pip.svg
+   :target: https://pypi.org/project/pip/
+
+.. image:: https://readthedocs.org/projects/pip/badge/?version=latest
+   :target: https://pip.pypa.io/en/latest
+
+pip is the `package installer`_ for Python. You can use pip to install packages from the `Python Package Index`_ and other indexes.
+
+Please take a look at our documentation for how to install and use pip:
+
+* `Installation`_
+* `Usage`_
+
+Updates are released regularly, with a new version every 3 months. More details can be found in our documentation:
+
+* `Release notes`_
+* `Release process`_
+
+If you find bugs, need help, or want to talk to the developers please use our mailing lists or chat rooms:
+
+* `Issue tracking`_
+* `Discourse channel`_
+* `User IRC`_
+
+If you want to get involved head over to GitHub to get the source code, look at our development documentation and feel free to jump on the developer mailing lists and chat rooms:
+
+* `GitHub page`_
+* `Dev documentation`_
+* `Dev mailing list`_
+* `Dev IRC`_
+
+Code of Conduct
+---------------
+
+Everyone interacting in the pip project's codebases, issue trackers, chat
+rooms, and mailing lists is expected to follow the `PyPA Code of Conduct`_.
+
+.. _package installer: https://packaging.python.org/guides/tool-recommendations/
+.. _Python Package Index: https://pypi.org
+.. _Installation: https://pip.pypa.io/en/stable/installing.html
+.. _Usage: https://pip.pypa.io/en/stable/
+.. _Release notes: https://pip.pypa.io/en/stable/news.html
+.. _Release process: https://pip.pypa.io/en/latest/development/release-process/
+.. _GitHub page: https://github.com/pypa/pip
+.. _Dev documentation: https://pip.pypa.io/en/latest/development
+.. _Issue tracking: https://github.com/pypa/pip/issues
+.. _Discourse channel: https://discuss.python.org/c/packaging
+.. _Dev mailing list: https://groups.google.com/forum/#!forum/pypa-dev
+.. _User IRC: https://webchat.freenode.net/?channels=%23pypa
+.. _Dev IRC: https://webchat.freenode.net/?channels=%23pypa-dev
+.. _PyPA Code of Conduct: https://www.pypa.io/en/latest/code-of-conduct/
+
+

jax_env/lib/python3.8/site-packages/pip-20.0.2.dist-info/RECORD

@@ -0,0 +1,246 @@
+../../../bin/pip,sha256=nMap_ot4S8UAEm2gA4OPEAOScwhi8NyPH6ggEPV0af8,260
+../../../bin/pip3,sha256=nMap_ot4S8UAEm2gA4OPEAOScwhi8NyPH6ggEPV0af8,260
+../../../bin/pip3.8,sha256=nMap_ot4S8UAEm2gA4OPEAOScwhi8NyPH6ggEPV0af8,260
+pip-20.0.2.dist-info/INSTALLER,sha256=zuuue4knoyJ-UwPPXg8fezS7VCrXJQrAP7zeNuwvFQg,4
+pip-20.0.2.dist-info/LICENSE.txt,sha256=W6Ifuwlk-TatfRU2LR7W1JMcyMj5_y1NkRkOEJvnRDE,1090
+pip-20.0.2.dist-info/METADATA,sha256=MSgjT2JTt8usp4Hopp5AGEmc-7sKR2Jd7HTMJqCoRhw,3352
+pip-20.0.2.dist-info/RECORD,,
+pip-20.0.2.dist-info/WHEEL,sha256=kGT74LWyRUZrL4VgLh6_g12IeVl_9u9ZVhadrgXZUEY,110
+pip-20.0.2.dist-info/entry_points.txt,sha256=HtfDOwpUlr9s73jqLQ6wF9V0_0qvUXJwCBz7Vwx0Ue0,125
+pip-20.0.2.dist-info/top_level.txt,sha256=zuuue4knoyJ-UwPPXg8fezS7VCrXJQrAP7zeNuwvFQg,4
+pip/__init__.py,sha256=U1AM82iShMaw90K6Yq0Q2-AZ1EsOcqQLQRB-rxwFtII,455
+pip/__main__.py,sha256=NM95x7KuQr-lwPoTjAC0d_QzLJsJjpmAoxZg0mP8s98,632
+pip/__pycache__/__init__.cpython-38.pyc,,
+pip/__pycache__/__main__.cpython-38.pyc,,
+pip/_internal/__init__.py,sha256=j5fiII6yCeZjpW7_7wAVRMM4DwE-gyARGVU4yAADDeE,517
+pip/_internal/__pycache__/__init__.cpython-38.pyc,,
+pip/_internal/__pycache__/build_env.cpython-38.pyc,,
+pip/_internal/__pycache__/cache.cpython-38.pyc,,
+pip/_internal/__pycache__/configuration.cpython-38.pyc,,
+pip/_internal/__pycache__/exceptions.cpython-38.pyc,,
+pip/_internal/__pycache__/legacy_resolve.cpython-38.pyc,,
+pip/_internal/__pycache__/locations.cpython-38.pyc,,
+pip/_internal/__pycache__/main.cpython-38.pyc,,
+pip/_internal/__pycache__/pep425tags.cpython-38.pyc,,
+pip/_internal/__pycache__/pyproject.cpython-38.pyc,,
+pip/_internal/__pycache__/self_outdated_check.cpython-38.pyc,,
+pip/_internal/__pycache__/wheel_builder.cpython-38.pyc,,
+pip/_internal/build_env.py,sha256=--aNgzIdYrCOclHMwoAdpclCpfdFE_jooRuCy5gczwg,7532
+pip/_internal/cache.py,sha256=16GrnDRLBQNlfKWIuIF6Sa-EFS78kez_w1WEjT3ykTI,11605
+pip/_internal/cli/__init__.py,sha256=FkHBgpxxb-_gd6r1FjnNhfMOzAUYyXoXKJ6abijfcFU,132
+pip/_internal/cli/__pycache__/__init__.cpython-38.pyc,,
+pip/_internal/cli/__pycache__/autocompletion.cpython-38.pyc,,
+pip/_internal/cli/__pycache__/base_command.cpython-38.pyc,,
+pip/_internal/cli/__pycache__/cmdoptions.cpython-38.pyc,,
+pip/_internal/cli/__pycache__/command_context.cpython-38.pyc,,
+pip/_internal/cli/__pycache__/main.cpython-38.pyc,,
+pip/_internal/cli/__pycache__/main_parser.cpython-38.pyc,,
+pip/_internal/cli/__pycache__/parser.cpython-38.pyc,,
+pip/_internal/cli/__pycache__/req_command.cpython-38.pyc,,
+pip/_internal/cli/__pycache__/status_codes.cpython-38.pyc,,
+pip/_internal/cli/autocompletion.py,sha256=ekGNtcDI0p7rFVc-7s4T9Tbss4Jgb7vsB649XJIblRg,6547
+pip/_internal/cli/base_command.py,sha256=v6yl5XNRqye8BT9ep8wvpMu6lylP_Hu6D95r_HqbpbQ,7948
+pip/_internal/cli/cmdoptions.py,sha256=f1TVHuu_fR3lLlMo6b367H_GsWFv26tLI9cAS-kZfE0,28114
+pip/_internal/cli/command_context.py,sha256=ygMVoTy2jpNilKT-6416gFSQpaBtrKRBbVbi2fy__EU,975
+pip/_internal/cli/main.py,sha256=8iq3bHe5lxJTB2EvKOqZ38NS0MmoS79_S1kgj4QuH8A,2610
+pip/_internal/cli/main_parser.py,sha256=W9OWeryh7ZkqELohaFh0Ko9sB98ZkSeDmnYbOZ1imBc,2819
+pip/_internal/cli/parser.py,sha256=O9djTuYQuSfObiY-NU6p4MJCfWsRUnDpE2YGA_fwols,9487
+pip/_internal/cli/req_command.py,sha256=pAUAglpTn0mUA6lRs7KN71yOm1KDabD0ySVTQTqWTSA,12463
+pip/_internal/cli/status_codes.py,sha256=F6uDG6Gj7RNKQJUDnd87QKqI16Us-t-B0wPF_4QMpWc,156
+pip/_internal/commands/__init__.py,sha256=uTSj58QlrSKeXqCUSdL-eAf_APzx5BHy1ABxb0j5ZNE,3714
+pip/_internal/commands/__pycache__/__init__.cpython-38.pyc,,
+pip/_internal/commands/__pycache__/check.cpython-38.pyc,,
+pip/_internal/commands/__pycache__/completion.cpython-38.pyc,,
+pip/_internal/commands/__pycache__/configuration.cpython-38.pyc,,
+pip/_internal/commands/__pycache__/debug.cpython-38.pyc,,
+pip/_internal/commands/__pycache__/download.cpython-38.pyc,,
+pip/_internal/commands/__pycache__/freeze.cpython-38.pyc,,
+pip/_internal/commands/__pycache__/hash.cpython-38.pyc,,
+pip/_internal/commands/__pycache__/help.cpython-38.pyc,,
+pip/_internal/commands/__pycache__/install.cpython-38.pyc,,
+pip/_internal/commands/__pycache__/list.cpython-38.pyc,,
+pip/_internal/commands/__pycache__/search.cpython-38.pyc,,
+pip/_internal/commands/__pycache__/show.cpython-38.pyc,,
+pip/_internal/commands/__pycache__/uninstall.cpython-38.pyc,,
+pip/_internal/commands/__pycache__/wheel.cpython-38.pyc,,
+pip/_internal/commands/check.py,sha256=mgLNYT3bd6Kmynwh4zzcBmVlFZ-urMo40jTgk6U405E,1505
+pip/_internal/commands/completion.py,sha256=UFQvq0Q4_B96z1bvnQyMOq82aPSu05RejbLmqeTZjC0,2975
+pip/_internal/commands/configuration.py,sha256=6riioZjMhsNSEct7dE-X8SobGodk3WERKJvuyjBje4Q,7226
+pip/_internal/commands/debug.py,sha256=a8llax2hRkxgK-tvwdJgaCaZCYPIx0fDvrlMDoYr8bQ,4209
+pip/_internal/commands/download.py,sha256=zX_0-IeFb4C8dxSmGHxk-6H5kehtyTSsdWpjNpAhSww,5007
+pip/_internal/commands/freeze.py,sha256=QS-4ib8jbKJ2wrDaDbTuyaB3Y_iJ5CQC2gAVHuAv9QU,3481
+pip/_internal/commands/hash.py,sha256=47teimfAPhpkaVbSDaafck51BT3XXYuL83lAqc5lOcE,1735
+pip/_internal/commands/help.py,sha256=Nhecq--ydFn80Gm1Zvbf9943EcRJfO0TnXUhsF0RO7s,1181
+pip/_internal/commands/install.py,sha256=T4P3J1rw7CQrZX4OUamtcoWMkTrJBfUe6gWpTfZW1bQ,27286
+pip/_internal/commands/list.py,sha256=2l0JiqHxjxDHNTCb2HZOjwwdo4duS1R0MsqZb6HSMKk,10660
+pip/_internal/commands/search.py,sha256=7Il8nKZ9mM7qF5jlnBoPvSIFY9f-0-5IbYoX3miTuZY,5148
+pip/_internal/commands/show.py,sha256=Vzsj2oX0JBl94MPyF3LV8YoMcigl8B2UsMM8zp0pH2s,6792
+pip/_internal/commands/uninstall.py,sha256=8mldFbrQecSoWDZRqxBgJkrlvx6Y9Iy7cs-2BIgtXt4,2983
+pip/_internal/commands/wheel.py,sha256=TMU5ZhjLo7BIZQApGPsYfoCsbGTnvP-N9jkgPJXhj1Y,7170
+pip/_internal/configuration.py,sha256=MgKrLFBJBkF3t2VJM4tvlnEspfSuS4scp_LhHWh53nY,14222
+pip/_internal/distributions/__init__.py,sha256=ECBUW5Gtu9TjJwyFLvim-i6kUMYVuikNh9I5asL6tbA,959
+pip/_internal/distributions/__pycache__/__init__.cpython-38.pyc,,
+pip/_internal/distributions/__pycache__/base.cpython-38.pyc,,
+pip/_internal/distributions/__pycache__/installed.cpython-38.pyc,,
+pip/_internal/distributions/__pycache__/sdist.cpython-38.pyc,,
+pip/_internal/distributions/__pycache__/wheel.cpython-38.pyc,,
+pip/_internal/distributions/base.py,sha256=ruprpM_L2T2HNi3KLUHlbHimZ1sWVw-3Q0Lb8O7TDAI,1425
+pip/_internal/distributions/installed.py,sha256=YqlkBKr6TVP1MAYS6SG8ojud21wVOYLMZ8jMLJe9MSU,760
+pip/_internal/distributions/sdist.py,sha256=D4XTMlCwgPlK69l62GLYkNSVTVe99fR5iAcVt2EbGok,4086
+pip/_internal/distributions/wheel.py,sha256=95uD-TfaYoq3KiKBdzk9YMN4RRqJ28LNoSTS2K46gek,1294
+pip/_internal/exceptions.py,sha256=6YRuwXAK6F1iyUWKIkCIpWWN2khkAn1sZOgrFA9S8Ro,10247
+pip/_internal/index/__init__.py,sha256=vpt-JeTZefh8a-FC22ZeBSXFVbuBcXSGiILhQZJaNpQ,30
+pip/_internal/index/__pycache__/__init__.cpython-38.pyc,,
+pip/_internal/index/__pycache__/collector.cpython-38.pyc,,
+pip/_internal/index/__pycache__/package_finder.cpython-38.pyc,,
+pip/_internal/index/collector.py,sha256=YS7Ix4oylU7ZbPTPFugh-244GSRqMvdHsGUG6nmz2gE,17892
+pip/_internal/index/package_finder.py,sha256=2Rg75AOpLj8BN1jyL8EI-Iw-Hv6ibJkrYVARCht3bX8,37542
+pip/_internal/legacy_resolve.py,sha256=L7R72I7CjVgJlPTggmA1j4b-H8NmxNu_dKVhrpGXGps,16277
+pip/_internal/locations.py,sha256=VifFEqhc7FWFV8QGoEM3CpECRY8Doq7kTytytxsEgx0,6734
+pip/_internal/main.py,sha256=IVBnUQ-FG7DK6617uEXRB5_QJqspAsBFmTmTesYkbdQ,437
+pip/_internal/models/__init__.py,sha256=3DHUd_qxpPozfzouoqa9g9ts1Czr5qaHfFxbnxriepM,63
+pip/_internal/models/__pycache__/__init__.cpython-38.pyc,,
+pip/_internal/models/__pycache__/candidate.cpython-38.pyc,,
+pip/_internal/models/__pycache__/format_control.cpython-38.pyc,,
+pip/_internal/models/__pycache__/index.cpython-38.pyc,,
+pip/_internal/models/__pycache__/link.cpython-38.pyc,,
+pip/_internal/models/__pycache__/scheme.cpython-38.pyc,,
+pip/_internal/models/__pycache__/search_scope.cpython-38.pyc,,
+pip/_internal/models/__pycache__/selection_prefs.cpython-38.pyc,,
+pip/_internal/models/__pycache__/target_python.cpython-38.pyc,,
+pip/_internal/models/__pycache__/wheel.cpython-38.pyc,,
+pip/_internal/models/candidate.py,sha256=Y58Bcm6oXUj0iS-yhmerlGo5CQJI2p0Ww9h6hR9zQDw,1150
+pip/_internal/models/format_control.py,sha256=ICzVjjGwfZYdX-eLLKHjMHLutEJlAGpfj09OG_eMqac,2673
+pip/_internal/models/index.py,sha256=K59A8-hVhBM20Xkahr4dTwP7OjkJyEqXH11UwHFVgqM,1060
+pip/_internal/models/link.py,sha256=y0H2ZOk0P6d1lfGUL2Pl09xFgZcRt5HwN2LElMifOpI,6827
+pip/_internal/models/scheme.py,sha256=vvhBrrno7eVDXcdKHiZWwxhPHf4VG5uSCEkC0QDR2RU,679
+pip/_internal/models/search_scope.py,sha256=2LXbU4wV8LwqdtXQXNXFYKv-IxiDI_QwSz9ZgbwtAfk,3898
+pip/_internal/models/selection_prefs.py,sha256=rPeif2KKjhTPXeMoQYffjqh10oWpXhdkxRDaPT1HO8k,1908
+pip/_internal/models/target_python.py,sha256=c-cFi6zCuo5HYbXNS3rVVpKRaHVh5yQlYEjEW23SidQ,3799
+pip/_internal/models/wheel.py,sha256=6KLuLKH5b0C5goWQXGSISRaq2UZtkHUEAU1y1Zsrwms,2766
+pip/_internal/network/__init__.py,sha256=jf6Tt5nV_7zkARBrKojIXItgejvoegVJVKUbhAa5Ioc,50
+pip/_internal/network/__pycache__/__init__.cpython-38.pyc,,
+pip/_internal/network/__pycache__/auth.cpython-38.pyc,,
+pip/_internal/network/__pycache__/cache.cpython-38.pyc,,
+pip/_internal/network/__pycache__/download.cpython-38.pyc,,
+pip/_internal/network/__pycache__/session.cpython-38.pyc,,
+pip/_internal/network/__pycache__/utils.cpython-38.pyc,,
+pip/_internal/network/__pycache__/xmlrpc.cpython-38.pyc,,
+pip/_internal/network/auth.py,sha256=K3G1ukKb3PiH8w_UnpXTz8qQsTULO-qdbfOE9zTo1fE,11119
+pip/_internal/network/cache.py,sha256=51CExcRkXWrgMZ7WsrZ6cmijKfViD5tVgKbBvJHO1IE,2394
+pip/_internal/network/download.py,sha256=3D9vdJmVwmCUMxzC-TaVI_GvVOpQna3BLEYNPCSx3Fc,6260
+pip/_internal/network/session.py,sha256=u1IXQfv21R1xv86ulyiB58-be4sYm90eFB0Wp8fVMYw,14702
+pip/_internal/network/utils.py,sha256=iiixo1OeaQ3niUWiBjg59PN6f1w7vvTww1vFriTD_IU,1959
+pip/_internal/network/xmlrpc.py,sha256=AL115M3vFJ8xiHVJneb8Hi0ZFeRvdPhblC89w25OG5s,1597
+pip/_internal/operations/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+pip/_internal/operations/__pycache__/__init__.cpython-38.pyc,,
+pip/_internal/operations/__pycache__/check.cpython-38.pyc,,
+pip/_internal/operations/__pycache__/freeze.cpython-38.pyc,,
+pip/_internal/operations/__pycache__/prepare.cpython-38.pyc,,
+pip/_internal/operations/build/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+pip/_internal/operations/build/__pycache__/__init__.cpython-38.pyc,,
+pip/_internal/operations/build/__pycache__/metadata.cpython-38.pyc,,
+pip/_internal/operations/build/__pycache__/metadata_legacy.cpython-38.pyc,,
+pip/_internal/operations/build/__pycache__/wheel.cpython-38.pyc,,
+pip/_internal/operations/build/__pycache__/wheel_legacy.cpython-38.pyc,,
+pip/_internal/operations/build/metadata.py,sha256=yHMi5gHYXcXyHcvUPWHdO-UyOo3McFWljn_nHfM1O9c,1307
+pip/_internal/operations/build/metadata_legacy.py,sha256=4n6N7BTysqVmEpITzT2UVClyt0Peij_Im8Qm965IWB4,3957
+pip/_internal/operations/build/wheel.py,sha256=ntltdNP6D2Tpr4V0agssu6rE0F9LaBpJkYT6zSdhEbw,1469
+pip/_internal/operations/build/wheel_legacy.py,sha256=DYSxQKutwSZnmNvWkwsl2HzE2XQBxV0i0wTphjtUe90,3349
+pip/_internal/operations/check.py,sha256=a6uHG0daoWpmSPCdL7iYJaGQYZ-CRvPvTnCv2PnIIs0,5353
+pip/_internal/operations/freeze.py,sha256=td4BeRnW10EXFTZrx6VgygO3CrjqD5B9f0BGzjQm-Ew,10180
+pip/_internal/operations/install/__init__.py,sha256=mX7hyD2GNBO2mFGokDQ30r_GXv7Y_PLdtxcUv144e-s,51
+pip/_internal/operations/install/__pycache__/__init__.cpython-38.pyc,,
+pip/_internal/operations/install/__pycache__/editable_legacy.cpython-38.pyc,,
+pip/_internal/operations/install/__pycache__/legacy.cpython-38.pyc,,
+pip/_internal/operations/install/__pycache__/wheel.cpython-38.pyc,,
+pip/_internal/operations/install/editable_legacy.py,sha256=rJ_xs2qtDUjpY2-n6eYlVyZiNoKbOtZXZrYrcnIELt4,1488
+pip/_internal/operations/install/legacy.py,sha256=eBV8gHbO9sBlBc-4nuR3Sd2nikHgEcnC9khfeLiypio,4566
+pip/_internal/operations/install/wheel.py,sha256=xdCjH6uIUyg39Pf8tUaMFUN4a7eozJAFMb_wKcgQlsY,23012
+pip/_internal/operations/prepare.py,sha256=ro2teBlbBpkRJhBKraP9CoJgVLpueSk62ziWhRToXww,20942
+pip/_internal/pep425tags.py,sha256=SlIQokevkoKnXhoK3PZvXiDoj8hFKoJ7thDifDtga3k,5490
+pip/_internal/pyproject.py,sha256=VJKsrXORGiGoDPVKCQhuu4tWlQSTOhoiRlVLRNu4rx4,7400
+pip/_internal/req/__init__.py,sha256=UVaYPlHZVGRBQQPjvGC_6jJDQtewXm0ws-8Lxhg_TiY,2671
+pip/_internal/req/__pycache__/__init__.cpython-38.pyc,,
+pip/_internal/req/__pycache__/constructors.cpython-38.pyc,,
+pip/_internal/req/__pycache__/req_file.cpython-38.pyc,,
+pip/_internal/req/__pycache__/req_install.cpython-38.pyc,,
+pip/_internal/req/__pycache__/req_set.cpython-38.pyc,,
+pip/_internal/req/__pycache__/req_tracker.cpython-38.pyc,,
+pip/_internal/req/__pycache__/req_uninstall.cpython-38.pyc,,
+pip/_internal/req/constructors.py,sha256=w5-kWWVCqlSqcIBitw86yq7XGMPpKrHDfQZSE2mJ_xc,14388
+pip/_internal/req/req_file.py,sha256=ECqRUicCw5Y08R1YynZAAp8dSKQhDXoc1Q-mY3a9b6I,18485
+pip/_internal/req/req_install.py,sha256=wjsIr4lDpbVSLqANKJI9mXwRVHaRxcnj8q30UiHoLRA,30442
+pip/_internal/req/req_set.py,sha256=GsrKmupRKhNMhjkofVfCEHEHfgEvYBxClaQH5xLBQHg,8066
+pip/_internal/req/req_tracker.py,sha256=27fvVG8Y2MJS1KpU2rBMnQyUEMHG4lkHT_bzbzQK-c0,4723
+pip/_internal/req/req_uninstall.py,sha256=DWnOsuyYGju6-sylyoCm7GtUNevn9qMAVhjAGLcdXUE,23609
+pip/_internal/self_outdated_check.py,sha256=3KO1pTJUuYaiV9X0t87I9PimkGL82HbhLWbocqKZpBU,8009
+pip/_internal/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+pip/_internal/utils/__pycache__/__init__.cpython-38.pyc,,
+pip/_internal/utils/__pycache__/appdirs.cpython-38.pyc,,
+pip/_internal/utils/__pycache__/compat.cpython-38.pyc,,
+pip/_internal/utils/__pycache__/deprecation.cpython-38.pyc,,
+pip/_internal/utils/__pycache__/distutils_args.cpython-38.pyc,,
+pip/_internal/utils/__pycache__/encoding.cpython-38.pyc,,
+pip/_internal/utils/__pycache__/entrypoints.cpython-38.pyc,,
+pip/_internal/utils/__pycache__/filesystem.cpython-38.pyc,,
+pip/_internal/utils/__pycache__/filetypes.cpython-38.pyc,,
+pip/_internal/utils/__pycache__/glibc.cpython-38.pyc,,
+pip/_internal/utils/__pycache__/hashes.cpython-38.pyc,,
+pip/_internal/utils/__pycache__/inject_securetransport.cpython-38.pyc,,
+pip/_internal/utils/__pycache__/logging.cpython-38.pyc,,
+pip/_internal/utils/__pycache__/marker_files.cpython-38.pyc,,
+pip/_internal/utils/__pycache__/misc.cpython-38.pyc,,
+pip/_internal/utils/__pycache__/models.cpython-38.pyc,,
+pip/_internal/utils/__pycache__/packaging.cpython-38.pyc,,
+pip/_internal/utils/__pycache__/pkg_resources.cpython-38.pyc,,
+pip/_internal/utils/__pycache__/setuptools_build.cpython-38.pyc,,
+pip/_internal/utils/__pycache__/subprocess.cpython-38.pyc,,
+pip/_internal/utils/__pycache__/temp_dir.cpython-38.pyc,,
+pip/_internal/utils/__pycache__/typing.cpython-38.pyc,,
+pip/_internal/utils/__pycache__/ui.cpython-38.pyc,,
+pip/_internal/utils/__pycache__/unpacking.cpython-38.pyc,,
+pip/_internal/utils/__pycache__/urls.cpython-38.pyc,,
+pip/_internal/utils/__pycache__/virtualenv.cpython-38.pyc,,
+pip/_internal/utils/__pycache__/wheel.cpython-38.pyc,,
+pip/_internal/utils/appdirs.py,sha256=PVo_7-IQWHa9qNuNbWSFiF2QGqeLbSAR4eLcYYhQ9ek,1307
+pip/_internal/utils/compat.py,sha256=D7FKGLBdQwWH-dHIGaoWMawDZWBYApvtJVL1kFPJ930,8869
+pip/_internal/utils/deprecation.py,sha256=pBnNogoA4UGTxa_JDnPXBRRYpKMbExAhXpBwAwklOBs,3318
+pip/_internal/utils/distutils_args.py,sha256=a56mblNxk9BGifbpEETG61mmBrqhjtjRkJ4HYn-oOEE,1350
+pip/_internal/utils/encoding.py,sha256=hxZz0t3Whw3d4MHQEiofxalTlfKwxFdLc8fpeGfhKo8,1320
+pip/_internal/utils/entrypoints.py,sha256=vHcNpnksCv6mllihU6hfifdsKPEjwcaJ1aLIXEaynaU,1152
+pip/_internal/utils/filesystem.py,sha256=PXa3vMcz4mbEKtkD0joFI8pBwddLQxhfPFOkVH5xjfE,5255
+pip/_internal/utils/filetypes.py,sha256=R2FwzoeX7b-rZALOXx5cuO8VPPMhUQ4ne7wm3n3IcWA,571
+pip/_internal/utils/glibc.py,sha256=LOeNGgawCKS-4ke9fii78fwXD73dtNav3uxz1Bf-Ab8,3297
+pip/_internal/utils/hashes.py,sha256=my-wSnAWEDvl_8rQaOQcVIWjwh1-f_QiEvGy9TPf53U,3942
+pip/_internal/utils/inject_securetransport.py,sha256=M17ZlFVY66ApgeASVjKKLKNz0LAfk-SyU0HZ4ZB6MmI,810
+pip/_internal/utils/logging.py,sha256=aJL7NldPhS5KGFof6Qt3o3MG5cjm5TOoo7bGRu9_wsg,13033
+pip/_internal/utils/marker_files.py,sha256=CO5djQlrPIozJpJybViH_insoAaBGY1aqEt6-cC-iW0,741
+pip/_internal/utils/misc.py,sha256=uIb58Hiu_g2HRORo2aMcgnW_7R5d-5wUAuoW0fA2ZME,26085
+pip/_internal/utils/models.py,sha256=IA0hw_T4awQzui0kqfIEASm5yLtgZAB08ag59Nip5G8,1148
+pip/_internal/utils/packaging.py,sha256=VtiwcAAL7LBi7tGL2je7LeW4bE11KMHGCsJ1NZY5XtM,3035
+pip/_internal/utils/pkg_resources.py,sha256=ZX-k7V5q_aNWyDse92nN7orN1aCpRLsaxzpkBZ1XKzU,1254
+pip/_internal/utils/setuptools_build.py,sha256=DouaVolV9olDDFIIN9IszaL-FHdNaZt10ufOZFH9ZAU,5070
+pip/_internal/utils/subprocess.py,sha256=Ph3x5eHQBxFotyGhpZN8asSMBud-BBkmgaNfARG-di8,9922
+pip/_internal/utils/temp_dir.py,sha256=87Ib8aNic_hoSDEmUYJHTQIn5-prL2AYL5u_yZ3s4sI,7768
+pip/_internal/utils/typing.py,sha256=xkYwOeHlf4zsHXBDC4310HtEqwhQcYXFPq2h35Tcrl0,1401
+pip/_internal/utils/ui.py,sha256=0FNxXlGtbpPtTviv2oXS9t8bQG_NBdfUgP4GbubhS9U,13911
+pip/_internal/utils/unpacking.py,sha256=M944JTSiapBOSKLWu7lbawpVHSE7flfzZTEr3TAG7v8,9438
+pip/_internal/utils/urls.py,sha256=aNV9wq5ClUmrz6sG-al7hEWJ4ToitOy7l82CmFGFNW8,1481
+pip/_internal/utils/virtualenv.py,sha256=Q3S1WPlI7JWpGOT2jUVJ8l2chm_k7VPJ9cHA_cUluEU,3396
+pip/_internal/utils/wheel.py,sha256=grTRwZtMQwApwbbSPmRVLtac6FKy6SVKeCXNkWyyePA,7302
+pip/_internal/vcs/__init__.py,sha256=viJxJRqRE_mVScum85bgQIXAd6o0ozFt18VpC-qIJrM,617
+pip/_internal/vcs/__pycache__/__init__.cpython-38.pyc,,
+pip/_internal/vcs/__pycache__/bazaar.cpython-38.pyc,,
+pip/_internal/vcs/__pycache__/git.cpython-38.pyc,,
+pip/_internal/vcs/__pycache__/mercurial.cpython-38.pyc,,
+pip/_internal/vcs/__pycache__/subversion.cpython-38.pyc,,
+pip/_internal/vcs/__pycache__/versioncontrol.cpython-38.pyc,,
+pip/_internal/vcs/bazaar.py,sha256=84q1-kj1_nJ9AMzMu8RmMp-riRZu81M7K9kowcYgi3U,3957
+pip/_internal/vcs/git.py,sha256=CdLz3DTsZsLMLPZpEuUwiS40npvDaVB1CNRzoXgcuJQ,14352
+pip/_internal/vcs/mercurial.py,sha256=2mg7BdYI_Fe00fF6omaNccFQLPHBsDBG5CAEzvqn5sA,5110
+pip/_internal/vcs/subversion.py,sha256=Fpwy71AmuqXnoKi6h1SrXRtPjEMn8fieuM1O4j01IBg,12292
+pip/_internal/vcs/versioncontrol.py,sha256=nqoaM1_rzx24WnHtihXA8RcPpnUae0sV2sR_LS_5HFA,22600
+pip/_internal/wheel_builder.py,sha256=gr9jE14W5ZuYblpldo-tpRuyG0e0AVmHLttImuAvXlE,9441
+pip/_vendor/__init__.py,sha256=-0H8nYPGWX7724kAmvYHNUE268FaYlWrNW518BygiUo,4979
+pip/_vendor/__pycache__/__init__.cpython-38.pyc,,

jax_env/lib/python3.8/site-packages/pip-20.0.2.dist-info/WHEEL

@@ -0,0 +1,6 @@
+Wheel-Version: 1.0
+Generator: bdist_wheel (0.34.2)
+Root-Is-Purelib: true
+Tag: py2-none-any
+Tag: py3-none-any
+

jax_env/lib/python3.8/site-packages/pip-20.0.2.dist-info/entry_points.txt

@@ -0,0 +1,5 @@
+[console_scripts]
+pip = pip._internal.cli.main:main
+pip3 = pip._internal.cli.main:main
+pip3.8 = pip._internal.cli.main:main
+

jax_env/lib/python3.8/site-packages/pip-20.0.2.dist-info/top_level.txt

@@ -0,0 +1 @@
+pip

jax_env/lib/python3.8/site-packages/pip/__init__.py

@@ -0,0 +1,18 @@
+from pip._internal.utils.typing import MYPY_CHECK_RUNNING
+
+if MYPY_CHECK_RUNNING:
+    from typing import List, Optional
+
+
+__version__ = "20.0.2"
+
+
+def main(args=None):
+    # type: (Optional[List[str]]) -> int
+    """This is an internal API only meant for use by pip's own console scripts.
+
+    For additional details, see https://github.com/pypa/pip/issues/7498.
+    """
+    from pip._internal.utils.entrypoints import _wrapper
+
+    return _wrapper(args)

jax_env/lib/python3.8/site-packages/pip/__main__.py

@@ -0,0 +1,19 @@
+from __future__ import absolute_import
+
+import os
+import sys
+
+# If we are running from a wheel, add the wheel to sys.path
+# This allows the usage python pip-*.whl/pip install pip-*.whl
+if __package__ == '':
+    # __file__ is pip-*.whl/pip/__main__.py
+    # first dirname call strips of '/__main__.py', second strips off '/pip'
+    # Resulting path is the name of the wheel itself
+    # Add that to sys.path so we can import pip
+    path = os.path.dirname(os.path.dirname(__file__))
+    sys.path.insert(0, path)
+
+from pip._internal.cli.main import main as _main  # isort:skip # noqa
+
+if __name__ == '__main__':
+    sys.exit(_main())

jax_env/lib/python3.8/site-packages/pip/_internal/__init__.py

@@ -0,0 +1,18 @@
+#!/usr/bin/env python
+import pip._internal.utils.inject_securetransport  # noqa
+from pip._internal.utils.typing import MYPY_CHECK_RUNNING
+
+if MYPY_CHECK_RUNNING:
+    from typing import Optional, List
+
+
+def main(args=None):
+    # type: (Optional[List[str]]) -> int
+    """This is preserved for old console scripts that may still be referencing
+    it.
+
+    For additional details, see https://github.com/pypa/pip/issues/7498.
+    """
+    from pip._internal.utils.entrypoints import _wrapper
+
+    return _wrapper(args)

jax_env/lib/python3.8/site-packages/pip/_internal/build_env.py

@@ -0,0 +1,221 @@
+"""Build Environment used for isolation during sdist building
+"""
+
+# The following comment should be removed at some point in the future.
+# mypy: strict-optional=False
+# mypy: disallow-untyped-defs=False
+
+import logging
+import os
+import sys
+import textwrap
+from collections import OrderedDict
+from distutils.sysconfig import get_python_lib
+from sysconfig import get_paths
+
+from pip._vendor.pkg_resources import Requirement, VersionConflict, WorkingSet
+
+from pip import __file__ as pip_location
+from pip._internal.utils.subprocess import call_subprocess
+from pip._internal.utils.temp_dir import TempDirectory
+from pip._internal.utils.typing import MYPY_CHECK_RUNNING
+from pip._internal.utils.ui import open_spinner
+
+if MYPY_CHECK_RUNNING:
+    from typing import Tuple, Set, Iterable, Optional, List
+    from pip._internal.index.package_finder import PackageFinder
+
+logger = logging.getLogger(__name__)
+
+
+class _Prefix:
+
+    def __init__(self, path):
+        # type: (str) -> None
+        self.path = path
+        self.setup = False
+        self.bin_dir = get_paths(
+            'nt' if os.name == 'nt' else 'posix_prefix',
+            vars={'base': path, 'platbase': path}
+        )['scripts']
+        # Note: prefer distutils' sysconfig to get the
+        # library paths so PyPy is correctly supported.
+        purelib = get_python_lib(plat_specific=False, prefix=path)
+        platlib = get_python_lib(plat_specific=True, prefix=path)
+        if purelib == platlib:
+            self.lib_dirs = [purelib]
+        else:
+            self.lib_dirs = [purelib, platlib]
+
+
+class BuildEnvironment(object):
+    """Creates and manages an isolated environment to install build deps
+    """
+
+    def __init__(self):
+        # type: () -> None
+        self._temp_dir = TempDirectory(kind="build-env")
+
+        self._prefixes = OrderedDict((
+            (name, _Prefix(os.path.join(self._temp_dir.path, name)))
+            for name in ('normal', 'overlay')
+        ))
+
+        self._bin_dirs = []  # type: List[str]
+        self._lib_dirs = []  # type: List[str]
+        for prefix in reversed(list(self._prefixes.values())):
+            self._bin_dirs.append(prefix.bin_dir)
+            self._lib_dirs.extend(prefix.lib_dirs)
+
+        # Customize site to:
+        # - ensure .pth files are honored
+        # - prevent access to system site packages
+        system_sites = {
+            os.path.normcase(site) for site in (
+                get_python_lib(plat_specific=False),
+                get_python_lib(plat_specific=True),
+            )
+        }
+        self._site_dir = os.path.join(self._temp_dir.path, 'site')
+        if not os.path.exists(self._site_dir):
+            os.mkdir(self._site_dir)
+        with open(os.path.join(self._site_dir, 'sitecustomize.py'), 'w') as fp:
+            fp.write(textwrap.dedent(
+                '''
+                import os, site, sys
+
+                # First, drop system-sites related paths.
+                original_sys_path = sys.path[:]
+                known_paths = set()
+                for path in {system_sites!r}:
+                    site.addsitedir(path, known_paths=known_paths)
+                system_paths = set(
+                    os.path.normcase(path)
+                    for path in sys.path[len(original_sys_path):]
+                )
+                original_sys_path = [
+                    path for path in original_sys_path
+                    if os.path.normcase(path) not in system_paths
+                ]
+                sys.path = original_sys_path
+
+                # Second, add lib directories.
+                # ensuring .pth file are processed.
+                for path in {lib_dirs!r}:
+                    assert not path in sys.path
+                    site.addsitedir(path)
+                '''
+            ).format(system_sites=system_sites, lib_dirs=self._lib_dirs))
+
+    def __enter__(self):
+        self._save_env = {
+            name: os.environ.get(name, None)
+            for name in ('PATH', 'PYTHONNOUSERSITE', 'PYTHONPATH')
+        }
+
+        path = self._bin_dirs[:]
+        old_path = self._save_env['PATH']
+        if old_path:
+            path.extend(old_path.split(os.pathsep))
+
+        pythonpath = [self._site_dir]
+
+        os.environ.update({
+            'PATH': os.pathsep.join(path),
+            'PYTHONNOUSERSITE': '1',
+            'PYTHONPATH': os.pathsep.join(pythonpath),
+        })
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        for varname, old_value in self._save_env.items():
+            if old_value is None:
+                os.environ.pop(varname, None)
+            else:
+                os.environ[varname] = old_value
+
+    def cleanup(self):
+        # type: () -> None
+        self._temp_dir.cleanup()
+
+    def check_requirements(self, reqs):
+        # type: (Iterable[str]) -> Tuple[Set[Tuple[str, str]], Set[str]]
+        """Return 2 sets:
+            - conflicting requirements: set of (installed, wanted) reqs tuples
+            - missing requirements: set of reqs
+        """
+        missing = set()
+        conflicting = set()
+        if reqs:
+            ws = WorkingSet(self._lib_dirs)
+            for req in reqs:
+                try:
+                    if ws.find(Requirement.parse(req)) is None:
+                        missing.add(req)
+                except VersionConflict as e:
+                    conflicting.add((str(e.args[0].as_requirement()),
+                                     str(e.args[1])))
+        return conflicting, missing
+
+    def install_requirements(
+        self,
+        finder,  # type: PackageFinder
+        requirements,  # type: Iterable[str]
+        prefix_as_string,  # type: str
+        message  # type: Optional[str]
+    ):
+        # type: (...) -> None
+        prefix = self._prefixes[prefix_as_string]
+        assert not prefix.setup
+        prefix.setup = True
+        if not requirements:
+            return
+        args = [
+            sys.executable, os.path.dirname(pip_location), 'install',
+            '--ignore-installed', '--no-user', '--prefix', prefix.path,
+            '--no-warn-script-location',
+        ]  # type: List[str]
+        if logger.getEffectiveLevel() <= logging.DEBUG:
+            args.append('-v')
+        for format_control in ('no_binary', 'only_binary'):
+            formats = getattr(finder.format_control, format_control)
+            args.extend(('--' + format_control.replace('_', '-'),
+                         ','.join(sorted(formats or {':none:'}))))
+
+        index_urls = finder.index_urls
+        if index_urls:
+            args.extend(['-i', index_urls[0]])
+            for extra_index in index_urls[1:]:
+                args.extend(['--extra-index-url', extra_index])
+        else:
+            args.append('--no-index')
+        for link in finder.find_links:
+            args.extend(['--find-links', link])
+
+        for host in finder.trusted_hosts:
+            args.extend(['--trusted-host', host])
+        if finder.allow_all_prereleases:
+            args.append('--pre')
+        args.append('--')
+        args.extend(requirements)
+        with open_spinner(message) as spinner:
+            call_subprocess(args, spinner=spinner)
+
+
+class NoOpBuildEnvironment(BuildEnvironment):
+    """A no-op drop-in replacement for BuildEnvironment
+    """
+
+    def __init__(self):
+        pass
+
+    def __enter__(self):
+        pass
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        pass
+
+    def cleanup(self):
+        pass
+
+    def install_requirements(self, finder, requirements, prefix, message):
+        raise NotImplementedError()

jax_env/lib/python3.8/site-packages/pip/_internal/cache.py

@@ -0,0 +1,329 @@
+"""Cache Management
+"""
+
+# The following comment should be removed at some point in the future.
+# mypy: strict-optional=False
+
+import hashlib
+import json
+import logging
+import os
+
+from pip._vendor.packaging.tags import interpreter_name, interpreter_version
+from pip._vendor.packaging.utils import canonicalize_name
+
+from pip._internal.exceptions import InvalidWheelFilename
+from pip._internal.models.link import Link
+from pip._internal.models.wheel import Wheel
+from pip._internal.utils.temp_dir import TempDirectory
+from pip._internal.utils.typing import MYPY_CHECK_RUNNING
+from pip._internal.utils.urls import path_to_url
+
+if MYPY_CHECK_RUNNING:
+    from typing import Optional, Set, List, Any, Dict
+
+    from pip._vendor.packaging.tags import Tag
+
+    from pip._internal.models.format_control import FormatControl
+
+logger = logging.getLogger(__name__)
+
+
+def _hash_dict(d):
+    # type: (Dict[str, str]) -> str
+    """Return a stable sha224 of a dictionary."""
+    s = json.dumps(d, sort_keys=True, separators=(",", ":"), ensure_ascii=True)
+    return hashlib.sha224(s.encode("ascii")).hexdigest()
+
+
+class Cache(object):
+    """An abstract class - provides cache directories for data from links
+
+
+        :param cache_dir: The root of the cache.
+        :param format_control: An object of FormatControl class to limit
+            binaries being read from the cache.
+        :param allowed_formats: which formats of files the cache should store.
+            ('binary' and 'source' are the only allowed values)
+    """
+
+    def __init__(self, cache_dir, format_control, allowed_formats):
+        # type: (str, FormatControl, Set[str]) -> None
+        super(Cache, self).__init__()
+        assert not cache_dir or os.path.isabs(cache_dir)
+        self.cache_dir = cache_dir or None
+        self.format_control = format_control
+        self.allowed_formats = allowed_formats
+
+        _valid_formats = {"source", "binary"}
+        assert self.allowed_formats.union(_valid_formats) == _valid_formats
+
+    def _get_cache_path_parts_legacy(self, link):
+        # type: (Link) -> List[str]
+        """Get parts of part that must be os.path.joined with cache_dir
+
+        Legacy cache key (pip < 20) for compatibility with older caches.
+        """
+
+        # We want to generate an url to use as our cache key, we don't want to
+        # just re-use the URL because it might have other items in the fragment
+        # and we don't care about those.
+        key_parts = [link.url_without_fragment]
+        if link.hash_name is not None and link.hash is not None:
+            key_parts.append("=".join([link.hash_name, link.hash]))
+        key_url = "#".join(key_parts)
+
+        # Encode our key url with sha224, we'll use this because it has similar
+        # security properties to sha256, but with a shorter total output (and
+        # thus less secure). However the differences don't make a lot of
+        # difference for our use case here.
+        hashed = hashlib.sha224(key_url.encode()).hexdigest()
+
+        # We want to nest the directories some to prevent having a ton of top
+        # level directories where we might run out of sub directories on some
+        # FS.
+        parts = [hashed[:2], hashed[2:4], hashed[4:6], hashed[6:]]
+
+        return parts
+
+    def _get_cache_path_parts(self, link):
+        # type: (Link) -> List[str]
+        """Get parts of part that must be os.path.joined with cache_dir
+        """
+
+        # We want to generate an url to use as our cache key, we don't want to
+        # just re-use the URL because it might have other items in the fragment
+        # and we don't care about those.
+        key_parts = {"url": link.url_without_fragment}
+        if link.hash_name is not None and link.hash is not None:
+            key_parts[link.hash_name] = link.hash
+        if link.subdirectory_fragment:
+            key_parts["subdirectory"] = link.subdirectory_fragment
+
+        # Include interpreter name, major and minor version in cache key
+        # to cope with ill-behaved sdists that build a different wheel
+        # depending on the python version their setup.py is being run on,
+        # and don't encode the difference in compatibility tags.
+        # https://github.com/pypa/pip/issues/7296
+        key_parts["interpreter_name"] = interpreter_name()
+        key_parts["interpreter_version"] = interpreter_version()
+
+        # Encode our key url with sha224, we'll use this because it has similar
+        # security properties to sha256, but with a shorter total output (and
+        # thus less secure). However the differences don't make a lot of
+        # difference for our use case here.
+        hashed = _hash_dict(key_parts)
+
+        # We want to nest the directories some to prevent having a ton of top
+        # level directories where we might run out of sub directories on some
+        # FS.
+        parts = [hashed[:2], hashed[2:4], hashed[4:6], hashed[6:]]
+
+        return parts
+
+    def _get_candidates(self, link, canonical_package_name):
+        # type: (Link, Optional[str]) -> List[Any]
+        can_not_cache = (
+            not self.cache_dir or
+            not canonical_package_name or
+            not link
+        )
+        if can_not_cache:
+            return []
+
+        formats = self.format_control.get_allowed_formats(
+            canonical_package_name
+        )
+        if not self.allowed_formats.intersection(formats):
+            return []
+
+        candidates = []
+        path = self.get_path_for_link(link)
+        if os.path.isdir(path):
+            for candidate in os.listdir(path):
+                candidates.append((candidate, path))
+        # TODO remove legacy path lookup in pip>=21
+        legacy_path = self.get_path_for_link_legacy(link)
+        if os.path.isdir(legacy_path):
+            for candidate in os.listdir(legacy_path):
+                candidates.append((candidate, legacy_path))
+        return candidates
+
+    def get_path_for_link_legacy(self, link):
+        # type: (Link) -> str
+        raise NotImplementedError()
+
+    def get_path_for_link(self, link):
+        # type: (Link) -> str
+        """Return a directory to store cached items in for link.
+        """
+        raise NotImplementedError()
+
+    def get(
+        self,
+        link,            # type: Link
+        package_name,    # type: Optional[str]
+        supported_tags,  # type: List[Tag]
+    ):
+        # type: (...) -> Link
+        """Returns a link to a cached item if it exists, otherwise returns the
+        passed link.
+        """
+        raise NotImplementedError()
+
+    def cleanup(self):
+        # type: () -> None
+        pass
+
+
+class SimpleWheelCache(Cache):
+    """A cache of wheels for future installs.
+    """
+
+    def __init__(self, cache_dir, format_control):
+        # type: (str, FormatControl) -> None
+        super(SimpleWheelCache, self).__init__(
+            cache_dir, format_control, {"binary"}
+        )
+
+    def get_path_for_link_legacy(self, link):
+        # type: (Link) -> str
+        parts = self._get_cache_path_parts_legacy(link)
+        return os.path.join(self.cache_dir, "wheels", *parts)
+
+    def get_path_for_link(self, link):
+        # type: (Link) -> str
+        """Return a directory to store cached wheels for link
+
+        Because there are M wheels for any one sdist, we provide a directory
+        to cache them in, and then consult that directory when looking up
+        cache hits.
+
+        We only insert things into the cache if they have plausible version
+        numbers, so that we don't contaminate the cache with things that were
+        not unique. E.g. ./package might have dozens of installs done for it
+        and build a version of 0.0...and if we built and cached a wheel, we'd
+        end up using the same wheel even if the source has been edited.
+
+        :param link: The link of the sdist for which this will cache wheels.
+        """
+        parts = self._get_cache_path_parts(link)
+
+        # Store wheels within the root cache_dir
+        return os.path.join(self.cache_dir, "wheels", *parts)
+
+    def get(
+        self,
+        link,            # type: Link
+        package_name,    # type: Optional[str]
+        supported_tags,  # type: List[Tag]
+    ):
+        # type: (...) -> Link
+        candidates = []
+
+        if not package_name:
+            return link
+
+        canonical_package_name = canonicalize_name(package_name)
+        for wheel_name, wheel_dir in self._get_candidates(
+            link, canonical_package_name
+        ):
+            try:
+                wheel = Wheel(wheel_name)
+            except InvalidWheelFilename:
+                continue
+            if canonicalize_name(wheel.name) != canonical_package_name:
+                logger.debug(
+                    "Ignoring cached wheel {} for {} as it "
+                    "does not match the expected distribution name {}.".format(
+                        wheel_name, link, package_name
+                    )
+                )
+                continue
+            if not wheel.supported(supported_tags):
+                # Built for a different python/arch/etc
+                continue
+            candidates.append(
+                (
+                    wheel.support_index_min(supported_tags),
+                    wheel_name,
+                    wheel_dir,
+                )
+            )
+
+        if not candidates:
+            return link
+
+        _, wheel_name, wheel_dir = min(candidates)
+        return Link(path_to_url(os.path.join(wheel_dir, wheel_name)))
+
+
+class EphemWheelCache(SimpleWheelCache):
+    """A SimpleWheelCache that creates it's own temporary cache directory
+    """
+
+    def __init__(self, format_control):
+        # type: (FormatControl) -> None
+        self._temp_dir = TempDirectory(kind="ephem-wheel-cache")
+
+        super(EphemWheelCache, self).__init__(
+            self._temp_dir.path, format_control
+        )
+
+    def cleanup(self):
+        # type: () -> None
+        self._temp_dir.cleanup()
+
+
+class WheelCache(Cache):
+    """Wraps EphemWheelCache and SimpleWheelCache into a single Cache
+
+    This Cache allows for gracefully degradation, using the ephem wheel cache
+    when a certain link is not found in the simple wheel cache first.
+    """
+
+    def __init__(self, cache_dir, format_control):
+        # type: (str, FormatControl) -> None
+        super(WheelCache, self).__init__(
+            cache_dir, format_control, {'binary'}
+        )
+        self._wheel_cache = SimpleWheelCache(cache_dir, format_control)
+        self._ephem_cache = EphemWheelCache(format_control)
+
+    def get_path_for_link_legacy(self, link):
+        # type: (Link) -> str
+        return self._wheel_cache.get_path_for_link_legacy(link)
+
+    def get_path_for_link(self, link):
+        # type: (Link) -> str
+        return self._wheel_cache.get_path_for_link(link)
+
+    def get_ephem_path_for_link(self, link):
+        # type: (Link) -> str
+        return self._ephem_cache.get_path_for_link(link)
+
+    def get(
+        self,
+        link,            # type: Link
+        package_name,    # type: Optional[str]
+        supported_tags,  # type: List[Tag]
+    ):
+        # type: (...) -> Link
+        retval = self._wheel_cache.get(
+            link=link,
+            package_name=package_name,
+            supported_tags=supported_tags,
+        )
+        if retval is not link:
+            return retval
+
+        return self._ephem_cache.get(
+            link=link,
+            package_name=package_name,
+            supported_tags=supported_tags,
+        )
+
+    def cleanup(self):
+        # type: () -> None
+        self._wheel_cache.cleanup()
+        self._ephem_cache.cleanup()

jax_env/lib/python3.8/site-packages/pip/_internal/cli/__init__.py

@@ -0,0 +1,4 @@
+"""Subpackage containing all of pip's command line interface related code
+"""
+
+# This file intentionally does not import submodules