Merge remote-tracking branch 'origin/main'

sanchit-gandhi · Aug 4, 2023 · 45bff9d · 45bff9d
2 parents 06f123d + cf0d4f2
commit 45bff9d
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 5 deletions.
diff --git a/benchmarks/run_pmap.py b/benchmarks/run_pmap.py
@@ -3,6 +3,7 @@
 import jax
 import jax.numpy as jnp
 from datasets import concatenate_datasets, load_dataset
+from flax import jax_utils
 from flax.training.common_utils import shard
 from transformers import FlaxWhisperForConditionalGeneration, WhisperProcessor
 
@@ -17,7 +18,7 @@
     dtype=jnp.bfloat16,
 )
 
-params = model.to_bf16(params)
+params = jax_utils.replicate(params)
 
 
 def generate_fn(batch):

diff --git a/whisper_jax/modeling_flax_whisper.py b/whisper_jax/modeling_flax_whisper.py
@@ -1621,8 +1621,8 @@ def prepare_inputs_for_generation(
         self,
         decoder_input_ids,
         max_length,
-        attention_mask: Optional[jnp.DeviceArray] = None,
-        decoder_attention_mask: Optional[jnp.DeviceArray] = None,
+        attention_mask: Optional[jax.Array] = None,
+        decoder_attention_mask: Optional[jax.Array] = None,
         encoder_outputs=None,
         **kwargs,
     ):

diff --git a/whisper_jax/pipeline.py b/whisper_jax/pipeline.py
@@ -34,8 +34,6 @@
 from .train_state import InferenceState
 
 
-jax.config.update("jax_array", True)
-
 logger = logging.get_logger(__name__)
 
 # 2D parameter and activation partitioning for DP