examples/research_projects/jax-projects/wav2vec2/run_wav2vec2_pretrain_flax.py/FlaxDataCollatorForWav2Vec2Pretraining/__call__
class FlaxDataCollatorForWav2Vec2Pretraining:    def __call__(self, features: List[Dict[str, Union[List[int], np.ndarray]]]) -> Dict[str, np.ndarray]:
        # reformat list to dict and set to pytorch format
        batch = self.feature_extractor.pad(
            features,
            max_length=self.max_length,
            padding=self.padding,
            pad_to_multiple_of=self.pad_to_multiple_of,
            return_tensors="np",
        )
        mask_indices_seq_length = self.model._get_feat_extract_output_lengths(batch["input_values"].shape[-1])

        batch_size = batch["input_values"].shape[0]

        attention_mask = None
        if batch["attention_mask"] is not None:
            output_lengths = self.model._get_feat_extract_output_lengths(batch["attention_mask"].sum(-1))
            attention_mask = np.zeros((batch_size, mask_indices_seq_length), dtype=np.int8)

            # these two operations makes sure that all values
            # before the output lengths indices are attended to
            attention_mask[(np.arange(attention_mask.shape[0]), output_lengths - 1)] = 1
            attention_mask = jnp.flip(jnp.flip(attention_mask, -1).cumsum(-1), -1).astype("bool")

        # sample randomly masked indices
        batch["mask_time_indices"] = _compute_mask_indices(
            (batch_size, mask_indices_seq_length),
            self.model.config.mask_time_prob,
            self.model.config.mask_time_length,
            attention_mask=attention_mask,
            min_masks=2,
        )

        # sample indices to take for negative vectors
        batch["sampled_negative_indices"] = _sample_negative_indices(
            (batch["mask_time_indices"].shape + (self.model.config.proj_codevector_dim,)),
            self.model.config.num_negatives,
            attention_mask=attention_mask,
        )

        return batch