sanchit-gandhi HF staff commited on
Commit
1d2ae3c
1 Parent(s): a55fa23
requirements.txt CHANGED
@@ -1,7 +1,8 @@
1
  --find-links https://storage.googleapis.com/jax-releases/libtpu_releases.html
2
  jax[tpu]
3
- git+https://github.com/sanchit-gandhi/whisper-jax.git
4
- requests
5
- yt-dlp>=2023.3.4
6
  torch
7
- transformers>=4.40.0
 
 
 
 
 
1
  --find-links https://storage.googleapis.com/jax-releases/libtpu_releases.html
2
  jax[tpu]
 
 
 
3
  torch
4
+ transformers>=4.40.0
5
+ flax
6
+ cached-property
7
+ requests
8
+ yt-dlp>=2023.3.4
whisper_jax/__init__.py ADDED
@@ -0,0 +1,21 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # coding=utf-8
2
+ # Copyright 2023 The HuggingFace Inc. team.
3
+ #
4
+ # Licensed under the Apache License, Version 2.0 (the "License");
5
+ # you may not use this file except in compliance with the License.
6
+ # You may obtain a copy of the License at
7
+ #
8
+ # http://www.apache.org/licenses/LICENSE-2.0
9
+ #
10
+ # Unless required by applicable law or agreed to in writing, software
11
+ # distributed under the License is distributed on an "AS IS" BASIS,
12
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
+ # See the License for the specific language governing permissions and
14
+ # limitations under the License.
15
+
16
+ __version__ = "0.0.1"
17
+
18
+ from .modeling_flax_whisper import FlaxWhisperForConditionalGeneration
19
+ from .partitioner import PjitPartitioner
20
+ from .pipeline import FlaxWhisperPipline
21
+ from .train_state import InferenceState
whisper_jax/layers.py ADDED
@@ -0,0 +1,1310 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # coding=utf-8
2
+ # Copyright 2023 The T5X Authors and The HuggingFace Inc. team.
3
+ #
4
+ # Licensed under the Apache License, Version 2.0 (the "License");
5
+ # you may not use this file except in compliance with the License.
6
+ # You may obtain a copy of the License at
7
+ #
8
+ # http://www.apache.org/licenses/LICENSE-2.0
9
+ #
10
+ # Unless required by applicable law or agreed to in writing, software
11
+ # distributed under the License is distributed on an "AS IS" BASIS,
12
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
+ # See the License for the specific language governing permissions and
14
+ # limitations under the License.
15
+
16
+ """Dense attention classes and mask/weighting functions."""
17
+
18
+ # pylint: disable=attribute-defined-outside-init,g-bare-generic
19
+
20
+ import dataclasses
21
+ import functools
22
+ import operator
23
+ from typing import Any, Callable, Iterable, List, Optional, Sequence, Tuple, Union
24
+
25
+ import jax
26
+ import jax.numpy as jnp
27
+ import numpy as np
28
+ from flax import linen as nn
29
+ from flax.linen import partitioning as nn_partitioning
30
+ from flax.linen.dtypes import promote_dtype
31
+ from jax import lax, random
32
+
33
+
34
+ # from flax.linen.partitioning import param_with_axes, with_sharding_constraint
35
+ param_with_axes = nn_partitioning.param_with_axes
36
+ with_sharding_constraint = nn_partitioning.with_sharding_constraint
37
+
38
+
39
+ # Type annotations
40
+ Array = jnp.ndarray
41
+ DType = jnp.dtype
42
+ PRNGKey = jnp.ndarray
43
+ Shape = Iterable[int]
44
+ Activation = Callable[..., Array]
45
+ PrecisionLike = Union[None, str, lax.Precision, Tuple[str, str], Tuple[lax.Precision, lax.Precision]]
46
+ DotGeneralT = Callable[..., Array]
47
+ ConvGeneralDilatedT = Callable[..., Array]
48
+ PaddingLike = Union[str, int, Sequence[Union[int, Tuple[int, int]]]]
49
+ LaxPadding = Union[str, Sequence[Tuple[int, int]]]
50
+
51
+ # Parameter initializers.
52
+ Initializer = Callable[[PRNGKey, Shape, DType], Array]
53
+ InitializerAxis = Union[int, Tuple[int, ...]]
54
+ NdInitializer = Callable[[PRNGKey, Shape, DType, InitializerAxis, InitializerAxis], Array]
55
+
56
+ default_embed_init = nn.initializers.variance_scaling(1.0, "fan_in", "normal", out_axis=0)
57
+
58
+
59
+ # ------------------------------------------------------------------------------
60
+ # Temporary inlined JAX N-d initializer code
61
+ # TODO(levskaya): remove once new JAX release is out.
62
+ # ------------------------------------------------------------------------------
63
+ def _compute_fans(shape: jax.core.NamedShape, in_axis=-2, out_axis=-1):
64
+ """Inlined JAX `nn.initializer._compute_fans`."""
65
+ if isinstance(in_axis, int):
66
+ in_size = shape[in_axis]
67
+ else:
68
+ in_size = int(np.prod([shape[i] for i in in_axis]))
69
+ if isinstance(out_axis, int):
70
+ out_size = shape[out_axis]
71
+ else:
72
+ out_size = int(np.prod([shape[i] for i in out_axis]))
73
+ receptive_field_size = shape.total / in_size / out_size
74
+ fan_in = in_size * receptive_field_size
75
+ fan_out = out_size * receptive_field_size
76
+ return fan_in, fan_out
77
+
78
+
79
+ def variance_scaling(scale, mode, distribution, in_axis=-2, out_axis=-1, dtype=jnp.float_):
80
+ """Inlined JAX `nn.initializer.variance_scaling`."""
81
+
82
+ def init(key, shape, dtype=dtype):
83
+ return jnp.zeros(shape, dtype=dtype)
84
+ dtype = jax.dtypes.canonicalize_dtype(dtype)
85
+ shape = jax.core.as_named_shape(shape)
86
+ fan_in, fan_out = _compute_fans(shape, in_axis, out_axis)
87
+ if mode == "fan_in":
88
+ denominator = fan_in
89
+ elif mode == "fan_out":
90
+ denominator = fan_out
91
+ elif mode == "fan_avg":
92
+ denominator = (fan_in + fan_out) / 2
93
+ else:
94
+ raise ValueError("invalid mode for variance scaling initializer: {}".format(mode))
95
+ variance = jnp.array(scale / denominator, dtype=dtype)
96
+
97
+ if distribution == "truncated_normal":
98
+ # constant is stddev of standard normal truncated to (-2, 2)
99
+ stddev = jnp.sqrt(variance) / jnp.array(0.87962566103423978, dtype)
100
+ return random.truncated_normal(key, -2, 2, shape, dtype) * stddev
101
+ elif distribution == "normal":
102
+ return random.normal(key, shape, dtype) * jnp.sqrt(variance)
103
+ elif distribution == "uniform":
104
+ return random.uniform(key, shape, dtype, -1) * jnp.sqrt(3 * variance)
105
+ else:
106
+ raise ValueError("invalid distribution for variance scaling " "initializer: {}".format(distribution))
107
+
108
+ return init
109
+
110
+
111
+ # ------------------------------------------------------------------------------
112
+
113
+
114
+ def nd_dense_init(scale, mode, distribution):
115
+ """Initializer with in_axis, out_axis set at call time."""
116
+
117
+ def init_fn(key, shape, dtype, in_axis, out_axis):
118
+ fn = variance_scaling(scale, mode, distribution, in_axis, out_axis)
119
+ return fn(key, shape, dtype)
120
+
121
+ return init_fn
122
+
123
+
124
+ def dot_product_attention(
125
+ query: Array,
126
+ key: Array,
127
+ value: Array,
128
+ bias: Optional[Array] = None,
129
+ dropout_rng: Optional[PRNGKey] = None,
130
+ dropout_rate: float = 0.0,
131
+ deterministic: bool = False,
132
+ dtype: DType = jnp.float32,
133
+ float32_logits: bool = False,
134
+ ):
135
+ """Computes dot-product attention given query, key, and value.
136
+
137
+ This is the core function for applying attention based on
138
+ https://arxiv.org/abs/1706.03762. It calculates the attention weights given
139
+ query and key and combines the values using the attention weights.
140
+
141
+ Args:
142
+ query: queries for calculating attention with shape of `[batch, q_length,
143
+ num_heads, qk_depth_per_head]`.
144
+ key: keys for calculating attention with shape of `[batch, kv_length,
145
+ num_heads, qk_depth_per_head]`.
146
+ value: values to be used in attention with shape of `[batch, kv_length,
147
+ num_heads, v_depth_per_head]`.
148
+ bias: bias for the attention weights. This should be broadcastable to the
149
+ shape `[batch, num_heads, q_length, kv_length]` This can be used for
150
+ incorporating causal masks, padding masks, proximity bias, etc.
151
+ dropout_rng: JAX PRNGKey: to be used for dropout
152
+ dropout_rate: dropout rate
153
+ deterministic: bool, deterministic or not (to apply dropout)
154
+ dtype: the dtype of the computation (default: float32)
155
+ float32_logits: bool, if True then compute logits in float32 to avoid
156
+ numerical issues with bfloat16.
157
+
158
+ Returns:
159
+ Output of shape `[batch, length, num_heads, v_depth_per_head]`.
160
+ """
161
+ assert key.ndim == query.ndim == value.ndim, "q, k, v must have same rank."
162
+ assert query.shape[:-3] == key.shape[:-3] == value.shape[:-3], "q, k, v batch dims must match."
163
+ assert query.shape[-2] == key.shape[-2] == value.shape[-2], "q, k, v num_heads must match."
164
+ assert key.shape[-3] == value.shape[-3], "k, v lengths must match."
165
+ assert query.shape[-1] == key.shape[-1], "q, k depths must match."
166
+
167
+ # Casting logits and softmax computation for float32 for model stability.
168
+ if float32_logits:
169
+ query = query.astype(jnp.float32)
170
+ key = key.astype(jnp.float32)
171
+
172
+ # `attn_weights`: [batch, num_heads, q_length, kv_length]
173
+ attn_weights = jnp.einsum("bqhd,bkhd->bhqk", query, key)
174
+
175
+ # Apply attention bias: masking, dropout, proximity bias, etc.
176
+ if bias is not None:
177
+ attn_weights = attn_weights + bias.astype(attn_weights.dtype)
178
+
179
+ # Normalize the attention weights across `kv_length` dimension.
180
+ attn_weights = jax.nn.softmax(attn_weights).astype(dtype)
181
+
182
+ # Apply attention dropout.
183
+ if not deterministic and dropout_rate > 0.0:
184
+ keep_prob = 1.0 - dropout_rate
185
+ # T5 broadcasts along the "length" dim, but unclear which one that
186
+ # corresponds to in positional dimensions here, assuming query dim.
187
+ dropout_shape = list(attn_weights.shape)
188
+ dropout_shape[-2] = 1
189
+ keep = random.bernoulli(dropout_rng, keep_prob, dropout_shape)
190
+ keep = jnp.broadcast_to(keep, attn_weights.shape)
191
+ multiplier = keep.astype(attn_weights.dtype) / jnp.asarray(keep_prob, dtype=dtype)
192
+ attn_weights = attn_weights * multiplier
193
+
194
+ # Take the linear combination of `value`.
195
+ return jnp.einsum("bhqk,bkhd->bqhd", attn_weights, value)
196
+
197
+
198
+ dynamic_vector_slice_in_dim = jax.vmap(lax.dynamic_slice_in_dim, in_axes=(None, 0, None, None))
199
+
200
+
201
+ class MultiHeadDotProductAttention(nn.Module):
202
+ """Multi-head dot-product attention.
203
+
204
+ Attributes:
205
+ num_heads: number of attention heads. Features (i.e. inputs_q.shape[-1])
206
+ should be divisible by the number of heads.
207
+ head_dim: dimension of each head.
208
+ dtype: the dtype of the computation.
209
+ dropout_rate: dropout rate
210
+ kernel_init: initializer for the kernel of the Dense layers.
211
+ float32_logits: bool, if True then compute logits in float32 to avoid
212
+ numerical issues with bfloat16.
213
+ """
214
+
215
+ num_heads: int
216
+ head_dim: int
217
+ dtype: DType = jnp.float32
218
+ dropout_rate: float = 0.0
219
+ kernel_init: NdInitializer = nd_dense_init(1.0, "fan_in", "normal")
220
+ float32_logits: bool = False # computes logits in float32 for stability.
221
+
222
+ @nn.compact
223
+ def __call__(
224
+ self,
225
+ inputs_q: Array,
226
+ inputs_kv: Array,
227
+ mask: Optional[Array] = None,
228
+ bias: Optional[Array] = None,
229
+ *,
230
+ decode: bool = False,
231
+ deterministic: bool = False,
232
+ ) -> Array:
233
+ """Applies multi-head dot product attention on the input data.
234
+
235
+ Projects the inputs into multi-headed query, key, and value vectors,
236
+ applies dot-product attention and project the results to an output vector.
237
+
238
+ There are two modes: decoding and non-decoding (e.g., training). The mode is
239
+ determined by `decode` argument. For decoding, this method is called twice,
240
+ first to initialize the cache and then for an actual decoding process. The
241
+ two calls are differentiated by the presence of 'cached_key' in the variable
242
+ dict. In the cache initialization stage, the cache variables are initialized
243
+ as zeros and will be filled in the subsequent decoding process.
244
+
245
+ In the cache initialization call, `inputs_q` has a shape [batch, length,
246
+ q_features] and `inputs_kv`: [batch, length, kv_features]. During the
247
+ incremental decoding stage, query, key and value all have the shape [batch,
248
+ 1, qkv_features] corresponding to a single step.
249
+
250
+ Args:
251
+ inputs_q: input queries of shape `[batch, q_length, q_features]`.
252
+ inputs_kv: key/values of shape `[batch, kv_length, kv_features]`.
253
+ mask: attention mask of shape `[batch, num_heads, q_length, kv_length]`.
254
+ bias: attention bias of shape `[batch, num_heads, q_length, kv_length]`.
255
+ decode: Whether to prepare and use an autoregressive cache.
256
+ deterministic: Disables dropout if set to True.
257
+
258
+ Returns:
259
+ output of shape `[batch, length, q_features]`.
260
+ """
261
+ projection = functools.partial(
262
+ DenseGeneral,
263
+ axis=-1,
264
+ features=(self.num_heads, self.head_dim),
265
+ kernel_axes=("embed", "heads", "kv"),
266
+ dtype=self.dtype,
267
+ )
268
+
269
+ # NOTE: T5 does not explicitly rescale the attention logits by
270
+ # 1/sqrt(depth_kq)! This is folded into the initializers of the
271
+ # linear transformations, which is equivalent under Adafactor.
272
+ depth_scaling = jnp.sqrt(self.head_dim).astype(self.dtype)
273
+
274
+ def query_init(*args):
275
+ return self.kernel_init(*args) / depth_scaling
276
+
277
+ # Project inputs_q to multi-headed q/k/v
278
+ # dimensions are then [batch, length, num_heads, head_dim]
279
+ query = projection(kernel_init=query_init, name="query")(inputs_q)
280
+ key = projection(kernel_init=self.kernel_init, name="key")(inputs_kv)
281
+ value = projection(kernel_init=self.kernel_init, name="value")(inputs_kv)
282
+
283
+ query = with_sharding_constraint(query, ("batch", "length", "heads", "kv"))
284
+ key = with_sharding_constraint(key, ("batch", "length", "heads", "kv"))
285
+ value = with_sharding_constraint(value, ("batch", "length", "heads", "kv"))
286
+
287
+ if decode:
288
+ # Detect if we're initializing by absence of existing cache data.
289
+ is_initialized = self.has_variable("cache", "cached_key")
290
+
291
+ # The key and value have dimension [batch, length, num_heads, head_dim],
292
+ # but we cache them as [batch, num_heads, head_dim, length] as a TPU
293
+ # fusion optimization. This also enables the "scatter via one-hot
294
+ # broadcast" trick, which means we do a one-hot broadcast instead of a
295
+ # scatter/gather operations, resulting in a 3-4x speedup in practice.
296
+ def swap_dims(x):
297
+ return x[:-3] + tuple(x[i] for i in [-2, -1, -3])
298
+
299
+ cached_key = self.variable("cache", "cached_key", jnp.zeros, swap_dims(key.shape), key.dtype)
300
+ cached_value = self.variable("cache", "cached_value", jnp.zeros, swap_dims(value.shape), value.dtype)
301
+ cache_index = self.variable("cache", "cache_index", lambda: jnp.array(0, dtype=jnp.int32))
302
+ if is_initialized:
303
+ batch, num_heads, head_dim, length = cached_key.value.shape
304
+ # During fast autoregressive decoding, we feed one position at a time,
305
+ # and cache the keys and values step by step.
306
+ # Sanity shape check of cached key against input query.
307
+ expected_shape = (batch, 1, num_heads, head_dim)
308
+ if expected_shape != query.shape:
309
+ raise ValueError(
310
+ "Autoregressive cache shape error, "
311
+ "expected query shape %s instead got %s." % (expected_shape, query.shape)
312
+ )
313
+
314
+ # Create a OHE of the current index. NOTE: the index is increased below.
315
+ cur_index = cache_index.value
316
+ one_hot_indices = jax.nn.one_hot(cur_index, length, dtype=key.dtype)
317
+ # In order to update the key, value caches with the current key and
318
+ # value, we move the length axis to the back, similar to what we did for
319
+ # the cached ones above.
320
+ # Note these are currently the key and value of a single position, since
321
+ # we feed one position at a time.
322
+ one_token_key = jnp.moveaxis(key, -3, -1)
323
+ one_token_value = jnp.moveaxis(value, -3, -1)
324
+ # Update key, value caches with our new 1d spatial slices.
325
+ # We implement an efficient scatter into the cache via one-hot
326
+ # broadcast and addition.
327
+ key = cached_key.value + one_token_key * one_hot_indices
328
+ value = cached_value.value + one_token_value * one_hot_indices
329
+ cached_key.value = key
330
+ cached_value.value = value
331
+ cache_index.value = cache_index.value + 1
332
+ # Move the keys and values back to their original shapes.
333
+ key = jnp.moveaxis(key, -1, -3)
334
+ value = jnp.moveaxis(value, -1, -3)
335
+
336
+ # Causal mask for cached decoder self-attention: our single query
337
+ # position should only attend to those key positions that have already
338
+ # been generated and cached, not the remaining zero elements.
339
+ mask = combine_masks(
340
+ mask,
341
+ jnp.broadcast_to(
342
+ jnp.arange(length) <= cur_index,
343
+ # (1, 1, length) represent (head dim, query length, key length)
344
+ # query length is 1 because during decoding we deal with one
345
+ # index.
346
+ # The same mask is applied to all batch elements and heads.
347
+ (batch, 1, 1, length),
348
+ ),
349
+ )
350
+
351
+ # Grab the correct relative attention bias during decoding. This is
352
+ # only required during single step decoding.
353
+ if bias is not None:
354
+ # The bias is a full attention matrix, but during decoding we only
355
+ # have to take a slice of it.
356
+ # This is equivalent to bias[..., cur_index:cur_index+1, :].
357
+ bias = dynamic_vector_slice_in_dim(jnp.squeeze(bias, axis=0), jnp.reshape(cur_index, (-1)), 1, -2)
358
+
359
+ # Convert the boolean attention mask to an attention bias.
360
+ if mask is not None:
361
+ # attention mask in the form of attention bias
362
+ attention_bias = lax.select(
363
+ mask > 0, jnp.full(mask.shape, 0.0).astype(self.dtype), jnp.full(mask.shape, -1e10).astype(self.dtype)
364
+ )
365
+ else:
366
+ attention_bias = None
367
+
368
+ # Add provided bias term (e.g. relative position embedding).
369
+ if bias is not None:
370
+ attention_bias = combine_biases(attention_bias, bias)
371
+
372
+ dropout_rng = None
373
+ if not deterministic and self.dropout_rate > 0.0:
374
+ dropout_rng = self.make_rng("dropout")
375
+
376
+ # Apply attention.
377
+ x = dot_product_attention(
378
+ query,
379
+ key,
380
+ value,
381
+ bias=attention_bias,
382
+ dropout_rng=dropout_rng,
383
+ dropout_rate=self.dropout_rate,
384
+ deterministic=deterministic,
385
+ dtype=self.dtype,
386
+ float32_logits=self.float32_logits,
387
+ )
388
+
389
+ # Back to the original inputs dimensions.
390
+ out = DenseGeneral(
391
+ features=inputs_q.shape[-1], # output dim is set to the input dim.
392
+ axis=(-2, -1),
393
+ kernel_init=self.kernel_init,
394
+ kernel_axes=("heads", "kv", "embed"),
395
+ dtype=self.dtype,
396
+ name="out",
397
+ )(x)
398
+ return out
399
+
400
+
401
+ def _normalize_axes(axes: Iterable[int], ndim: int) -> Tuple[int]:
402
+ # A tuple by convention. len(axes_tuple) then also gives the rank efficiently.
403
+ return tuple([ax if ax >= 0 else ndim + ax for ax in axes])
404
+
405
+
406
+ def _canonicalize_tuple(x):
407
+ if isinstance(x, Iterable):
408
+ return tuple(x)
409
+ else:
410
+ return (x,)
411
+
412
+
413
+ # ------------------------------------------------------------------------------
414
+ # DenseGeneral for attention layers.
415
+ # ------------------------------------------------------------------------------
416
+ class DenseGeneral(nn.Module):
417
+ """A linear transformation (without bias) with flexible axes.
418
+
419
+ Attributes:
420
+ features: tuple with numbers of output features.
421
+ axis: tuple with axes to apply the transformation on.
422
+ dtype: the dtype of the computation (default: float32).
423
+ kernel_init: initializer function for the weight matrix.
424
+ """
425
+
426
+ features: Union[Iterable[int], int]
427
+ axis: Union[Iterable[int], int] = -1
428
+ dtype: DType = jnp.float32
429
+ params_dtype: DType = jnp.float32
430
+ kernel_init: NdInitializer = nd_dense_init(1.0, "fan_in", "normal")
431
+ kernel_axes: Tuple[str, ...] = ()
432
+ use_bias: bool = True
433
+ bias_init: Any = nn.initializers.zeros
434
+
435
+ @nn.compact
436
+ def __call__(self, inputs: Array) -> Array:
437
+ """Applies a linear transformation to the inputs along multiple dimensions.
438
+
439
+ Args:
440
+ inputs: The nd-array to be transformed.
441
+
442
+ Returns:
443
+ The transformed input.
444
+ """
445
+ features = _canonicalize_tuple(self.features)
446
+ axis = _canonicalize_tuple(self.axis)
447
+
448
+ inputs = jnp.asarray(inputs, self.dtype)
449
+ axis = _normalize_axes(axis, inputs.ndim)
450
+
451
+ kernel_shape = tuple([inputs.shape[ax] for ax in axis]) + features
452
+ kernel_in_axis = np.arange(len(axis))
453
+ kernel_out_axis = np.arange(len(axis), len(axis) + len(features))
454
+ kernel = param_with_axes(
455
+ "kernel",
456
+ self.kernel_init,
457
+ kernel_shape,
458
+ self.params_dtype,
459
+ kernel_in_axis,
460
+ kernel_out_axis,
461
+ axes=self.kernel_axes,
462
+ )
463
+ if self.use_bias:
464
+ bias = param_with_axes("bias", self.bias_init, features, self.params_dtype, axes=(self.kernel_axes[-1],))
465
+ kernel = jnp.asarray(kernel, self.dtype)
466
+
467
+ contract_ind = tuple(range(0, len(axis)))
468
+ y = lax.dot_general(inputs, kernel, ((axis, contract_ind), ((), ())))
469
+ if self.use_bias:
470
+ bias = jnp.asarray(bias, self.dtype)
471
+ # y += jnp.reshape(bias, (1,) * (y.ndim - 1) + (-1,))
472
+ y += jnp.reshape(bias, (1,) * (len(features) - y.ndim) + bias.shape[:])
473
+ return y
474
+
475
+
476
+ def _convert_to_activation_function(fn_or_string: Union[str, Callable]) -> Callable:
477
+ """Convert a string to an activation function."""
478
+ if fn_or_string == "linear":
479
+ return lambda x: x
480
+ elif isinstance(fn_or_string, str):
481
+ return getattr(nn, fn_or_string)
482
+ elif callable(fn_or_string):
483
+ return fn_or_string
484
+ else:
485
+ raise ValueError("don't know how to convert %s to an activation function" % (fn_or_string,))
486
+
487
+
488
+ class MlpBlock(nn.Module):
489
+ """Transformer MLP / feed-forward block.
490
+
491
+ Attributes:
492
+ intermediate_dim: Shared dimension of hidden layers.
493
+ activations: Type of activations for each layer. Each element is either
494
+ 'linear', a string function name in flax.linen, or a function.
495
+ kernel_init: Kernel function, passed to the dense layers.
496
+ deterministic: Whether the dropout layers should be deterministic.
497
+ intermediate_dropout_rate: Dropout rate used after the intermediate layers.
498
+ dtype: Type for the dense layer.
499
+ """
500
+
501
+ intermediate_dim: int = 2048
502
+ activations: Sequence[Union[str, Callable]] = ("relu",)
503
+ kernel_init: NdInitializer = nd_dense_init(1.0, "fan_in", "truncated_normal")
504
+ intermediate_dropout_rate: float = 0.1
505
+ dtype: Any = jnp.float32
506
+
507
+ @nn.compact
508
+ def __call__(self, inputs, decode: bool = False, deterministic: bool = False):
509
+ """Applies Transformer MlpBlock module."""
510
+ # Iterate over specified MLP input activation functions.
511
+ # e.g. ('relu',) or ('gelu', 'linear') for gated-gelu.
512
+ activations = []
513
+ for idx, act_fn in enumerate(self.activations):
514
+ dense_name = "wi" if len(self.activations) == 1 else f"wi_{idx}"
515
+ x = DenseGeneral(
516
+ self.intermediate_dim,
517
+ dtype=self.dtype,
518
+ kernel_init=self.kernel_init,
519
+ kernel_axes=("embed", "mlp"),
520
+ name=dense_name,
521
+ )(inputs)
522
+ x = _convert_to_activation_function(act_fn)(x)
523
+ activations.append(x)
524
+
525
+ # Take elementwise product of above intermediate activations.
526
+ x = functools.reduce(operator.mul, activations)
527
+ # Apply dropout and final dense output projection.
528
+ x = nn.Dropout(rate=self.intermediate_dropout_rate, broadcast_dims=(-2,))(
529
+ x, deterministic=deterministic
530
+ ) # Broadcast along length.
531
+ x = with_sharding_constraint(x, ("batch", "length", "mlp"))
532
+ output = DenseGeneral(
533
+ inputs.shape[-1], dtype=self.dtype, kernel_init=self.kernel_init, kernel_axes=("mlp", "embed"), name="wo"
534
+ )(x)
535
+ return output
536
+
537
+
538
+ class Embed(nn.Module):
539
+ """A parameterized function from integers [0, n) to d-dimensional vectors.
540
+
541
+ Attributes:
542
+ num_embeddings: number of embeddings.
543
+ features: number of feature dimensions for each embedding.
544
+ dtype: the dtype of the embedding vectors (default: float32).
545
+ embedding_init: embedding initializer.
546
+ one_hot: performs the gather with a one-hot contraction rather than a true
547
+ gather. This is currently needed for SPMD partitioning.
548
+ """
549
+
550
+ num_embeddings: int
551
+ features: int
552
+ cast_input_dtype: Optional[DType] = None
553
+ dtype: DType = jnp.float32
554
+ params_dtype: DType = jnp.float32
555
+ attend_dtype: Optional[DType] = None
556
+ embedding_init: Initializer = default_embed_init
557
+ one_hot: bool = True
558
+ embedding: Array = dataclasses.field(init=False)
559
+
560
+ def setup(self):
561
+ self.embedding = param_with_axes(
562
+ "embedding",
563
+ self.embedding_init,
564
+ (self.num_embeddings, self.features),
565
+ self.params_dtype,
566
+ axes=("vocab", "embed"),
567
+ )
568
+
569
+ def __call__(self, inputs: Array) -> Array:
570
+ """Embeds the inputs along the last dimension.
571
+
572
+ Args:
573
+ inputs: input data, all dimensions are considered batch dimensions.
574
+
575
+ Returns:
576
+ Output which is embedded input data. The output shape follows the input,
577
+ with an additional `features` dimension appended.
578
+ """
579
+ if self.cast_input_dtype:
580
+ inputs = inputs.astype(self.cast_input_dtype)
581
+ if not jnp.issubdtype(inputs.dtype, jnp.integer):
582
+ raise ValueError("Input type must be an integer or unsigned integer.")
583
+ if self.one_hot:
584
+ iota = lax.iota(jnp.int32, self.num_embeddings)
585
+ one_hot = jnp.array(inputs[..., jnp.newaxis] == iota, dtype=self.dtype)
586
+ output = jnp.dot(one_hot, jnp.asarray(self.embedding, self.dtype))
587
+ else:
588
+ output = jnp.asarray(self.embedding, self.dtype)[inputs]
589
+ output = with_sharding_constraint(output, ("batch", "length", "embed"))
590
+ return output
591
+
592
+ def attend(self, query: Array) -> Array:
593
+ """Attend over the embedding using a query array.
594
+
595
+ Args:
596
+ query: array with last dimension equal the feature depth `features` of the
597
+ embedding.
598
+
599
+ Returns:
600
+ An array with final dim `num_embeddings` corresponding to the batched
601
+ inner-product of the array of query vectors against each embedding.
602
+ Commonly used for weight-sharing between embeddings and logit transform
603
+ in NLP models.
604
+ """
605
+ dtype = self.attend_dtype if self.attend_dtype is not None else self.dtype
606
+ return jnp.dot(query, jnp.asarray(self.embedding, dtype).T)
607
+
608
+
609
+ class RelativePositionBiases(nn.Module):
610
+ """Adds T5-style relative positional embeddings to the attention logits.
611
+
612
+ Attributes:
613
+ num_buckets: Number of buckets to bucket distances between key and query
614
+ positions into.
615
+ max_distance: Maximum distance before everything is lumped into the last
616
+ distance bucket.
617
+ num_heads: Number of heads in the attention layer. Each head will get a
618
+ different relative position weighting.
619
+ dtype: Type of arrays through this module.
620
+ embedding_init: initializer for relative embedding table.
621
+ """
622
+
623
+ num_buckets: int
624
+ max_distance: int
625
+ num_heads: int
626
+ dtype: Any
627
+ embedding_init: Callable[..., Array] = nn.linear.default_embed_init
628
+
629
+ @staticmethod
630
+ def _relative_position_bucket(relative_position, bidirectional=True, num_buckets=32, max_distance=128):
631
+ """Translate relative position to a bucket number for relative attention.
632
+
633
+ The relative position is defined as memory_position - query_position, i.e.
634
+ the distance in tokens from the attending position to the attended-to
635
+ position. If bidirectional=False, then positive relative positions are
636
+ invalid.
637
+ We use smaller buckets for small absolute relative_position and larger
638
+ buckets for larger absolute relative_positions. All relative
639
+ positions >=max_distance map to the same bucket. All relative
640
+ positions <=-max_distance map to the same bucket. This should allow for
641
+ more graceful generalization to longer sequences than the model has been
642
+ trained on.
643
+
644
+ Args:
645
+ relative_position: an int32 array
646
+ bidirectional: a boolean - whether the attention is bidirectional
647
+ num_buckets: an integer
648
+ max_distance: an integer
649
+
650
+ Returns:
651
+ a Tensor with the same shape as relative_position, containing int32
652
+ values in the range [0, num_buckets)
653
+ """
654
+ ret = 0
655
+ n = -relative_position
656
+ if bidirectional:
657
+ num_buckets //= 2
658
+ ret += (n < 0).astype(np.int32) * num_buckets
659
+ n = np.abs(n)
660
+ else:
661
+ n = np.maximum(n, 0)
662
+ # now n is in the range [0, inf)
663
+ max_exact = num_buckets // 2
664
+ is_small = n < max_exact
665
+ val_if_large = max_exact + (
666
+ np.log(n.astype(np.float32) / max_exact + np.finfo(np.float32).eps)
667
+ / np.log(max_distance / max_exact)
668
+ * (num_buckets - max_exact)
669
+ ).astype(np.int32)
670
+ val_if_large = np.minimum(val_if_large, num_buckets - 1)
671
+ ret += np.where(is_small, n, val_if_large)
672
+ return ret
673
+
674
+ @nn.compact
675
+ def __call__(self, qlen, klen, bidirectional=True):
676
+ """Produce relative position embedding attention biases.
677
+
678
+ Args:
679
+ qlen: attention query length.
680
+ klen: attention key length.
681
+ bidirectional: whether to allow positive memory-query relative position
682
+ embeddings.
683
+
684
+ Returns:
685
+ output: `(1, len, q_len, k_len)` attention bias
686
+ """
687
+ # TODO(levskaya): should we be computing this w. numpy as a program
688
+ # constant?
689
+ context_position = np.arange(qlen, dtype=jnp.int32)[:, None]
690
+ memory_position = np.arange(klen, dtype=jnp.int32)[None, :]
691
+ relative_position = memory_position - context_position # shape (qlen, klen)
692
+ rp_bucket = self._relative_position_bucket(
693
+ relative_position,
694
+ bidirectional=bidirectional,
695
+ num_buckets=self.num_buckets,
696
+ max_distance=self.max_distance,
697
+ )
698
+ relative_attention_bias = param_with_axes(
699
+ "rel_embedding",
700
+ self.embedding_init,
701
+ (self.num_heads, self.num_buckets),
702
+ jnp.float32,
703
+ axes=("heads", "relpos_buckets"),
704
+ )
705
+
706
+ relative_attention_bias = jnp.asarray(relative_attention_bias, self.dtype)
707
+ # Instead of using a slow gather, we create a leading-dimension one-hot
708
+ # array from rp_bucket and use it to perform the gather-equivalent via a
709
+ # contraction, i.e.:
710
+ # (num_head, num_buckets) x (num_buckets one-hot, qlen, klen).
711
+ # This is equivalent to relative_attention_bias[:, rp_bucket]
712
+ bcast_iota = lax.broadcasted_iota(jnp.int32, (self.num_buckets, 1, 1), 0)
713
+ rp_bucket_one_hot = jnp.array(rp_bucket[jnp.newaxis, ...] == bcast_iota, dtype=self.dtype)
714
+ # --> shape (qlen, klen, num_heads)
715
+ values = lax.dot_general(
716
+ relative_attention_bias, rp_bucket_one_hot, (((1,), (0,)), ((), ())) # rhs, lhs contracting dims
717
+ ) # no batched dims
718
+ # Add a singleton batch dimension.
719
+ # --> shape (1, num_heads, qlen, klen)
720
+ return values[jnp.newaxis, ...]
721
+
722
+
723
+ # ------------------------------------------------------------------------------
724
+ # T5 Layernorm - no subtraction of mean or bias.
725
+ # ------------------------------------------------------------------------------
726
+ # class LayerNorm(nn.Module):
727
+ # """T5 Layer normalization operating on the last axis of the input data."""
728
+ # epsilon: float = 1e-6
729
+ # dtype: Any = jnp.float32
730
+ # scale_init: Initializer = nn.initializers.ones
731
+
732
+ # @nn.compact
733
+ # def __call__(self, x: jnp.ndarray) -> jnp.ndarray:
734
+ # """Applies layer normalization on the input."""
735
+ # x = jnp.asarray(x, jnp.float32)
736
+ # features = x.shape[-1]
737
+ # mean2 = jnp.mean(lax.square(x), axis=-1, keepdims=True)
738
+ # y = jnp.asarray(x * lax.rsqrt(mean2 + self.epsilon), self.dtype)
739
+ # scale = param_with_axes(
740
+ # 'scale', self.scale_init, (features,), jnp.float32, axes=('embed',))
741
+
742
+ # scale = jnp.asarray(scale, self.dtype)
743
+ # return y * scale
744
+
745
+
746
+ class LayerNorm(nn.Module):
747
+ """Layer normalization (https://arxiv.org/abs/1607.06450).
748
+ Operates on the last axis of the input data.
749
+ It normalizes the activations of the layer for each given example in a
750
+ batch independently, rather than across a batch like Batch Normalization.
751
+ i.e. applies a transformation that maintains the mean activation within
752
+ each example close to 0 and the activation standard deviation close to 1.
753
+ Attributes:
754
+ epsilon: A small float added to variance to avoid dividing by zero.
755
+ dtype: the dtype of the computation (default: float32).
756
+ use_bias: If True, bias (beta) is added.
757
+ use_scale: If True, multiply by scale (gamma). When the next layer is linear
758
+ (also e.g. nn.relu), this can be disabled since the scaling will be done
759
+ by the next layer.
760
+ bias_init: Initializer for bias, by default, zero.
761
+ scale_init: Initializer for scale, by default, one.
762
+ """
763
+
764
+ epsilon: float = 1e-6
765
+ dtype: Any = jnp.float32
766
+ params_dtype: DType = jnp.float32
767
+ use_bias: bool = True
768
+ use_scale: bool = True
769
+ bias_init: Callable[[PRNGKey, Shape, Any], Array] = nn.initializers.zeros
770
+ scale_init: Callable[[PRNGKey, Shape, Any], Array] = nn.initializers.ones
771
+
772
+ @nn.compact
773
+ def __call__(self, x):
774
+ """Applies layer normalization on the input.
775
+ Args:
776
+ x: the inputs
777
+ Returns:
778
+ Normalized inputs (the same shape as inputs).
779
+ """
780
+ x = jnp.asarray(x, jnp.float32)
781
+ features = x.shape[-1]
782
+ mean = jnp.mean(x, axis=-1, keepdims=True)
783
+ mean2 = jnp.mean(lax.square(x), axis=-1, keepdims=True)
784
+ var = mean2 - lax.square(mean)
785
+ mul = lax.rsqrt(var + self.epsilon)
786
+ if self.use_scale:
787
+ scale = param_with_axes("scale", self.scale_init, (features,), self.params_dtype, axes=("embed",))
788
+ mul = mul * jnp.asarray(scale, self.dtype)
789
+ y = (x - mean) * mul
790
+ if self.use_bias:
791
+ bias = param_with_axes("bias", self.bias_init, (features,), self.params_dtype, axes=("embed",))
792
+ y = y + jnp.asarray(bias, self.dtype)
793
+ return jnp.asarray(y, self.dtype)
794
+
795
+
796
+ # ------------------------------------------------------------------------------
797
+ # Mask-making utility functions.
798
+ # ------------------------------------------------------------------------------
799
+ def make_attention_mask(
800
+ query_input: Array,
801
+ key_input: Array,
802
+ pairwise_fn: Callable = jnp.multiply,
803
+ extra_batch_dims: int = 0,
804
+ dtype: DType = jnp.float32,
805
+ ) -> Array:
806
+ """Mask-making helper for attention weights.
807
+
808
+ In case of 1d inputs (i.e., `[batch, len_q]`, `[batch, len_kv]`, the
809
+ attention weights will be `[batch, heads, len_q, len_kv]` and this
810
+ function will produce `[batch, 1, len_q, len_kv]`.
811
+
812
+ Args:
813
+ query_input: a batched, flat input of query_length size
814
+ key_input: a batched, flat input of key_length size
815
+ pairwise_fn: broadcasting elementwise comparison function
816
+ extra_batch_dims: number of extra batch dims to add singleton axes for, none
817
+ by default
818
+ dtype: mask return dtype
819
+
820
+ Returns:
821
+ A `[batch, 1, len_q, len_kv]` shaped mask for 1d attention.
822
+ """
823
+ # [batch, len_q, len_kv]
824
+ mask = pairwise_fn(
825
+ # [batch, len_q] -> [batch, len_q, 1]
826
+ jnp.expand_dims(query_input, axis=-1),
827
+ # [batch, len_q] -> [batch, 1, len_kv]
828
+ jnp.expand_dims(key_input, axis=-2),
829
+ )
830
+
831
+ # [batch, 1, len_q, len_kv]. This creates the head dim.
832
+ mask = jnp.expand_dims(mask, axis=-3)
833
+ mask = jnp.expand_dims(mask, axis=tuple(range(extra_batch_dims)))
834
+ return mask.astype(dtype)
835
+
836
+
837
+ def make_causal_mask(x: Array, extra_batch_dims: int = 0, dtype: DType = jnp.float32) -> Array:
838
+ """Make a causal mask for self-attention.
839
+
840
+ In case of 1d inputs (i.e., `[batch, len]`, the self-attention weights
841
+ will be `[batch, heads, len, len]` and this function will produce a
842
+ causal mask of shape `[batch, 1, len, len]`.
843
+
844
+ Note that a causal mask does not depend on the values of x; it only depends on
845
+ the shape. If x has padding elements, they will not be treated in a special
846
+ manner.
847
+
848
+ Args:
849
+ x: input array of shape `[batch, len]`
850
+ extra_batch_dims: number of batch dims to add singleton axes for, none by
851
+ default
852
+ dtype: mask return dtype
853
+
854
+ Returns:
855
+ A `[batch, 1, len, len]` shaped causal mask for 1d attention.
856
+ """
857
+ idxs = jnp.broadcast_to(jnp.arange(x.shape[-1], dtype=jnp.int32), x.shape)
858
+ return make_attention_mask(idxs, idxs, jnp.greater_equal, extra_batch_dims=extra_batch_dims, dtype=dtype)
859
+
860
+
861
+ def combine_masks(*masks: Optional[Array], dtype: DType = jnp.float32):
862
+ """Combine attention masks.
863
+
864
+ Args:
865
+ *masks: set of attention mask arguments to combine, some can be None.
866
+ dtype: final mask dtype
867
+
868
+ Returns:
869
+ Combined mask, reduced by logical and, returns None if no masks given.
870
+ """
871
+ masks = [m for m in masks if m is not None]
872
+ if not masks:
873
+ return None
874
+ assert all(
875
+ (x.ndim == masks[0].ndim for x in masks)
876
+ ), f"masks must have same rank: {tuple((x.ndim for x in masks))}"
877
+ mask, *other_masks = masks
878
+ for other_mask in other_masks:
879
+ mask = jnp.logical_and(mask, other_mask)
880
+ return mask.astype(dtype)
881
+
882
+
883
+ def combine_biases(*masks: Optional[Array]):
884
+ """Combine attention biases.
885
+
886
+ Args:
887
+ *masks: set of attention bias arguments to combine, some can be None.
888
+
889
+ Returns:
890
+ Combined mask, reduced by summation, returns None if no masks given.
891
+ """
892
+ masks = [m for m in masks if m is not None]
893
+ if not masks:
894
+ return None
895
+ assert all(
896
+ (x.ndim == masks[0].ndim for x in masks)
897
+ ), f"masks must have same rank: {tuple((x.ndim for x in masks))}"
898
+ mask, *other_masks = masks
899
+ for other_mask in other_masks:
900
+ mask = mask + other_mask
901
+ return mask
902
+
903
+
904
+ def make_decoder_mask(
905
+ decoder_target_tokens: Array,
906
+ dtype: DType,
907
+ decoder_causal_attention: Optional[Array] = None,
908
+ decoder_segment_ids: Optional[Array] = None,
909
+ ) -> Array:
910
+ """Compute the self-attention mask for a decoder.
911
+
912
+ Decoder mask is formed by combining a causal mask, a padding mask and an
913
+ optional packing mask. If decoder_causal_attention is passed, it makes the
914
+ masking non-causal for positions that have value of 1.
915
+
916
+ A prefix LM is applied to a dataset which has a notion of "inputs" and
917
+ "targets", e.g., a machine translation task. The inputs and targets are
918
+ concatenated to form a new target. `decoder_target_tokens` is the concatenated
919
+ decoder output tokens.
920
+
921
+ The "inputs" portion of the concatenated sequence can attend to other "inputs"
922
+ tokens even for those at a later time steps. In order to control this
923
+ behavior, `decoder_causal_attention` is necessary. This is a binary mask with
924
+ a value of 1 indicating that the position belonged to "inputs" portion of the
925
+ original dataset.
926
+
927
+ Example:
928
+
929
+ Suppose we have a dataset with two examples.
930
+
931
+ ds = [{"inputs": [6, 7], "targets": [8]},
932
+ {"inputs": [3, 4], "targets": [5]}]
933
+
934
+ After the data preprocessing with packing, the two examples are packed into
935
+ one example with the following three fields (some fields are skipped for
936
+ simplicity).
937
+
938
+ decoder_target_tokens = [[6, 7, 8, 3, 4, 5, 0]]
939
+ decoder_segment_ids = [[1, 1, 1, 2, 2, 2, 0]]
940
+ decoder_causal_attention = [[1, 1, 0, 1, 1, 0, 0]]
941
+
942
+ where each array has [batch, length] shape with batch size being 1. Then,
943
+ this function computes the following mask.
944
+
945
+ mask = [[[[1, 1, 0, 0, 0, 0, 0],
946
+ [1, 1, 0, 0, 0, 0, 0],
947
+ [1, 1, 1, 0, 0, 0, 0],
948
+ [0, 0, 0, 1, 1, 0, 0],
949
+ [0, 0, 0, 1, 1, 0, 0],
950
+ [0, 0, 0, 1, 1, 1, 0],
951
+ [0, 0, 0, 0, 0, 0, 0]]]]
952
+
953
+ mask[b, 1, :, :] represents the mask for the example `b` in the batch.
954
+ Because mask is for a self-attention layer, the mask's shape is a square of
955
+ shape [query length, key length].
956
+
957
+ mask[b, 1, i, j] = 1 means that the query token at position i can attend to
958
+ the key token at position j.
959
+
960
+ Args:
961
+ decoder_target_tokens: decoder output tokens. [batch, length]
962
+ dtype: dtype of the output mask.
963
+ decoder_causal_attention: a binary mask indicating which position should
964
+ only attend to earlier positions in the sequence. Others will attend
965
+ bidirectionally. [batch, length]
966
+ decoder_segment_ids: decoder segmentation info for packed examples. [batch,
967
+ length]
968
+
969
+ Returns:
970
+ the combined decoder mask.
971
+ """
972
+ masks = []
973
+ # The same mask is applied to all attention heads. So the head dimension is 1,
974
+ # i.e., the mask will be broadcast along the heads dim.
975
+ # [batch, 1, length, length]
976
+ causal_mask = make_causal_mask(decoder_target_tokens, dtype=dtype)
977
+
978
+ # Positions with value 1 in `decoder_causal_attneition` can attend
979
+ # bidirectionally.
980
+ if decoder_causal_attention is not None:
981
+ # [batch, 1, length, length]
982
+ inputs_mask = make_attention_mask(
983
+ decoder_causal_attention, decoder_causal_attention, jnp.logical_and, dtype=dtype
984
+ )
985
+ masks.append(jnp.logical_or(causal_mask, inputs_mask).astype(dtype))
986
+ else:
987
+ masks.append(causal_mask)
988
+
989
+ # Padding mask.
990
+ masks.append(make_attention_mask(decoder_target_tokens > 0, decoder_target_tokens > 0, dtype=dtype))
991
+
992
+ # Packing mask
993
+ if decoder_segment_ids is not None:
994
+ masks.append(make_attention_mask(decoder_segment_ids, decoder_segment_ids, jnp.equal, dtype=dtype))
995
+
996
+ return combine_masks(*masks, dtype=dtype)
997
+
998
+
999
+ def canonicalize_padding(padding: PaddingLike, rank: int) -> LaxPadding:
1000
+ """ "Canonicalizes conv padding to a jax.lax supported format."""
1001
+ if isinstance(padding, str):
1002
+ return padding
1003
+ if isinstance(padding, int):
1004
+ return [(padding, padding)] * rank
1005
+ if isinstance(padding, Sequence) and len(padding) == rank:
1006
+ new_pad = []
1007
+ for p in padding:
1008
+ if isinstance(p, int):
1009
+ new_pad.append((p, p))
1010
+ elif isinstance(p, tuple) and len(p) == 2:
1011
+ new_pad.append(p)
1012
+ else:
1013
+ break
1014
+ if len(new_pad) == rank:
1015
+ return new_pad
1016
+ raise ValueError(
1017
+ f"Invalid padding format: {padding}, should be str, int,"
1018
+ f" or a sequence of len {rank} where each element is an"
1019
+ f" int or pair of ints."
1020
+ )
1021
+
1022
+
1023
+ def _conv_dimension_numbers(input_shape):
1024
+ """Computes the dimension numbers based on the input shape."""
1025
+ ndim = len(input_shape)
1026
+ lhs_spec = (0, ndim - 1) + tuple(range(1, ndim - 1))
1027
+ rhs_spec = (ndim - 1, ndim - 2) + tuple(range(0, ndim - 2))
1028
+ out_spec = lhs_spec
1029
+ return lax.ConvDimensionNumbers(lhs_spec, rhs_spec, out_spec)
1030
+
1031
+
1032
+ class _Conv(nn.Module):
1033
+ """Convolution Module wrapping `lax.conv_general_dilated[_local]`.
1034
+
1035
+ Attributes:
1036
+ features: number of convolution filters.
1037
+ kernel_size: shape of the convolutional kernel. For 1D convolution,
1038
+ the kernel size can be passed as an integer. For all other cases, it must
1039
+ be a sequence of integers.
1040
+ strides: an integer or a sequence of `n` integers, representing the
1041
+ inter-window strides (default: 1).
1042
+ padding: either the string `'SAME'`, the string `'VALID'`, the string
1043
+ `'CIRCULAR'` (periodic boundary conditions), or a sequence of `n` `(low,
1044
+ high)` integer pairs that give the padding to apply before and after each
1045
+ spatial dimension. A single int is interpeted as applying the same padding
1046
+ in all dims and passign a single int in a sequence causes the same padding
1047
+ to be used on both sides. `'CAUSAL'` padding for a 1D convolution will
1048
+ left-pad the convolution axis, resulting in same-sized output.
1049
+ input_dilation: an integer or a sequence of `n` integers, giving the
1050
+ dilation factor to apply in each spatial dimension of `inputs`
1051
+ (default: 1). Convolution with input dilation `d` is equivalent to
1052
+ transposed convolution with stride `d`.
1053
+ kernel_dilation: an integer or a sequence of `n` integers, giving the
1054
+ dilation factor to apply in each spatial dimension of the convolution
1055
+ kernel (default: 1). Convolution with kernel dilation
1056
+ is also known as 'atrous convolution'.
1057
+ feature_group_count: integer, default 1. If specified divides the input
1058
+ features into groups.
1059
+ use_bias: whether to add a bias to the output (default: True).
1060
+ mask: Optional mask for the weights during masked convolution. The mask must
1061
+ be the same shape as the convolution weight matrix.
1062
+ dtype: the dtype of the computation (default: infer from input and params).
1063
+ params_dtype: the dtype passed to parameter initializers (default: float32).
1064
+ precision: numerical precision of the computation see `jax.lax.Precision`
1065
+ for details.
1066
+ kernel_init: initializer for the convolutional kernel.
1067
+ bias_init: initializer for the bias.
1068
+ """
1069
+
1070
+ features: int
1071
+ kernel_size: Sequence[int]
1072
+ strides: Union[None, int, Sequence[int]] = 1
1073
+ padding: PaddingLike = "SAME"
1074
+ input_dilation: Union[None, int, Sequence[int]] = 1
1075
+ kernel_dilation: Union[None, int, Sequence[int]] = 1
1076
+ feature_group_count: int = 1
1077
+ use_bias: bool = True
1078
+ mask: Optional[Array] = None
1079
+ dtype: Optional[DType] = None
1080
+ params_dtype: DType = jnp.float32
1081
+ precision: PrecisionLike = None
1082
+ kernel_init: Callable[[PRNGKey, Shape, DType], Array] = nn.initializers.lecun_normal()
1083
+ bias_init: Callable[[PRNGKey, Shape, DType], Array] = nn.initializers.zeros
1084
+ conv_general_dilated: ConvGeneralDilatedT = lax.conv_general_dilated
1085
+ kernel_axes: Tuple[str, ...] = ()
1086
+
1087
+ @property
1088
+ def shared_weights(self) -> bool: # type: ignore
1089
+ """Defines whether weights are shared or not between different pixels.
1090
+
1091
+ Returns:
1092
+ `True` to use shared weights in convolution (regular convolution).
1093
+ `False` to use different weights at different pixels, a.k.a.
1094
+ "locally connected layer", "unshared convolution", or "local convolution".
1095
+
1096
+ """
1097
+ ...
1098
+
1099
+ @nn.compact
1100
+ def __call__(self, inputs: Array) -> Array:
1101
+ """Applies a (potentially unshared) convolution to the inputs.
1102
+
1103
+ Args:
1104
+ inputs: input data with dimensions (*batch_dims, spatial_dims...,
1105
+ features). This is the channels-last convention, i.e. NHWC for a 2d
1106
+ convolution and NDHWC for a 3D convolution. Note: this is different from
1107
+ the input convention used by `lax.conv_general_dilated`, which puts the
1108
+ spatial dimensions last.
1109
+ Note: If the input has more than 1 batch dimension, all batch dimensions
1110
+ are flattened into a single dimension for the convolution and restored
1111
+ before returning. In some cases directly vmap'ing the layer may yield
1112
+ better performance than this default flattening approach. If the input
1113
+ lacks a batch dimension it will be added for the convolution and removed
1114
+ n return, an allowance made to enable writing single-example code.
1115
+
1116
+ Returns:
1117
+ The convolved data.
1118
+ """
1119
+
1120
+ if isinstance(self.kernel_size, int):
1121
+ raise TypeError(
1122
+ "Expected Conv kernel_size to be a"
1123
+ " tuple/list of integers (eg.: [3, 3]) but got"
1124
+ f" {self.kernel_size}."
1125
+ )
1126
+ else:
1127
+ kernel_size = tuple(self.kernel_size)
1128
+
1129
+ def maybe_broadcast(x: Optional[Union[int, Sequence[int]]]) -> Tuple[int, ...]:
1130
+ if x is None:
1131
+ # backward compatibility with using None as sentinel for
1132
+ # broadcast 1
1133
+ x = 1
1134
+ if isinstance(x, int):
1135
+ return (x,) * len(kernel_size)
1136
+ return tuple(x)
1137
+
1138
+ # Combine all input batch dimensions into a single leading batch axis.
1139
+ num_batch_dimensions = inputs.ndim - (len(kernel_size) + 1)
1140
+ if num_batch_dimensions != 1:
1141
+ input_batch_shape = inputs.shape[:num_batch_dimensions]
1142
+ total_batch_size = int(np.prod(input_batch_shape))
1143
+ flat_input_shape = (total_batch_size,) + inputs.shape[num_batch_dimensions:]
1144
+ inputs = jnp.reshape(inputs, flat_input_shape)
1145
+
1146
+ # self.strides or (1,) * (inputs.ndim - 2)
1147
+ strides = maybe_broadcast(self.strides)
1148
+ input_dilation = maybe_broadcast(self.input_dilation)
1149
+ kernel_dilation = maybe_broadcast(self.kernel_dilation)
1150
+
1151
+ padding_lax = canonicalize_padding(self.padding, len(kernel_size))
1152
+ if padding_lax == "CIRCULAR":
1153
+ kernel_size_dilated = [(k - 1) * d + 1 for k, d in zip(kernel_size, kernel_dilation)]
1154
+ zero_pad: List[Tuple[int, int]] = [(0, 0)]
1155
+ pads = zero_pad + [((k - 1) // 2, k // 2) for k in kernel_size_dilated] + [(0, 0)]
1156
+ inputs = jnp.pad(inputs, pads, mode="wrap")
1157
+ padding_lax = "VALID"
1158
+ elif padding_lax == "CAUSAL":
1159
+ if len(kernel_size) != 1:
1160
+ raise ValueError("Causal padding is only implemented for 1D convolutions.")
1161
+ left_pad = kernel_dilation[0] * (kernel_size[0] - 1)
1162
+ pads = [(0, 0), (left_pad, 0), (0, 0)]
1163
+ inputs = jnp.pad(inputs, pads)
1164
+ padding_lax = "VALID"
1165
+
1166
+ dimension_numbers = _conv_dimension_numbers(inputs.shape)
1167
+ in_features = jnp.shape(inputs)[-1]
1168
+
1169
+ if self.shared_weights:
1170
+ # One shared convolutional kernel for all pixels in the output.
1171
+ assert in_features % self.feature_group_count == 0
1172
+ kernel_shape = kernel_size + (in_features // self.feature_group_count, self.features)
1173
+
1174
+ else:
1175
+ if self.feature_group_count != 1:
1176
+ raise NotImplementedError(
1177
+ f"`lax.conv_general_dilated_local` does not support "
1178
+ f"`feature_group_count != 1`, got `{self.feature_group_count}`."
1179
+ )
1180
+
1181
+ # Need to know the spatial output shape of a standard convolution to
1182
+ # create the unshared convolution kernel.
1183
+ conv_output_shape = jax.eval_shape(
1184
+ lambda lhs, rhs: self.conv_general_dilated( # pylint: disable=g-long-lambda
1185
+ lhs=lhs,
1186
+ rhs=rhs,
1187
+ window_strides=strides,
1188
+ padding=padding_lax,
1189
+ dimension_numbers=dimension_numbers,
1190
+ lhs_dilation=input_dilation,
1191
+ rhs_dilation=kernel_dilation,
1192
+ ),
1193
+ inputs,
1194
+ jax.ShapedArray(kernel_size + (in_features, self.features), inputs.dtype),
1195
+ ).shape
1196
+
1197
+ # One (unshared) convolutional kernel per each pixel in the output.
1198
+ kernel_shape = conv_output_shape[1:-1] + (np.prod(kernel_size) * in_features, self.features)
1199
+
1200
+ if self.mask is not None and self.mask.shape != kernel_shape:
1201
+ raise ValueError(
1202
+ "Mask needs to have the same shape as weights. " f"Shapes are: {self.mask.shape}, {kernel_shape}"
1203
+ )
1204
+
1205
+ kernel = param_with_axes(
1206
+ "kernel",
1207
+ self.kernel_init,
1208
+ kernel_shape,
1209
+ self.params_dtype,
1210
+ axes=self.kernel_axes,
1211
+ )
1212
+
1213
+ if self.mask is not None:
1214
+ kernel *= self.mask
1215
+
1216
+ if self.use_bias:
1217
+ if self.shared_weights:
1218
+ # One bias weight per output channel, shared between pixels.
1219
+ bias_shape = (self.features,)
1220
+ else:
1221
+ # One bias weight per output entry, unshared betwen pixels.
1222
+ bias_shape = conv_output_shape[1:]
1223
+
1224
+ bias = param_with_axes(
1225
+ "bias",
1226
+ self.bias_init,
1227
+ bias_shape,
1228
+ self.params_dtype,
1229
+ axes=(self.kernel_axes[-1],),
1230
+ )
1231
+ else:
1232
+ bias = None
1233
+
1234
+ inputs, kernel, bias = promote_dtype(inputs, kernel, bias, dtype=self.dtype)
1235
+ if self.shared_weights:
1236
+ y = self.conv_general_dilated(
1237
+ inputs,
1238
+ kernel,
1239
+ strides,
1240
+ padding_lax,
1241
+ lhs_dilation=input_dilation,
1242
+ rhs_dilation=kernel_dilation,
1243
+ dimension_numbers=dimension_numbers,
1244
+ feature_group_count=self.feature_group_count,
1245
+ precision=self.precision,
1246
+ )
1247
+ else:
1248
+ y = lax.conv_general_dilated_local(
1249
+ lhs=inputs,
1250
+ rhs=kernel,
1251
+ window_strides=strides,
1252
+ padding=padding_lax,
1253
+ filter_shape=kernel_size,
1254
+ lhs_dilation=input_dilation,
1255
+ rhs_dilation=kernel_dilation,
1256
+ dimension_numbers=dimension_numbers,
1257
+ precision=self.precision,
1258
+ )
1259
+
1260
+ if self.use_bias:
1261
+ bias = bias.reshape((1,) * (y.ndim - bias.ndim) + bias.shape)
1262
+ y += bias
1263
+
1264
+ if num_batch_dimensions != 1:
1265
+ output_shape = input_batch_shape + y.shape[1:]
1266
+ y = jnp.reshape(y, output_shape)
1267
+ return y
1268
+
1269
+
1270
+ class Conv(_Conv):
1271
+ """Convolution Module wrapping `lax.conv_general_dilated`.
1272
+
1273
+ Attributes:
1274
+ features: number of convolution filters.
1275
+ kernel_size: shape of the convolutional kernel. For 1D convolution,
1276
+ the kernel size can be passed as an integer. For all other cases, it must
1277
+ be a sequence of integers.
1278
+ strides: an integer or a sequence of `n` integers, representing the
1279
+ inter-window strides (default: 1).
1280
+ padding: either the string `'SAME'`, the string `'VALID'`, the string
1281
+ `'CIRCULAR'` (periodic boundary conditions), or a sequence of `n` `(low,
1282
+ high)` integer pairs that give the padding to apply before and after each
1283
+ spatial dimension. A single int is interpeted as applying the same padding
1284
+ in all dims and passign a single int in a sequence causes the same padding
1285
+ to be used on both sides. `'CAUSAL'` padding for a 1D convolution will
1286
+ left-pad the convolution axis, resulting in same-sized output.
1287
+ input_dilation: an integer or a sequence of `n` integers, giving the
1288
+ dilation factor to apply in each spatial dimension of `inputs`
1289
+ (default: 1). Convolution with input dilation `d` is equivalent to
1290
+ transposed convolution with stride `d`.
1291
+ kernel_dilation: an integer or a sequence of `n` integers, giving the
1292
+ dilation factor to apply in each spatial dimension of the convolution
1293
+ kernel (default: 1). Convolution with kernel dilation
1294
+ is also known as 'atrous convolution'.
1295
+ feature_group_count: integer, default 1. If specified divides the input
1296
+ features into groups.
1297
+ use_bias: whether to add a bias to the output (default: True).
1298
+ mask: Optional mask for the weights during masked convolution. The mask must
1299
+ be the same shape as the convolution weight matrix.
1300
+ dtype: the dtype of the computation (default: infer from input and params).
1301
+ params_dtype: the dtype passed to parameter initializers (default: float32).
1302
+ precision: numerical precision of the computation see `jax.lax.Precision`
1303
+ for details.
1304
+ kernel_init: initializer for the convolutional kernel.
1305
+ bias_init: initializer for the bias.
1306
+ """
1307
+
1308
+ @property
1309
+ def shared_weights(self) -> bool:
1310
+ return True
whisper_jax/modeling_flax_whisper.py ADDED
@@ -0,0 +1,1686 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # coding=utf-8
2
+ # Copyright 2023 The OpenAI Authors and The HuggingFace Inc. team. All rights reserved.
3
+ #
4
+ # Licensed under the Apache License, Version 2.0 (the "License");
5
+ # you may not use this file except in compliance with the License.
6
+ # You may obtain a copy of the License at
7
+ #
8
+ # http://www.apache.org/licenses/LICENSE-2.0
9
+ #
10
+ # Unless required by applicable law or agreed to in writing, software
11
+ # distributed under the License is distributed on an "AS IS" BASIS,
12
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
+ # See the License for the specific language governing permissions and
14
+ # limitations under the License.
15
+ """ Flax whisper model."""
16
+
17
+ import random
18
+ from functools import partial
19
+ from typing import Optional, Tuple
20
+
21
+ import flax.linen as nn
22
+ import jax
23
+ import jax.numpy as jnp
24
+ from flax.core.frozen_dict import FrozenDict, freeze, unfreeze
25
+ from flax.linen import combine_masks, make_causal_mask
26
+ from flax.linen.attention import dot_product_attention_weights
27
+ from flax.traverse_util import flatten_dict, unflatten_dict
28
+ from jax import lax
29
+ from jax.random import PRNGKey
30
+ from transformers import WhisperConfig
31
+ from transformers.generation.flax_logits_process import (
32
+ FlaxLogitsProcessor,
33
+ FlaxLogitsProcessorList,
34
+ FlaxWhisperTimeStampLogitsProcessor,
35
+ )
36
+ from transformers.modeling_flax_outputs import (
37
+ FlaxBaseModelOutput,
38
+ FlaxBaseModelOutputWithPastAndCrossAttentions,
39
+ FlaxCausalLMOutputWithCrossAttentions,
40
+ FlaxSeq2SeqLMOutput,
41
+ FlaxSeq2SeqModelOutput,
42
+ )
43
+ from transformers.modeling_flax_utils import (
44
+ ACT2FN,
45
+ FlaxPreTrainedModel,
46
+ append_call_sample_docstring,
47
+ append_replace_return_docstrings,
48
+ overwrite_call_docstring,
49
+ )
50
+ from transformers.utils import (
51
+ add_start_docstrings,
52
+ add_start_docstrings_to_model_forward,
53
+ logging,
54
+ replace_return_docstrings,
55
+ )
56
+
57
+ from whisper_jax import layers
58
+ from whisper_jax.layers import with_sharding_constraint
59
+
60
+
61
+ logger = logging.get_logger(__name__)
62
+
63
+
64
+ _CHECKPOINT_FOR_DOC = "openai/whisper-tiny"
65
+ _CONFIG_FOR_DOC = "WhisperConfig"
66
+
67
+
68
+ WHISPER_START_DOCSTRING = r"""
69
+ This model inherits from [`FlaxPreTrainedModel`]. Check the superclass documentation for the generic methods the
70
+ library implements for all its models (such as downloading or saving, resizing the input embeddings, pruning heads
71
+ etc.) This model is also a Flax Linen
72
+ [flax.nn.Module](https://flax.readthedocs.io/en/latest/_autosummary/flax.nn.module.html) subclass. Use it as a
73
+ regular Flax Module and refer to the Flax documentation for all matter related to general usage and behavior.
74
+ Finally, this model supports inherent JAX features such as:
75
+ - [Just-In-Time (JIT) compilation](https://jax.readthedocs.io/en/latest/jax.html#just-in-time-compilation-jit)
76
+ - [Automatic Differentiation](https://jax.readthedocs.io/en/latest/jax.html#automatic-differentiation)
77
+ - [Vectorization](https://jax.readthedocs.io/en/latest/jax.html#vectorization-vmap)
78
+ - [Parallelization](https://jax.readthedocs.io/en/latest/jax.html#parallelization-pmap)
79
+
80
+ Parameters:
81
+ config ([`WhisperConfig`]): Model configuration class with all the parameters of the model.
82
+ Initializing with a config file does not load the weights associated with the model, only the
83
+ configuration. Check out the [`~FlaxPreTrainedModel.from_pretrained`] method to load the model weights.
84
+ dtype (`jax.numpy.dtype`, *optional*, defaults to `jax.numpy.float32`):
85
+ The data type of the computation. Can be one of `jax.numpy.float32`, `jax.numpy.float16` (on GPUs) and
86
+ `jax.numpy.bfloat16` (on TPUs). This can be used to enable mixed-precision training or half-precision
87
+ inference on GPUs or TPUs. If specified all the computation will be performed with the given `dtype`.
88
+ **Note that this only specifies the dtype of the computation and does not influence the dtype of model
89
+ parameters.** If you wish to change the dtype of the model parameters, see [`~FlaxPreTrainedModel.to_fp16`]
90
+ and [`~FlaxPreTrainedModel.to_bf16`].
91
+ """
92
+
93
+ WHISPER_INPUTS_DOCSTRING = r"""
94
+ Args:
95
+ input_features (`numpy.ndarray` of shape `(batch_size, feature_size, sequence_length)`):
96
+ Float values mel features extracted from the raw speech waveform. Raw speech waveform can be obtained by
97
+ loading a `.flac` or `.wav` audio file into an array of type `List[float]` or a `numpy.ndarray`, *e.g.* via
98
+ the soundfile library (`pip install soundfile`). To prepare the array into `input_features`, the
99
+ [`WhisperFeatureExtractor`] should be used for extracting the features, padding and conversion into a
100
+ tensor of type `numpy.ndarray`. See [`~WhisperFeatureExtractor.__call__`]
101
+ attention_mask (`numpy.ndarray` of shape `(batch_size, sequence_length)`, *optional*):
102
+ Whisper does not support masking of the `input_features`, this argument is preserved for compatibility, but
103
+ is not used. By default the silence in the input log mel spectrogram are ignored.
104
+ decoder_input_ids (`numpy.ndarray` of shape `(batch_size, target_sequence_length)`, *optional*):
105
+ Indices of decoder input sequence tokens in the vocabulary. Indices can be obtained using
106
+ [`WhisperTokenizer`]. See [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for details.
107
+ [What are decoder input IDs?](../glossary#decoder-input-ids) Whisper uses the `decoder_start_token_id` as
108
+ the starting token for `decoder_input_ids` generation.
109
+ decoder_attention_mask (`numpy.ndarray` of shape `(batch_size, target_sequence_length)`, *optional*):
110
+ Default behavior: generate a tensor that ignores pad tokens in `decoder_input_ids`. Causal mask will also
111
+ be used by default. If you want to change padding behavior, you should modify to your needs. See diagram 1
112
+ in [the paper](https://arxiv.org/abs/1910.13461) for more information on the default strategy.
113
+ position_ids (`numpy.ndarray` of shape `(batch_size, sequence_length)`, *optional*):
114
+ Whisper does not use `position_ids` in the encoder as `input_features` is always the same size and doesn't
115
+ use masking, but this argument is preserved for compatibility. By default the silence in the input log mel
116
+ spectrogram are ignored.
117
+ decoder_position_ids (`numpy.ndarray` of shape `(batch_size, sequence_length)`, *optional*):
118
+ Indices of positions of each decoder input sequence tokens in the position embeddings. Selected in the
119
+ range `[0, config.max_position_embeddings - 1]`.
120
+ output_attentions (`bool`, *optional*):
121
+ Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
122
+ tensors for more detail.
123
+ output_hidden_states (`bool`, *optional*):
124
+ Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
125
+ more detail.
126
+ return_dict (`bool`, *optional*):
127
+ Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
128
+ """
129
+
130
+ WHISPER_ENCODE_INPUTS_DOCSTRING = r"""
131
+ Args:
132
+ input_features (`numpy.ndarray` of shape `(batch_size, feature_size, sequence_length)`):
133
+ Float values mel features extracted from the raw speech waveform. Raw speech waveform can be obtained by
134
+ loading a `.flac` or `.wav` audio file into an array of type `List[float]` or a `numpy.ndarray`, *e.g.* via
135
+ the soundfile library (`pip install soundfile`). To prepare the array into `input_features`, the
136
+ [`WhisperFeatureExtractor`] should be used for extracting the mel features, padding and conversion into a
137
+ tensor of type `numpy.ndarray`. See [`~WhisperFeatureExtractor.__call__`].
138
+ attention_mask (`numpy.ndarray` of shape `(batch_size, sequence_length)`, *optional*):
139
+ Whisper does not support masking of the `input_features`, this argument is preserved for compatibility, but
140
+ is not used. By default the silence in the input log mel spectrogram are ignored.
141
+ output_attentions (`bool`, *optional*):
142
+ Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
143
+ tensors for more detail.
144
+ output_hidden_states (`bool`, *optional*):
145
+ Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
146
+ more detail.
147
+ return_dict (`bool`, *optional*):
148
+ Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
149
+ """
150
+
151
+ WHISPER_DECODE_INPUTS_DOCSTRING = r"""
152
+ Args:
153
+ decoder_input_ids (`numpy.ndarray` of shape `(batch_size, target_sequence_length)`):
154
+ Indices of decoder input sequence tokens in the vocabulary. Indices can be obtained using
155
+ [`WhisperTokenizer`]. See [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for details.
156
+ [What are decoder input IDs?](../glossary#decoder-input-ids)
157
+ encoder_outputs (`tuple(tuple(numpy.ndarray)`):
158
+ Tuple consists of (`last_hidden_state`, *optional*: `hidden_states`, *optional*: `attentions`)
159
+ `last_hidden_state` of shape `(batch_size, sequence_length, hidden_size)`, *optional*) is a sequence of
160
+ hidden-states at the output of the last layer of the encoder. Used in the cross-attention of the decoder.
161
+ encoder_attention_mask (`numpy.ndarray` of shape `(batch_size, sequence_length)`, *optional*):
162
+ Whisper does not support masking of the `input_features`, this argument is preserved for compatibility,
163
+ but it is not used. By default the silence in the input log mel spectrogram are ignored.
164
+ decoder_attention_mask (`numpy.ndarray` of shape `(batch_size, target_sequence_length)`, *optional*):
165
+ Default behavior: generate a tensor that ignores pad tokens in `decoder_input_ids`. Causal mask will also
166
+ be used by default. If you want to change padding behavior, you should modify to your needs. See diagram 1
167
+ in [the paper](https://arxiv.org/abs/1910.13461) for more information on the default strategy.
168
+ decoder_position_ids (`numpy.ndarray` of shape `(batch_size, sequence_length)`, *optional*):
169
+ Indices of positions of each decoder input sequence tokens in the position embeddings. Selected in the
170
+ range `[0, config.max_position_embeddings - 1]`.
171
+ past_key_values (`Dict[str, numpy.ndarray]`, *optional*, returned by `init_cache` or when passing previous `past_key_values`):
172
+ Dictionary of pre-computed hidden-states (key and values in the attention blocks) that can be used for fast
173
+ auto-regressive decoding. Pre-computed key and value hidden-states are of shape *[batch_size, max_length]*.
174
+ output_attentions (`bool`, *optional*):
175
+ Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
176
+ tensors for more detail.
177
+ output_hidden_states (`bool`, *optional*):
178
+ Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
179
+ more detail.
180
+ return_dict (`bool`, *optional*):
181
+ Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
182
+ """
183
+
184
+
185
+ class FlaxStaticForceTokensLogitsProcessor(FlaxLogitsProcessor):
186
+ r"""
187
+ [`FlaxLogitsProcessor`] that takes a list of pairs of integers which indicates a mapping from generation indices to
188
+ token indices that will be forced before sampling. The processor will set their log probs to 0 and all other tokens
189
+ to `-inf` so that they are sampled at their corresponding index. This is a static version of the `transformers` logit
190
+ processor [`FlaxForceTokensLogitsProcessor`] that is compatible with sharded forced tokens.
191
+
192
+ Args:
193
+ force_token_map (`list`):
194
+ Map giving token ids and indices where they will be forced to be sampled.
195
+ """
196
+
197
+ def __init__(self, force_token_map):
198
+ # The generic `transformers` logit processor builds `force_token_array` as a dictionary - this is not a valid
199
+ # JAX type, and so we switch to using a JAX array instead
200
+ force_token_map = jnp.array(force_token_map)
201
+ # Converts the array of format [[index, token]] containing the tokens to be forced to an array, where the
202
+ # index of the array corresponds to the index of the token to be forced. For XLA compatibility,
203
+ # indexes without forced tokens will have a negative value. Note that the last token we ever need to force in
204
+ # Whisper is at position 3, so we only construct an array up to this index. The native version constructs a tensor
205
+ # dynamically according to the length of the `force_token_map`. Array shapes need to be concrete for XLA compatibility,
206
+ # so this is not permitted here.
207
+ force_token_array = jnp.ones(3, dtype=jnp.int32) * -1
208
+ for index, token in force_token_map:
209
+ force_token_array = force_token_array.at[index].set(token)
210
+ self.force_token_array = jnp.int32(force_token_array)
211
+
212
+ def __call__(self, input_ids: jnp.ndarray, scores: jnp.ndarray, cur_len: int) -> jnp.ndarray:
213
+ def _force_token(generation_idx):
214
+ batch_size = scores.shape[0]
215
+ current_token = self.force_token_array[generation_idx]
216
+
217
+ new_scores = jnp.ones_like(scores, dtype=scores.dtype) * -float("inf")
218
+ updates = jnp.zeros((batch_size, 1), dtype=scores.dtype)
219
+ new_scores = lax.dynamic_update_slice(new_scores, updates, (0, current_token))
220
+ return new_scores
221
+
222
+ scores = lax.cond(
223
+ cur_len >= self.force_token_array.shape[0],
224
+ # If the current length is geq than the length of force_token_array, the processor does nothing.
225
+ lambda: scores,
226
+ # Otherwise, it may force a certain token.
227
+ lambda: lax.cond(
228
+ self.force_token_array[cur_len] >= 0,
229
+ # Only valid (positive) tokens are forced
230
+ lambda: _force_token(cur_len),
231
+ # Otherwise, the processor does nothing.
232
+ lambda: scores,
233
+ ),
234
+ )
235
+ return scores
236
+
237
+
238
+ class FlaxWhisperAttention(nn.Module):
239
+ config: WhisperConfig
240
+ embed_dim: int
241
+ num_heads: int
242
+ dropout: float = 0.0
243
+ causal: bool = False
244
+ bias: bool = True
245
+ dtype: jnp.dtype = jnp.float32
246
+ params_dtype: jnp.dtype = jnp.float32
247
+
248
+ def setup(self) -> None:
249
+ self.head_dim = self.embed_dim // self.num_heads
250
+ if self.head_dim * self.num_heads != self.embed_dim:
251
+ raise ValueError(
252
+ f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim}"
253
+ f" and `num_heads`: {self.num_heads})."
254
+ )
255
+
256
+ dense = partial(
257
+ layers.DenseGeneral,
258
+ self.embed_dim,
259
+ axis=-1,
260
+ dtype=self.dtype,
261
+ params_dtype=self.params_dtype,
262
+ kernel_axes=("embed", "joined_kv"),
263
+ )
264
+
265
+ self.q_proj = dense(use_bias=self.bias)
266
+ self.k_proj = dense(use_bias=False)
267
+ self.v_proj = dense(use_bias=self.bias)
268
+
269
+ self.out_proj = layers.DenseGeneral(
270
+ self.embed_dim,
271
+ axis=-1,
272
+ dtype=self.dtype,
273
+ params_dtype=self.params_dtype,
274
+ kernel_axes=("joined_kv", "embed"),
275
+ use_bias=self.bias,
276
+ )
277
+
278
+ if self.causal:
279
+ self.causal_mask = make_causal_mask(
280
+ jnp.ones((1, self.config.max_target_positions), dtype="bool"), dtype="bool"
281
+ )
282
+
283
+ def __call__(
284
+ self,
285
+ hidden_states: jnp.ndarray,
286
+ key_value_states: Optional[jnp.ndarray] = None,
287
+ attention_mask: Optional[jnp.ndarray] = None,
288
+ init_cache: bool = False,
289
+ deterministic: bool = True,
290
+ ) -> Tuple[jnp.ndarray]:
291
+ is_cross_attention = key_value_states is not None
292
+ batch_size = hidden_states.shape[0]
293
+
294
+ query_states = self.q_proj(hidden_states)
295
+
296
+ if is_cross_attention:
297
+ key_states = self.k_proj(key_value_states)
298
+ value_states = self.v_proj(key_value_states)
299
+ else:
300
+ key_states = self.k_proj(hidden_states)
301
+ value_states = self.v_proj(hidden_states)
302
+
303
+ query_states = self._split_heads(query_states)
304
+ key_states = self._split_heads(key_states)
305
+ value_states = self._split_heads(value_states)
306
+
307
+ query_states = with_sharding_constraint(query_states, ("batch", "length", "heads", "kv"))
308
+ key_states = with_sharding_constraint(key_states, ("batch", "length", "heads", "kv"))
309
+ value_states = with_sharding_constraint(value_states, ("batch", "length", "heads", "kv"))
310
+
311
+ if self.causal:
312
+ query_length, key_length = query_states.shape[1], key_states.shape[1]
313
+ if self.has_variable("cache", "cached_key"):
314
+ mask_shift = self.variables["cache"]["cache_index"]
315
+ # max_length of cached_key is last dim
316
+ max_decoder_length = self.variables["cache"]["cached_key"].shape[-1]
317
+ causal_mask = lax.dynamic_slice(
318
+ self.causal_mask,
319
+ (0, 0, mask_shift, 0),
320
+ (1, 1, query_length, max_decoder_length),
321
+ )
322
+ else:
323
+ causal_mask = self.causal_mask[:, :, :query_length, :key_length]
324
+ causal_mask = jnp.broadcast_to(causal_mask, (batch_size,) + causal_mask.shape[1:])
325
+
326
+ # combine masks if needed
327
+ if attention_mask is not None and self.causal:
328
+ attention_mask = jnp.broadcast_to(jnp.expand_dims(attention_mask, axis=(-3, -2)), causal_mask.shape)
329
+ attention_mask = combine_masks(attention_mask, causal_mask)
330
+ elif self.causal:
331
+ attention_mask = causal_mask
332
+ elif attention_mask is not None:
333
+ attention_mask = jnp.expand_dims(attention_mask, axis=(-3, -2))
334
+
335
+ # During fast autoregressive decoding, we feed one position at a time,
336
+ # and cache the keys and values step by step.
337
+
338
+ if self.causal and (self.has_variable("cache", "cached_key") or init_cache):
339
+ key_states, value_states, attention_mask = self._concatenate_to_cache(
340
+ key_states, value_states, query_states, attention_mask
341
+ )
342
+
343
+ # Convert the boolean attention mask to an attention bias.
344
+ if attention_mask is not None:
345
+ # attention mask in the form of attention bias
346
+ attention_bias = lax.select(
347
+ attention_mask > 0,
348
+ jnp.full(attention_mask.shape, 0.0).astype(self.dtype),
349
+ jnp.full(attention_mask.shape, jnp.finfo(self.dtype).min).astype(self.dtype),
350
+ )
351
+ else:
352
+ attention_bias = None
353
+
354
+ dropout_rng = None
355
+ if not deterministic and self.dropout > 0.0:
356
+ dropout_rng = self.make_rng("dropout")
357
+
358
+ attn_weights = dot_product_attention_weights(
359
+ query_states,
360
+ key_states,
361
+ bias=attention_bias,
362
+ dropout_rng=dropout_rng,
363
+ dropout_rate=self.dropout,
364
+ broadcast_dropout=True,
365
+ deterministic=deterministic,
366
+ dtype=self.dtype,
367
+ precision=None,
368
+ )
369
+
370
+ attn_output = jnp.einsum("...hqk,...khd->...qhd", attn_weights, value_states)
371
+ attn_output = self._merge_heads(attn_output)
372
+ attn_output = self.out_proj(attn_output)
373
+
374
+ return attn_output, attn_weights
375
+
376
+ def _split_heads(self, hidden_state) -> jnp.ndarray:
377
+ return hidden_state.reshape(hidden_state.shape[:2] + (self.num_heads, self.head_dim))
378
+
379
+ def _merge_heads(self, hidden_state) -> jnp.ndarray:
380
+ return hidden_state.reshape(hidden_state.shape[:2] + (self.embed_dim,))
381
+
382
+ @nn.compact
383
+ def _concatenate_to_cache(self, key, value, query, attention_mask):
384
+ # The following code is largely copied from: https://github.com/google-research/t5x/blob/63d9addf628c6d8c547a407a32095fcb527bb20b/t5x/examples/scalable_t5/layers.py#L280-L284
385
+ is_initialized = self.has_variable("cache", "cached_key")
386
+
387
+ # The key and value have dimension [batch_size, seq_length, num_heads, head_dim],
388
+ # but we cache them as [batch_size, num_heads, head_dim, seq_length] as a TPU
389
+ # fusion optimization. This also enables the "scatter via one-hot
390
+ # broadcast" trick, which means we do a one-hot broadcast instead of a
391
+ # scatter/gather operations, resulting in a 3-4x speedup in practice.
392
+ def swap_dims(x):
393
+ return x[:-3] + tuple(x[i] for i in [-2, -1, -3])
394
+
395
+ cached_key = self.variable("cache", "cached_key", jnp.zeros, swap_dims(key.shape), key.dtype)
396
+ cached_value = self.variable("cache", "cached_value", jnp.zeros, swap_dims(value.shape), value.dtype)
397
+ cache_index = self.variable("cache", "cache_index", lambda: jnp.array(0, dtype=jnp.int32))
398
+
399
+ if is_initialized:
400
+ batch_size, num_heads, head_dim, seq_length = cached_key.value.shape
401
+ # During fast autoregressive decoding, we feed one position at a time,
402
+ # and cache the keys and values step by step.
403
+ # Sanity shape check of cached key against input query.
404
+ num_updated_cache_vectors = query.shape[1]
405
+ expected_shape = (batch_size, 1, num_heads, head_dim)
406
+ if num_updated_cache_vectors == 1 and expected_shape != query.shape:
407
+ raise ValueError(
408
+ f"Autoregressive cache shape error, expected query shape {expected_shape} instead got {query.shape}"
409
+ )
410
+
411
+ # Create a OHE of the current index. NOTE: the index is increased below.
412
+ cur_index = cache_index.value
413
+
414
+ # In order to update the key, value caches with the current key and
415
+ # value, we move the seq_length axis to the back, similar to what we did for
416
+ # the cached ones above.
417
+ # Note these are currently the key and value of a single position, since
418
+ # we feed one position at a time.
419
+ one_token_key = jnp.moveaxis(key, -3, -1)
420
+ one_token_value = jnp.moveaxis(value, -3, -1)
421
+
422
+ # Update key, value caches with our new 1d spatial slices.
423
+ # We implement an efficient scatter into the cache via one-hot
424
+ # broadcast and addition.
425
+ if num_updated_cache_vectors > 1:
426
+ indices = jnp.eye(num_updated_cache_vectors, seq_length)[None, None]
427
+ key = cached_key.value + jnp.matmul(one_token_key, indices)
428
+ value = cached_value.value + jnp.matmul(one_token_value, indices)
429
+ else:
430
+ one_hot_indices = jax.nn.one_hot(cur_index, seq_length, dtype=key.dtype)
431
+ key = cached_key.value + one_token_key * one_hot_indices
432
+ value = cached_value.value + one_token_value * one_hot_indices
433
+
434
+ cached_key.value = key
435
+ cached_value.value = value
436
+ cache_index.value = cache_index.value + num_updated_cache_vectors
437
+
438
+ # Move the keys and values back to their original shapes.
439
+ key = jnp.moveaxis(key, -1, -3)
440
+ value = jnp.moveaxis(value, -1, -3)
441
+
442
+ # causal mask for cached decoder self-attention: our single query position should only
443
+ # attend to those key positions that have already been generated and cached, not the
444
+ # remaining zero elements.
445
+ pad_mask = jnp.broadcast_to(
446
+ jnp.arange(seq_length) < cur_index + num_updated_cache_vectors,
447
+ (batch_size,) + (1, num_updated_cache_vectors, seq_length),
448
+ )
449
+ attention_mask = combine_masks(pad_mask, attention_mask)
450
+
451
+ return key, value, attention_mask
452
+
453
+
454
+ # Copied from transformers.models.mbart.modeling_flax_mbart.FlaxMBartEncoderLayer with MBart->Whisper
455
+ class FlaxWhisperEncoderLayer(nn.Module):
456
+ config: WhisperConfig
457
+ dtype: jnp.dtype = jnp.float32
458
+ params_dtype: jnp.dtype = jnp.float32
459
+
460
+ def setup(self) -> None:
461
+ self.embed_dim = self.config.d_model
462
+ self.self_attn = FlaxWhisperAttention(
463
+ config=self.config,
464
+ embed_dim=self.embed_dim,
465
+ num_heads=self.config.encoder_attention_heads,
466
+ dropout=self.config.attention_dropout,
467
+ dtype=self.dtype,
468
+ params_dtype=self.params_dtype,
469
+ )
470
+ self.self_attn_layer_norm = layers.LayerNorm(dtype=self.dtype, epsilon=1e-05, params_dtype=self.params_dtype)
471
+ self.dropout_layer = nn.Dropout(rate=self.config.dropout)
472
+ self.activation_fn = ACT2FN[self.config.activation_function]
473
+ self.activation_dropout_layer = nn.Dropout(rate=self.config.activation_dropout)
474
+ self.fc1 = layers.DenseGeneral(
475
+ self.config.encoder_ffn_dim,
476
+ dtype=self.dtype,
477
+ params_dtype=self.params_dtype,
478
+ kernel_axes=("embed", "mlp"),
479
+ )
480
+ self.fc2 = layers.DenseGeneral(
481
+ self.embed_dim,
482
+ dtype=self.dtype,
483
+ params_dtype=self.params_dtype,
484
+ kernel_axes=("mlp", "embed"),
485
+ )
486
+ self.final_layer_norm = layers.LayerNorm(dtype=self.dtype, epsilon=1e-05, params_dtype=self.params_dtype)
487
+
488
+ def __call__(
489
+ self,
490
+ hidden_states: jnp.ndarray,
491
+ attention_mask: jnp.ndarray,
492
+ output_attentions: bool = True,
493
+ deterministic: bool = True,
494
+ ) -> Tuple[jnp.ndarray]:
495
+ hidden_states = with_sharding_constraint(hidden_states, ("batch", "length", "embed"))
496
+
497
+ residual = hidden_states
498
+
499
+ layernorm_output = self.self_attn_layer_norm(hidden_states)
500
+ layernorm_output = with_sharding_constraint(layernorm_output, ("batch", "length", "embed"))
501
+
502
+ attn_output, attn_weights = self.self_attn(hidden_states=layernorm_output, attention_mask=attention_mask)
503
+ attn_output = self.dropout_layer(attn_output, deterministic=deterministic)
504
+ attn_output = residual + attn_output
505
+ attn_output = with_sharding_constraint(attn_output, ("batch", "length", "embed"))
506
+
507
+ residual = attn_output
508
+
509
+ post_layer_norm = self.final_layer_norm(attn_output)
510
+ post_layer_norm = with_sharding_constraint(post_layer_norm, ("batch", "length", "embed"))
511
+
512
+ fc1_output = self.activation_fn(self.fc1(post_layer_norm))
513
+ fc1_output = self.activation_dropout_layer(fc1_output, deterministic=deterministic)
514
+ fc1_output = with_sharding_constraint(fc1_output, ("batch", "length", "mlp"))
515
+
516
+ hidden_states = self.fc2(fc1_output)
517
+ hidden_states = self.dropout_layer(hidden_states, deterministic=deterministic)
518
+ hidden_states = residual + hidden_states
519
+ hidden_states = with_sharding_constraint(hidden_states, ("batch", "length", "embed"))
520
+
521
+ outputs = (hidden_states,)
522
+
523
+ if output_attentions:
524
+ outputs += (attn_weights,)
525
+
526
+ return outputs
527
+
528
+
529
+ # Copied from transformers.models.mbart.modeling_flax_mbart.FlaxMBartEncoderLayerCollection with MBart->Whisper
530
+ class FlaxWhisperEncoderLayerCollection(nn.Module):
531
+ config: WhisperConfig
532
+ dtype: jnp.dtype = jnp.float32 # the dtype of the computation
533
+ params_dtype: jnp.dtype = jnp.float32
534
+
535
+ def setup(self):
536
+ self.layers = [
537
+ FlaxWhisperEncoderLayer(self.config, name=str(i), dtype=self.dtype, params_dtype=self.params_dtype)
538
+ for i in range(self.config.encoder_layers)
539
+ ]
540
+ self.layerdrop = self.config.encoder_layerdrop
541
+
542
+ def __call__(
543
+ self,
544
+ hidden_states,
545
+ attention_mask,
546
+ deterministic: bool = True,
547
+ output_attentions: bool = False,
548
+ output_hidden_states: bool = False,
549
+ return_dict: bool = True,
550
+ ):
551
+ all_attentions = () if output_attentions else None
552
+ all_hidden_states = () if output_hidden_states else None
553
+
554
+ for encoder_layer in self.layers:
555
+ if output_hidden_states:
556
+ all_hidden_states = all_hidden_states + (hidden_states,)
557
+ # add LayerDrop (see https://arxiv.org/abs/1909.11556 for description)
558
+ dropout_probability = random.uniform(0, 1)
559
+ if not deterministic and (dropout_probability < self.layerdrop): # skip the layer
560
+ layer_outputs = (None, None)
561
+ else:
562
+ layer_outputs = encoder_layer(
563
+ hidden_states,
564
+ attention_mask,
565
+ output_attentions,
566
+ deterministic,
567
+ )
568
+ hidden_states = layer_outputs[0]
569
+ if output_attentions:
570
+ all_attentions = all_attentions + (layer_outputs[1],)
571
+
572
+ if output_hidden_states:
573
+ all_hidden_states += (hidden_states,)
574
+
575
+ outputs = (hidden_states, all_hidden_states, all_attentions)
576
+
577
+ if not return_dict:
578
+ return tuple(v for v in outputs if v is not None)
579
+
580
+ return FlaxBaseModelOutput(
581
+ last_hidden_state=hidden_states, hidden_states=all_hidden_states, attentions=all_attentions
582
+ )
583
+
584
+
585
+ # Copied from transformers.models.mbart.modeling_flax_mbart.FlaxMBartDecoderLayer with MBart->Whisper
586
+ class FlaxWhisperDecoderLayer(nn.Module):
587
+ config: WhisperConfig
588
+ dtype: jnp.dtype = jnp.float32
589
+ params_dtype: jnp.dtype = jnp.float32
590
+
591
+ def setup(self) -> None:
592
+ self.embed_dim = self.config.d_model
593
+ self.self_attn = FlaxWhisperAttention(
594
+ config=self.config,
595
+ embed_dim=self.embed_dim,
596
+ num_heads=self.config.decoder_attention_heads,
597
+ dropout=self.config.attention_dropout,
598
+ causal=True,
599
+ dtype=self.dtype,
600
+ params_dtype=self.params_dtype,
601
+ )
602
+ self.dropout_layer = nn.Dropout(rate=self.config.dropout)
603
+ self.activation_fn = ACT2FN[self.config.activation_function]
604
+ self.activation_dropout_layer = nn.Dropout(rate=self.config.activation_dropout)
605
+
606
+ self.self_attn_layer_norm = layers.LayerNorm(dtype=self.dtype, epsilon=1e-05, params_dtype=self.params_dtype)
607
+ self.encoder_attn = FlaxWhisperAttention(
608
+ config=self.config,
609
+ embed_dim=self.embed_dim,
610
+ num_heads=self.config.decoder_attention_heads,
611
+ dropout=self.config.attention_dropout,
612
+ dtype=self.dtype,
613
+ params_dtype=self.params_dtype,
614
+ )
615
+ self.encoder_attn_layer_norm = layers.LayerNorm(
616
+ dtype=self.dtype, epsilon=1e-05, params_dtype=self.params_dtype
617
+ )
618
+ self.fc1 = layers.DenseGeneral(
619
+ self.config.decoder_ffn_dim,
620
+ dtype=self.dtype,
621
+ params_dtype=self.params_dtype,
622
+ kernel_axes=("embed", "mlp"),
623
+ )
624
+ self.fc2 = layers.DenseGeneral(
625
+ self.embed_dim,
626
+ dtype=self.dtype,
627
+ params_dtype=self.params_dtype,
628
+ kernel_axes=("mlp", "embed"),
629
+ )
630
+ self.final_layer_norm = layers.LayerNorm(dtype=self.dtype, epsilon=1e-05, params_dtype=self.params_dtype)
631
+
632
+ def __call__(
633
+ self,
634
+ hidden_states: jnp.ndarray,
635
+ attention_mask: jnp.ndarray,
636
+ encoder_hidden_states: Optional[jnp.ndarray] = None,
637
+ encoder_attention_mask: Optional[jnp.ndarray] = None,
638
+ init_cache: bool = False,
639
+ output_attentions: bool = True,
640
+ deterministic: bool = True,
641
+ ) -> Tuple[jnp.ndarray]:
642
+ hidden_states = with_sharding_constraint(hidden_states, ("batch", "length", "embed"))
643
+
644
+ residual = hidden_states
645
+
646
+ layer_norm_output = self.self_attn_layer_norm(hidden_states)
647
+ layer_norm_output = with_sharding_constraint(layer_norm_output, ("batch", "length", "embed"))
648
+
649
+ # Self Attention
650
+ self_attn_output, self_attn_weights = self.self_attn(
651
+ hidden_states=layer_norm_output, attention_mask=attention_mask, init_cache=init_cache
652
+ )
653
+ self_attn_output = self.dropout_layer(self_attn_output, deterministic=deterministic)
654
+ self_attn_output = residual + self_attn_output
655
+ self_attn_output = with_sharding_constraint(self_attn_output, ("batch", "length", "embed"))
656
+
657
+ # Cross-Attention Block
658
+ cross_attn_weights = None
659
+ if encoder_hidden_states is not None:
660
+ residual = self_attn_output
661
+
662
+ encoder_layer_norm_output = self.encoder_attn_layer_norm(self_attn_output)
663
+ encoder_layer_norm_output = with_sharding_constraint(
664
+ encoder_layer_norm_output, ("batch", "length", "embed")
665
+ )
666
+
667
+ cross_attn_output, cross_attn_weights = self.encoder_attn(
668
+ hidden_states=encoder_layer_norm_output,
669
+ key_value_states=encoder_hidden_states,
670
+ attention_mask=encoder_attention_mask,
671
+ )
672
+ cross_attn_output = self.dropout_layer(cross_attn_output, deterministic=deterministic)
673
+ cross_attn_output = residual + cross_attn_output
674
+ cross_attn_output = with_sharding_constraint(cross_attn_output, ("batch", "length", "embed"))
675
+
676
+ # Fully Connected
677
+ residual = cross_attn_output
678
+
679
+ post_layer_norm = self.final_layer_norm(cross_attn_output)
680
+ post_layer_norm = with_sharding_constraint(post_layer_norm, ("batch", "length", "embed"))
681
+
682
+ fc1_output = self.activation_fn(self.fc1(post_layer_norm))
683
+ fc1_output = self.activation_dropout_layer(fc1_output, deterministic=deterministic)
684
+ fc1_output = with_sharding_constraint(fc1_output, ("batch", "length", "mlp"))
685
+
686
+ hidden_states = self.fc2(fc1_output)
687
+ hidden_states = self.dropout_layer(hidden_states, deterministic=deterministic)
688
+ hidden_states = residual + hidden_states
689
+ hidden_states = with_sharding_constraint(hidden_states, ("batch", "length", "embed"))
690
+
691
+ outputs = (hidden_states,)
692
+
693
+ if output_attentions:
694
+ outputs += (self_attn_weights, cross_attn_weights)
695
+
696
+ return outputs
697
+
698
+
699
+ # Copied from transformers.models.mbart.modeling_flax_mbart.FlaxMBartDecoderLayerCollection with MBart->Whisper
700
+ class FlaxWhisperDecoderLayerCollection(nn.Module):
701
+ config: WhisperConfig
702
+ dtype: jnp.dtype = jnp.float32 # the dtype of the computation
703
+ params_dtype: jnp.dtype = jnp.float32
704
+
705
+ def setup(self):
706
+ self.layers = [
707
+ FlaxWhisperDecoderLayer(self.config, name=str(i), dtype=self.dtype, params_dtype=self.params_dtype)
708
+ for i in range(self.config.decoder_layers)
709
+ ]
710
+ self.layerdrop = self.config.decoder_layerdrop
711
+
712
+ def __call__(
713
+ self,
714
+ hidden_states,
715
+ attention_mask,
716
+ encoder_hidden_states: Optional[jnp.ndarray] = None,
717
+ encoder_attention_mask: Optional[jnp.ndarray] = None,
718
+ deterministic: bool = True,
719
+ init_cache: bool = False,
720
+ output_attentions: bool = False,
721
+ output_hidden_states: bool = False,
722
+ return_dict: bool = True,
723
+ ):
724
+ # decoder layers
725
+ all_hidden_states = () if output_hidden_states else None
726
+ all_self_attns = () if output_attentions else None
727
+ all_cross_attentions = () if (output_attentions and encoder_hidden_states is not None) else None
728
+
729
+ for decoder_layer in self.layers:
730
+ if output_hidden_states:
731
+ all_hidden_states += (hidden_states,)
732
+ # add LayerDrop (see https://arxiv.org/abs/1909.11556 for description)
733
+ dropout_probability = random.uniform(0, 1)
734
+ if not deterministic and (dropout_probability < self.layerdrop):
735
+ layer_outputs = (None, None, None)
736
+ else:
737
+ layer_outputs = decoder_layer(
738
+ hidden_states,
739
+ attention_mask=attention_mask,
740
+ encoder_hidden_states=encoder_hidden_states,
741
+ encoder_attention_mask=encoder_attention_mask,
742
+ init_cache=init_cache,
743
+ output_attentions=output_attentions,
744
+ deterministic=deterministic,
745
+ )
746
+
747
+ hidden_states = layer_outputs[0]
748
+ if output_attentions:
749
+ all_self_attns += (layer_outputs[1],)
750
+
751
+ if encoder_hidden_states is not None:
752
+ all_cross_attentions += (layer_outputs[2],)
753
+
754
+ # add hidden states from the last decoder layer
755
+ if output_hidden_states:
756
+ all_hidden_states += (hidden_states,)
757
+
758
+ outputs = [hidden_states, all_hidden_states, all_self_attns, all_cross_attentions]
759
+
760
+ if not return_dict:
761
+ return tuple(v for v in outputs if v is not None)
762
+
763
+ return FlaxBaseModelOutputWithPastAndCrossAttentions(
764
+ last_hidden_state=hidden_states,
765
+ hidden_states=all_hidden_states,
766
+ attentions=all_self_attns,
767
+ cross_attentions=all_cross_attentions,
768
+ )
769
+
770
+
771
+ class FlaxWhisperEncoder(nn.Module):
772
+ config: WhisperConfig
773
+ dtype: jnp.dtype = jnp.float32
774
+ params_dtype: jnp.dtype = jnp.float32
775
+
776
+ def setup(self) -> None:
777
+ self.conv1 = layers.Conv(
778
+ self.config.d_model,
779
+ kernel_size=(3,),
780
+ padding=1,
781
+ dtype=self.dtype,
782
+ params_dtype=self.params_dtype,
783
+ kernel_axes=("channels", "num_mel", "embed"),
784
+ )
785
+ self.conv2 = layers.Conv(
786
+ self.config.d_model,
787
+ kernel_size=(3,),
788
+ strides=2,
789
+ padding=1,
790
+ dtype=self.dtype,
791
+ params_dtype=self.params_dtype,
792
+ kernel_axes=("channels", "embed", "num_mel"),
793
+ )
794
+
795
+ self.dropout_layer = nn.Dropout(rate=self.config.dropout)
796
+
797
+ self.layers = FlaxWhisperEncoderLayerCollection(
798
+ self.config,
799
+ dtype=self.dtype,
800
+ params_dtype=self.params_dtype,
801
+ )
802
+ self.embed_positions = layers.Embed(
803
+ self.config.max_source_positions, self.config.d_model, dtype=self.dtype, params_dtype=self.params_dtype
804
+ )
805
+
806
+ self.layer_norm = layers.LayerNorm(dtype=self.dtype, epsilon=1e-05, params_dtype=self.params_dtype)
807
+
808
+ def __call__(
809
+ self,
810
+ input_features: jnp.ndarray,
811
+ output_attentions: bool = False,
812
+ output_hidden_states: bool = False,
813
+ return_dict: bool = True,
814
+ deterministic: bool = True,
815
+ ) -> Tuple[jnp.ndarray]:
816
+ if input_features.shape[1:] != (self.config.num_mel_bins, self.config.max_source_positions * 2):
817
+ raise ValueError(
818
+ "input_features.shape[1:], must be equal to (self.config.num_mel_bins,"
819
+ f" self.config.max_source_positions * 2) (got {input_features.shape[1:]}, but should be"
820
+ f" ({self.config.num_mel_bins}, {self.config.max_source_positions * 2}))"
821
+ )
822
+
823
+ input_features = input_features.transpose(0, 2, 1)
824
+ hidden_states = jax.nn.gelu(self.conv1(input_features), approximate=False)
825
+ hidden_states = with_sharding_constraint(hidden_states, ("batch", "embed", "num_mel"))
826
+ hidden_states = jax.nn.gelu(self.conv2(hidden_states), approximate=False)
827
+ hidden_states = with_sharding_constraint(hidden_states, ("batch", "length", "embed"))
828
+
829
+ embed_positions = self.embed_positions(jnp.arange(self.config.max_source_positions))
830
+ hidden_states = hidden_states + embed_positions
831
+
832
+ hidden_states = self.dropout_layer(hidden_states, deterministic=deterministic)
833
+
834
+ outputs = self.layers(
835
+ hidden_states,
836
+ attention_mask=None,
837
+ deterministic=deterministic,
838
+ output_attentions=output_attentions,
839
+ output_hidden_states=output_hidden_states,
840
+ return_dict=return_dict,
841
+ )
842
+
843
+ last_hidden_states = outputs[0]
844
+ last_hidden_states = self.layer_norm(last_hidden_states)
845
+
846
+ # update the last element in `hidden_states` after applying `layernorm` above
847
+ hidden_states = None
848
+ if output_hidden_states:
849
+ hidden_states = outputs[1]
850
+ hidden_states = hidden_states[:-1] + (last_hidden_states,)
851
+
852
+ if not return_dict:
853
+ outputs = (last_hidden_states, hidden_states) + (outputs[2:] if output_hidden_states else outputs[1:])
854
+ return tuple(v for v in outputs if v is not None)
855
+
856
+ return FlaxBaseModelOutput(
857
+ last_hidden_state=last_hidden_states,
858
+ hidden_states=hidden_states,
859
+ attentions=outputs.attentions,
860
+ )
861
+
862
+
863
+ class FlaxWhisperDecoder(nn.Module):
864
+ config: WhisperConfig
865
+ dtype: jnp.dtype = jnp.float32
866
+ params_dtype: jnp.dtype = jnp.float32
867
+
868
+ def setup(self) -> None:
869
+ self.embed_tokens = layers.Embed(
870
+ self.config.vocab_size, self.config.d_model, dtype=self.dtype, params_dtype=self.params_dtype
871
+ )
872
+ self.embed_positions = layers.Embed(
873
+ self.config.max_target_positions, self.config.d_model, dtype=self.dtype, params_dtype=self.params_dtype
874
+ )
875
+
876
+ self.layers = FlaxWhisperDecoderLayerCollection(self.config, dtype=self.dtype, params_dtype=self.params_dtype)
877
+
878
+ self.dropout_layer = nn.Dropout(rate=self.config.dropout)
879
+
880
+ self.layer_norm = layers.LayerNorm(dtype=self.dtype, epsilon=1e-5, params_dtype=self.params_dtype)
881
+
882
+ def __call__(
883
+ self,
884
+ input_ids: jnp.ndarray,
885
+ attention_mask: jnp.ndarray,
886
+ position_ids: jnp.ndarray,
887
+ encoder_hidden_states: Optional[jnp.ndarray] = None,
888
+ init_cache: bool = False,
889
+ output_attentions: bool = False,
890
+ output_hidden_states: bool = False,
891
+ return_dict: bool = True,
892
+ deterministic: bool = True,
893
+ ) -> Tuple[jnp.ndarray]:
894
+ input_embeds = self.embed_tokens(input_ids)
895
+ position_embeds = self.embed_positions(position_ids)
896
+
897
+ hidden_states = input_embeds + position_embeds
898
+ hidden_states = self.dropout_layer(hidden_states, deterministic=deterministic)
899
+
900
+ outputs = self.layers(
901
+ hidden_states,
902
+ attention_mask=attention_mask,
903
+ encoder_hidden_states=encoder_hidden_states,
904
+ deterministic=deterministic,
905
+ init_cache=init_cache,
906
+ output_attentions=output_attentions,
907
+ output_hidden_states=output_hidden_states,
908
+ return_dict=return_dict,
909
+ )
910
+
911
+ last_hidden_states = outputs[0]
912
+ last_hidden_states = self.layer_norm(last_hidden_states)
913
+
914
+ # update the last element in `hidden_states` after applying `layernorm` above
915
+ hidden_states = None
916
+ if output_hidden_states:
917
+ hidden_states = outputs[1]
918
+ hidden_states = hidden_states[:-1] + (last_hidden_states,)
919
+
920
+ if not return_dict:
921
+ outputs = (last_hidden_states, hidden_states) + (outputs[2:] if output_hidden_states else outputs[1:])
922
+ return tuple(v for v in outputs if v is not None)
923
+
924
+ return FlaxBaseModelOutputWithPastAndCrossAttentions(
925
+ last_hidden_state=last_hidden_states,
926
+ hidden_states=hidden_states,
927
+ attentions=outputs.attentions,
928
+ cross_attentions=outputs.cross_attentions,
929
+ )
930
+
931
+
932
+ class FlaxWhisperModule(nn.Module):
933
+ config: WhisperConfig
934
+ dtype: jnp.dtype = jnp.float32
935
+ params_dtype: jnp.dtype = jnp.float32
936
+
937
+ def setup(self) -> None:
938
+ self.encoder = FlaxWhisperEncoder(self.config, dtype=self.dtype, params_dtype=self.params_dtype)
939
+ self.decoder = FlaxWhisperDecoder(self.config, dtype=self.dtype, params_dtype=self.params_dtype)
940
+
941
+ def __call__(
942
+ self,
943
+ input_features: jnp.ndarray,
944
+ decoder_input_ids: jnp.ndarray,
945
+ decoder_attention_mask: jnp.ndarray,
946
+ decoder_position_ids: jnp.ndarray,
947
+ output_attentions: bool = False,
948
+ output_hidden_states: bool = False,
949
+ return_dict: bool = True,
950
+ deterministic: bool = True,
951
+ ):
952
+ encoder_outputs = self.encoder(
953
+ input_features,
954
+ output_attentions=output_attentions,
955
+ output_hidden_states=output_hidden_states,
956
+ return_dict=return_dict,
957
+ deterministic=deterministic,
958
+ )
959
+
960
+ decoder_outputs = self.decoder(
961
+ input_ids=decoder_input_ids,
962
+ attention_mask=decoder_attention_mask,
963
+ position_ids=decoder_position_ids,
964
+ encoder_hidden_states=encoder_outputs[0],
965
+ output_attentions=output_attentions,
966
+ output_hidden_states=output_hidden_states,
967
+ return_dict=return_dict,
968
+ deterministic=deterministic,
969
+ )
970
+
971
+ if not return_dict:
972
+ return decoder_outputs + encoder_outputs
973
+
974
+ return FlaxSeq2SeqModelOutput(
975
+ last_hidden_state=decoder_outputs.last_hidden_state,
976
+ decoder_hidden_states=decoder_outputs.hidden_states,
977
+ decoder_attentions=decoder_outputs.attentions,
978
+ cross_attentions=decoder_outputs.cross_attentions,
979
+ encoder_last_hidden_state=encoder_outputs.last_hidden_state,
980
+ encoder_hidden_states=encoder_outputs.hidden_states,
981
+ encoder_attentions=encoder_outputs.attentions,
982
+ )
983
+
984
+ def _get_encoder_module(self):
985
+ return self.encoder
986
+
987
+ def _get_decoder_module(self):
988
+ return self.decoder
989
+
990
+
991
+ class FlaxWhisperPreTrainedModel(FlaxPreTrainedModel):
992
+ config_class = WhisperConfig
993
+ base_model_prefix: str = "model"
994
+ main_input_name = "input_features"
995
+ module_class: nn.Module = None
996
+
997
+ def __init__(
998
+ self,
999
+ config: WhisperConfig,
1000
+ input_shape: Tuple[int, int, int] = None,
1001
+ seed: int = 0,
1002
+ dtype: jnp.dtype = jnp.float32,
1003
+ params_dtype: jnp.dtype = jnp.float32,
1004
+ _do_init: bool = True,
1005
+ **kwargs,
1006
+ ):
1007
+ if input_shape is None:
1008
+ input_shape = (1, config.num_mel_bins, 2 * config.max_source_positions)
1009
+
1010
+ module = self.module_class(config=config, dtype=dtype, params_dtype=params_dtype, **kwargs)
1011
+ super().__init__(config, module, input_shape=input_shape, seed=seed, dtype=dtype, _do_init=_do_init)
1012
+
1013
+ def init_weights(self, rng: jax.random.PRNGKey, input_shape: Tuple, params: FrozenDict = None) -> FrozenDict:
1014
+ # init input tensors
1015
+ input_features = jnp.zeros(input_shape, dtype="f4")
1016
+ input_features = input_features.at[(..., -1)].set(self.config.eos_token_id)
1017
+
1018
+ decoder_input_ids = jnp.zeros((input_shape[0], 1), dtype="i4")
1019
+ decoder_attention_mask = jnp.ones_like(decoder_input_ids)
1020
+
1021
+ batch_size, sequence_length = decoder_input_ids.shape
1022
+ decoder_position_ids = jnp.broadcast_to(jnp.arange(sequence_length)[None, :], (batch_size, sequence_length))
1023
+
1024
+ params_rng, dropout_rng = jax.random.split(rng)
1025
+ rngs = {"params": params_rng, "dropout": dropout_rng}
1026
+
1027
+ random_params = self.module.init(
1028
+ rngs,
1029
+ input_features=input_features,
1030
+ decoder_input_ids=decoder_input_ids,
1031
+ decoder_attention_mask=decoder_attention_mask,
1032
+ decoder_position_ids=decoder_position_ids,
1033
+ )["params"]
1034
+
1035
+ if params is not None:
1036
+ random_params = flatten_dict(unfreeze(random_params))
1037
+ params = flatten_dict(unfreeze(params))
1038
+ for missing_key in self._missing_keys:
1039
+ params[missing_key] = random_params[missing_key]
1040
+ self._missing_keys = set()
1041
+ return freeze(unflatten_dict(params))
1042
+ else:
1043
+ return random_params
1044
+
1045
+ # Copied from transformers.models.bart.modeling_flax_bart.FlaxBartPreTrainedModel.init_cache with Bart->Whisper
1046
+ def init_cache(self, batch_size, max_length, encoder_outputs):
1047
+ r"""
1048
+ Args:
1049
+ batch_size (`int`):
1050
+ batch_size used for fast auto-regressive decoding. Defines the batch size of the initialized cache.
1051
+ max_length (`int`):
1052
+ maximum possible length for auto-regressive decoding. Defines the sequence length of the initialized
1053
+ cache.
1054
+ encoder_outputs (`Union[FlaxBaseModelOutput, tuple(tuple(jnp.ndarray)]`):
1055
+ `encoder_outputs` consists of (`last_hidden_state`, *optional*: `hidden_states`, *optional*:
1056
+ `attentions`). `last_hidden_state` of shape `(batch_size, sequence_length, hidden_size)`, *optional*)
1057
+ is a sequence of hidden-states at the output of the last layer of the encoder. Used in the
1058
+ cross-attention of the decoder.
1059
+ """
1060
+ # init input variables to retrieve cache
1061
+ decoder_input_ids = jnp.ones((batch_size, max_length), dtype="i4")
1062
+ decoder_attention_mask = jnp.ones_like(decoder_input_ids)
1063
+ decoder_position_ids = jnp.broadcast_to(
1064
+ jnp.arange(jnp.atleast_2d(decoder_input_ids).shape[-1]), decoder_input_ids.shape
1065
+ )
1066
+
1067
+ def _decoder_forward(module, decoder_input_ids, decoder_attention_mask, decoder_position_ids, **kwargs):
1068
+ decoder_module = module._get_decoder_module()
1069
+ return decoder_module(
1070
+ decoder_input_ids,
1071
+ decoder_attention_mask,
1072
+ decoder_position_ids,
1073
+ **kwargs,
1074
+ )
1075
+
1076
+ init_variables = self.module.init(
1077
+ jax.random.PRNGKey(0),
1078
+ decoder_input_ids=decoder_input_ids,
1079
+ decoder_attention_mask=decoder_attention_mask,
1080
+ decoder_position_ids=decoder_position_ids,
1081
+ encoder_hidden_states=encoder_outputs[0],
1082
+ init_cache=True,
1083
+ method=_decoder_forward, # we only need to call the decoder to init the cache
1084
+ )
1085
+ return unfreeze(init_variables["cache"])
1086
+
1087
+ @add_start_docstrings(WHISPER_ENCODE_INPUTS_DOCSTRING)
1088
+ @replace_return_docstrings(output_type=FlaxBaseModelOutput, config_class=WhisperConfig)
1089
+ def encode(
1090
+ self,
1091
+ input_features: jnp.ndarray,
1092
+ attention_mask: Optional[jnp.ndarray] = None,
1093
+ output_attentions: Optional[bool] = None,
1094
+ output_hidden_states: Optional[bool] = None,
1095
+ return_dict: Optional[bool] = None,
1096
+ train: bool = False,
1097
+ params: dict = None,
1098
+ dropout_rng: PRNGKey = None,
1099
+ **kwargs,
1100
+ ):
1101
+ r"""
1102
+ Returns:
1103
+
1104
+ Example:
1105
+
1106
+ ```python
1107
+ >>> from transformers import WhisperProcessor, FlaxWhisperForConditionalGeneration
1108
+ >>> from datasets import load_dataset
1109
+
1110
+ >>> processor = WhisperProcessor.from_pretrained("openai/whisper-tiny.en")
1111
+ >>> model = FlaxWhisperForConditionalGeneration.from_pretrained("openai/whisper-tiny.en", from_pt=True)
1112
+ >>> ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")
1113
+ >>> inputs = processor(ds[0]["audio"]["array"], return_tensors="np")
1114
+ >>> input_features = inputs.input_features
1115
+ >>> encoder_outputs = model.encode(input_features=input_features)
1116
+ ```"""
1117
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
1118
+ output_hidden_states = (
1119
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
1120
+ )
1121
+ return_dict = return_dict if return_dict is not None else self.config.return_dict
1122
+
1123
+ # Handle any PRNG if needed
1124
+ rngs = {}
1125
+ if dropout_rng is not None:
1126
+ rngs["dropout"] = dropout_rng
1127
+
1128
+ def _encoder_forward(module, input_features, **kwargs):
1129
+ encode_module = module._get_encoder_module()
1130
+ return encode_module(input_features, **kwargs)
1131
+
1132
+ return self.module.apply(
1133
+ {"params": params or self.params},
1134
+ input_features=jnp.array(input_features, dtype="f4"),
1135
+ output_attentions=output_attentions,
1136
+ output_hidden_states=output_hidden_states,
1137
+ return_dict=return_dict,
1138
+ deterministic=not train,
1139
+ rngs=rngs,
1140
+ method=_encoder_forward,
1141
+ )
1142
+
1143
+ @add_start_docstrings(WHISPER_DECODE_INPUTS_DOCSTRING)
1144
+ @replace_return_docstrings(output_type=FlaxBaseModelOutputWithPastAndCrossAttentions, config_class=WhisperConfig)
1145
+ def decode(
1146
+ self,
1147
+ decoder_input_ids,
1148
+ encoder_outputs,
1149
+ encoder_attention_mask: Optional[jnp.ndarray] = None,
1150
+ decoder_attention_mask: Optional[jnp.ndarray] = None,
1151
+ decoder_position_ids: Optional[jnp.ndarray] = None,
1152
+ past_key_values: dict = None,
1153
+ output_attentions: Optional[bool] = None,
1154
+ output_hidden_states: Optional[bool] = None,
1155
+ return_dict: Optional[bool] = None,
1156
+ train: bool = False,
1157
+ params: dict = None,
1158
+ dropout_rng: PRNGKey = None,
1159
+ ):
1160
+ r"""
1161
+ Returns:
1162
+
1163
+ Example:
1164
+
1165
+ ```python
1166
+ >>> from transformers import WhisperProcessor, FlaxWhisperForConditionalGeneration
1167
+ >>> from datasets import load_dataset
1168
+
1169
+ >>> processor = WhisperProcessor.from_pretrained("openai/whisper-tiny.en")
1170
+ >>> model = FlaxWhisperForConditionalGeneration.from_pretrained("openai/whisper-tiny.en", from_pt=True)
1171
+ >>> ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")
1172
+ >>> inputs = processor(ds[0]["audio"]["array"], return_tensors="np")
1173
+ >>> input_features = inputs.input_features
1174
+ >>> encoder_outputs = model.encode(input_features=input_features)
1175
+ >>> decoder_start_token_id = model.config.decoder_start_token_id
1176
+
1177
+ >>> decoder_input_ids = jnp.ones((inputs.input_ids.shape[0], 1), dtype="i4") * decoder_start_token_id
1178
+
1179
+ >>> outputs = model.decode(decoder_input_ids, encoder_outputs)
1180
+ >>> last_decoder_hidden_states = outputs.last_hidden_state
1181
+ ```"""
1182
+
1183
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
1184
+ output_hidden_states = (
1185
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
1186
+ )
1187
+ return_dict = return_dict if return_dict is not None else self.config.return_dict
1188
+
1189
+ encoder_hidden_states = encoder_outputs[0]
1190
+
1191
+ batch_size, sequence_length = decoder_input_ids.shape
1192
+ if decoder_position_ids is None:
1193
+ if past_key_values is not None:
1194
+ raise ValueError("Make sure to provide `decoder_position_ids` when passing `past_key_values`.")
1195
+
1196
+ if decoder_attention_mask is not None:
1197
+ decoder_position_ids = (decoder_attention_mask.cumsum(-1) * decoder_attention_mask) - 1
1198
+ else:
1199
+ decoder_position_ids = jnp.broadcast_to(
1200
+ jnp.arange(sequence_length)[None, :], (batch_size, sequence_length)
1201
+ )
1202
+
1203
+ if decoder_attention_mask is None:
1204
+ decoder_attention_mask = jnp.ones((batch_size, sequence_length))
1205
+
1206
+ # Handle any PRNG if needed
1207
+ rngs = {}
1208
+ if dropout_rng is not None:
1209
+ rngs["dropout"] = dropout_rng
1210
+
1211
+ inputs = {"params": params or self.params}
1212
+
1213
+ # if past_key_values are passed then cache is already initialized a private flag init_cache has to be
1214
+ # passed down to ensure cache is used. It has to be made sure that cache is marked as mutable so that
1215
+ # it can be changed by FlaxWhisperAttention module
1216
+ if past_key_values:
1217
+ inputs["cache"] = past_key_values
1218
+ mutable = ["cache"]
1219
+ else:
1220
+ mutable = False
1221
+
1222
+ def _decoder_forward(module, decoder_input_ids, decoder_attention_mask, decoder_position_ids, **kwargs):
1223
+ decoder_module = module._get_decoder_module()
1224
+ return decoder_module(
1225
+ input_ids=decoder_input_ids,
1226
+ attention_mask=decoder_attention_mask,
1227
+ position_ids=decoder_position_ids,
1228
+ **kwargs,
1229
+ )
1230
+
1231
+ outputs = self.module.apply(
1232
+ inputs,
1233
+ decoder_input_ids=jnp.array(decoder_input_ids, dtype="i4"),
1234
+ decoder_attention_mask=jnp.array(decoder_attention_mask, dtype="i4"),
1235
+ decoder_position_ids=jnp.array(decoder_position_ids, dtype="i4"),
1236
+ encoder_hidden_states=encoder_hidden_states,
1237
+ output_attentions=output_attentions,
1238
+ output_hidden_states=output_hidden_states,
1239
+ return_dict=return_dict,
1240
+ deterministic=not train,
1241
+ rngs=rngs,
1242
+ mutable=mutable,
1243
+ method=_decoder_forward,
1244
+ )
1245
+
1246
+ # add updated cache to model output
1247
+ if past_key_values is not None and return_dict:
1248
+ outputs, past = outputs
1249
+ outputs["past_key_values"] = unfreeze(past["cache"])
1250
+ return outputs
1251
+ elif past_key_values is not None and not return_dict:
1252
+ outputs, past = outputs
1253
+ outputs = outputs[:1] + (unfreeze(past["cache"]),) + outputs[1:]
1254
+
1255
+ return outputs
1256
+
1257
+ @add_start_docstrings_to_model_forward(WHISPER_INPUTS_DOCSTRING)
1258
+ def __call__(
1259
+ self,
1260
+ input_features: jnp.ndarray,
1261
+ decoder_input_ids: jnp.ndarray,
1262
+ attention_mask: Optional[jnp.ndarray] = None,
1263
+ decoder_attention_mask: Optional[jnp.ndarray] = None,
1264
+ position_ids: Optional[jnp.ndarray] = None,
1265
+ decoder_position_ids: Optional[jnp.ndarray] = None,
1266
+ output_attentions: Optional[bool] = None,
1267
+ output_hidden_states: Optional[bool] = None,
1268
+ return_dict: Optional[bool] = None,
1269
+ train: bool = False,
1270
+ params: dict = None,
1271
+ dropout_rng: PRNGKey = None,
1272
+ ):
1273
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
1274
+ output_hidden_states = (
1275
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
1276
+ )
1277
+ return_dict = return_dict if return_dict is not None else self.config.return_dict
1278
+
1279
+ # prepare decoder inputs
1280
+ if decoder_position_ids is None:
1281
+ if decoder_attention_mask is not None:
1282
+ decoder_position_ids = (decoder_attention_mask.cumsum(-1) * decoder_attention_mask) - 1
1283
+ else:
1284
+ batch_size, sequence_length = decoder_input_ids.shape
1285
+ decoder_position_ids = jnp.broadcast_to(
1286
+ jnp.arange(sequence_length)[None, :], (batch_size, sequence_length)
1287
+ )
1288
+ if decoder_attention_mask is None:
1289
+ decoder_attention_mask = jnp.ones_like(decoder_input_ids)
1290
+
1291
+ # Handle any PRNG if needed
1292
+ rngs = {"dropout": dropout_rng} if dropout_rng is not None else {}
1293
+
1294
+ return self.module.apply(
1295
+ {"params": params or self.params},
1296
+ input_features=jnp.array(input_features, dtype="f4"),
1297
+ decoder_input_ids=jnp.array(decoder_input_ids, dtype="i4"),
1298
+ decoder_attention_mask=jnp.array(decoder_attention_mask, dtype="i4"),
1299
+ decoder_position_ids=jnp.array(decoder_position_ids, dtype="i4"),
1300
+ output_attentions=output_attentions,
1301
+ output_hidden_states=output_hidden_states,
1302
+ return_dict=return_dict,
1303
+ deterministic=not train,
1304
+ rngs=rngs,
1305
+ )
1306
+
1307
+
1308
+ @add_start_docstrings(
1309
+ "The bare Whisper Model transformer outputting raw hidden-states without any specific head on top.",
1310
+ WHISPER_START_DOCSTRING,
1311
+ )
1312
+ class FlaxWhisperModel(FlaxWhisperPreTrainedModel):
1313
+ config: WhisperConfig
1314
+ dtype: jnp.dtype = jnp.float32 # the dtype of the computation
1315
+ params_dtype: jnp.dtype = jnp.float32
1316
+ module_class = FlaxWhisperModule
1317
+
1318
+
1319
+ append_call_sample_docstring(FlaxWhisperModel, _CHECKPOINT_FOR_DOC, FlaxSeq2SeqModelOutput, _CONFIG_FOR_DOC)
1320
+
1321
+
1322
+ class FlaxWhisperForConditionalGenerationModule(nn.Module):
1323
+ config: WhisperConfig
1324
+ dtype: jnp.dtype = jnp.float32
1325
+ params_dtype: jnp.dtype = jnp.float32
1326
+
1327
+ def setup(self) -> None:
1328
+ self.model = FlaxWhisperModule(config=self.config, dtype=self.dtype, params_dtype=self.params_dtype)
1329
+ self.lm_head = layers.DenseGeneral(
1330
+ self.config.vocab_size,
1331
+ use_bias=False,
1332
+ dtype=self.dtype,
1333
+ params_dtype=self.params_dtype,
1334
+ kernel_axes=("embed", "vocab"),
1335
+ )
1336
+
1337
+ def _get_encoder_module(self):
1338
+ return self.model.encoder
1339
+
1340
+ def _get_decoder_module(self):
1341
+ return self.model.decoder
1342
+
1343
+ def __call__(
1344
+ self,
1345
+ input_features,
1346
+ decoder_input_ids,
1347
+ decoder_attention_mask: jnp.ndarray = None,
1348
+ decoder_position_ids: jnp.ndarray = None,
1349
+ position_ids: jnp.ndarray = None,
1350
+ attention_mask: jnp.ndarray = None,
1351
+ output_attentions: bool = False,
1352
+ output_hidden_states: bool = False,
1353
+ return_dict: bool = True,
1354
+ deterministic: bool = True,
1355
+ ):
1356
+ outputs = self.model(
1357
+ input_features=input_features,
1358
+ decoder_input_ids=decoder_input_ids,
1359
+ decoder_attention_mask=decoder_attention_mask,
1360
+ decoder_position_ids=decoder_position_ids,
1361
+ output_attentions=output_attentions,
1362
+ output_hidden_states=output_hidden_states,
1363
+ return_dict=return_dict,
1364
+ deterministic=deterministic,
1365
+ )
1366
+
1367
+ hidden_states = outputs[0]
1368
+
1369
+ if self.config.tie_word_embeddings:
1370
+ shared_embedding = self.model.decoder.embed_tokens.variables["params"]["embedding"]
1371
+ lm_logits = self.lm_head.apply({"params": {"kernel": shared_embedding.T}}, hidden_states)
1372
+ else:
1373
+ lm_logits = self.lm_head(hidden_states)
1374
+
1375
+ if not return_dict:
1376
+ output = (lm_logits,) + outputs[1:]
1377
+ return output
1378
+
1379
+ return FlaxSeq2SeqLMOutput(
1380
+ logits=lm_logits,
1381
+ decoder_hidden_states=outputs.decoder_hidden_states,
1382
+ decoder_attentions=outputs.decoder_attentions,
1383
+ cross_attentions=outputs.cross_attentions,
1384
+ encoder_last_hidden_state=outputs.encoder_last_hidden_state,
1385
+ encoder_hidden_states=outputs.encoder_hidden_states,
1386
+ encoder_attentions=outputs.encoder_attentions,
1387
+ )
1388
+
1389
+
1390
+ @add_start_docstrings("The Whisper Model with a language modeling head.", WHISPER_START_DOCSTRING)
1391
+ class FlaxWhisperForConditionalGeneration(FlaxWhisperPreTrainedModel):
1392
+ module_class = FlaxWhisperForConditionalGenerationModule
1393
+ dtype: jnp.dtype = jnp.float32
1394
+ params_dtype: jnp.dtype = jnp.float32
1395
+
1396
+ @add_start_docstrings(WHISPER_DECODE_INPUTS_DOCSTRING)
1397
+ @replace_return_docstrings(output_type=FlaxCausalLMOutputWithCrossAttentions, config_class=WhisperConfig)
1398
+ def decode(
1399
+ self,
1400
+ decoder_input_ids,
1401
+ encoder_outputs,
1402
+ encoder_attention_mask: Optional[jnp.ndarray] = None,
1403
+ decoder_attention_mask: Optional[jnp.ndarray] = None,
1404
+ decoder_position_ids: Optional[jnp.ndarray] = None,
1405
+ past_key_values: dict = None,
1406
+ output_attentions: Optional[bool] = None,
1407
+ output_hidden_states: Optional[bool] = None,
1408
+ return_dict: Optional[bool] = None,
1409
+ train: bool = False,
1410
+ params: dict = None,
1411
+ dropout_rng: PRNGKey = None,
1412
+ ):
1413
+ r"""
1414
+ Returns:
1415
+
1416
+ Example:
1417
+
1418
+ ```python
1419
+ >>> from transformers import WhisperProcessor, FlaxWhisperForConditionalGeneration
1420
+ >>> from datasets import load_dataset
1421
+
1422
+ >>> processor = WhisperProcessor.from_pretrained("openai/whisper-tiny.en")
1423
+ >>> model = FlaxWhisperForConditionalGeneration.from_pretrained("openai/whisper-tiny.en", from_pt=True)
1424
+ >>> ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")
1425
+ >>> inputs = processor(ds[0]["audio"]["array"], return_tensors="np")
1426
+ >>> input_features = inputs.input_features
1427
+ >>> encoder_outputs = model.encode(input_features=input_features)
1428
+ >>> decoder_start_token_id = model.config.decoder_start_token_id
1429
+
1430
+ >>> decoder_input_ids = jnp.ones((inputs.input_ids.shape[0], 1), dtype="i4") * decoder_start_token_id
1431
+
1432
+ >>> outputs = model.decode(decoder_input_ids, encoder_outputs)
1433
+ >>> last_decoder_hidden_states = outputs.last_hidden_state
1434
+ ```"""
1435
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
1436
+ output_hidden_states = (
1437
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
1438
+ )
1439
+ return_dict = return_dict if return_dict is not None else self.config.return_dict
1440
+
1441
+ encoder_hidden_states = encoder_outputs[0]
1442
+
1443
+ batch_size, sequence_length = decoder_input_ids.shape
1444
+ if decoder_position_ids is None:
1445
+ if past_key_values is not None:
1446
+ raise ValueError("Make sure to provide `decoder_position_ids` when passing `past_key_values`.")
1447
+
1448
+ if decoder_attention_mask is not None:
1449
+ decoder_position_ids = (decoder_attention_mask.cumsum(-1) * decoder_attention_mask) - 1
1450
+ else:
1451
+ decoder_position_ids = jnp.broadcast_to(
1452
+ jnp.arange(sequence_length)[None, :], (batch_size, sequence_length)
1453
+ )
1454
+ if decoder_attention_mask is None:
1455
+ decoder_attention_mask = jnp.ones((batch_size, sequence_length), dtype="i4")
1456
+
1457
+ # Handle any PRNG if needed
1458
+ rngs = {}
1459
+ if dropout_rng is not None:
1460
+ rngs["dropout"] = dropout_rng
1461
+
1462
+ inputs = {"params": params or self.params}
1463
+
1464
+ # if past_key_values are passed then cache is already initialized a private flag init_cache has to be
1465
+ # passed down to ensure cache is used. It has to be made sure that cache is marked as mutable so that
1466
+ # it can be changed by FlaxWhisperAttention module
1467
+ if past_key_values:
1468
+ inputs["cache"] = past_key_values
1469
+ mutable = ["cache"]
1470
+ else:
1471
+ mutable = False
1472
+
1473
+ def _decoder_forward(module, decoder_input_ids, decoder_attention_mask, decoder_position_ids, **kwargs):
1474
+ decoder_module = module._get_decoder_module()
1475
+ outputs = decoder_module(
1476
+ input_ids=decoder_input_ids,
1477
+ attention_mask=decoder_attention_mask,
1478
+ position_ids=decoder_position_ids,
1479
+ **kwargs,
1480
+ )
1481
+ hidden_states = outputs[0]
1482
+
1483
+ if self.config.tie_word_embeddings:
1484
+ shared_embedding = module.model.decoder.embed_tokens.variables["params"]["embedding"]
1485
+ lm_logits = module.lm_head.apply({"params": {"kernel": shared_embedding.T}}, hidden_states)
1486
+ else:
1487
+ lm_logits = module.lm_head(hidden_states)
1488
+
1489
+ return lm_logits, outputs
1490
+
1491
+ outputs = self.module.apply(
1492
+ inputs,
1493
+ decoder_input_ids=jnp.array(decoder_input_ids, dtype="i4"),
1494
+ decoder_attention_mask=jnp.array(decoder_attention_mask, dtype="i4"),
1495
+ decoder_position_ids=jnp.array(decoder_position_ids, dtype="i4"),
1496
+ encoder_hidden_states=encoder_hidden_states,
1497
+ output_attentions=output_attentions,
1498
+ output_hidden_states=output_hidden_states,
1499
+ return_dict=return_dict,
1500
+ deterministic=not train,
1501
+ rngs=rngs,
1502
+ mutable=mutable,
1503
+ method=_decoder_forward,
1504
+ )
1505
+
1506
+ if past_key_values is None:
1507
+ lm_logits, decoder_outputs = outputs
1508
+ else:
1509
+ (lm_logits, decoder_outputs), past = outputs
1510
+
1511
+ if return_dict:
1512
+ outputs = FlaxCausalLMOutputWithCrossAttentions(
1513
+ logits=lm_logits,
1514
+ hidden_states=decoder_outputs.hidden_states,
1515
+ attentions=decoder_outputs.attentions,
1516
+ cross_attentions=decoder_outputs.cross_attentions,
1517
+ )
1518
+ else:
1519
+ outputs = (lm_logits,) + decoder_outputs[1:]
1520
+
1521
+ # add updated cache to model output
1522
+ if past_key_values is not None and return_dict:
1523
+ outputs["past_key_values"] = unfreeze(past["cache"])
1524
+ return outputs
1525
+ elif past_key_values is not None and not return_dict:
1526
+ outputs = outputs[:1] + (unfreeze(past["cache"]),) + outputs[1:]
1527
+
1528
+ return outputs
1529
+
1530
+ def generate(
1531
+ self,
1532
+ input_features,
1533
+ generation_config=None,
1534
+ logits_processor=None,
1535
+ return_timestamps=None,
1536
+ task=None,
1537
+ language=None,
1538
+ is_multilingual=None,
1539
+ **kwargs,
1540
+ ):
1541
+ if generation_config is None:
1542
+ generation_config = self.generation_config
1543
+
1544
+ if return_timestamps is not None:
1545
+ generation_config.return_timestamps = return_timestamps
1546
+
1547
+ if task is not None:
1548
+ generation_config.task = task
1549
+
1550
+ if is_multilingual is not None:
1551
+ generation_config.is_multilingual = is_multilingual
1552
+
1553
+ if language is not None:
1554
+ generation_config.language = language
1555
+
1556
+ if kwargs is not None and "decoder_input_ids" in kwargs:
1557
+ decoder_input_length = len(kwargs["decoder_input_ids"])
1558
+ else:
1559
+ decoder_input_length = 1
1560
+
1561
+ forced_decoder_ids = []
1562
+
1563
+ if hasattr(generation_config, "is_multilingual") and generation_config.is_multilingual:
1564
+ if hasattr(generation_config, "language"):
1565
+ forced_decoder_ids.append((1, generation_config.lang_to_id[generation_config.language]))
1566
+ else:
1567
+ forced_decoder_ids.append((1, None))
1568
+
1569
+ if hasattr(generation_config, "task"):
1570
+ forced_decoder_ids.append((2, generation_config.task_to_id[generation_config.task]))
1571
+ else:
1572
+ forced_decoder_ids.append((2, generation_config.task_to_id["transcribe"]))
1573
+
1574
+ if (
1575
+ hasattr(generation_config, "return_timestamps") and generation_config.return_timestamps
1576
+ ) or return_timestamps:
1577
+ logits_processor = [
1578
+ FlaxWhisperTimeStampLogitsProcessor(generation_config, self.config, decoder_input_length)
1579
+ ]
1580
+ else:
1581
+ if forced_decoder_ids and forced_decoder_ids[-1][0] != generation_config.no_timestamps_token_id:
1582
+ idx = forced_decoder_ids[-1][0] + 1 if forced_decoder_ids else 1
1583
+ forced_decoder_ids.append((idx, generation_config.no_timestamps_token_id))
1584
+
1585
+ if len(forced_decoder_ids) > 0:
1586
+ generation_config.forced_decoder_ids = forced_decoder_ids
1587
+
1588
+ return super().generate(
1589
+ input_features,
1590
+ generation_config,
1591
+ logits_processor=logits_processor,
1592
+ **kwargs,
1593
+ )
1594
+
1595
+ def pipeline_generate(
1596
+ self,
1597
+ input_features,
1598
+ forced_decoder_ids,
1599
+ return_timestamps=False,
1600
+ generation_config=None,
1601
+ **kwargs,
1602
+ ):
1603
+ if generation_config is None:
1604
+ generation_config = self.generation_config
1605
+
1606
+ # override the generation config forced decoder ids in preference of the ones we have set
1607
+ generation_config.forced_decoder_ids = None
1608
+
1609
+ logits_processor = FlaxLogitsProcessorList()
1610
+
1611
+ logits_processor.append(FlaxStaticForceTokensLogitsProcessor(forced_decoder_ids))
1612
+
1613
+ if hasattr(generation_config, "return_timestamps") and return_timestamps:
1614
+ logits_processor.append(FlaxWhisperTimeStampLogitsProcessor(generation_config, self.config, 1))
1615
+
1616
+ return super().generate(
1617
+ input_features,
1618
+ generation_config,
1619
+ logits_processor=logits_processor,
1620
+ **kwargs,
1621
+ )
1622
+
1623
+ def prepare_inputs_for_generation(
1624
+ self,
1625
+ decoder_input_ids,
1626
+ max_length,
1627
+ attention_mask: Optional[jax.Array] = None,
1628
+ decoder_attention_mask: Optional[jax.Array] = None,
1629
+ encoder_outputs=None,
1630
+ **kwargs,
1631
+ ):
1632
+ # initializing the cache
1633
+ batch_size, seq_length = decoder_input_ids.shape
1634
+
1635
+ past_key_values = self.init_cache(batch_size, max_length, encoder_outputs)
1636
+ # Note that usually one would have to put 0's in the attention_mask for x > input_ids.shape[-1] and x < cache_length.
1637
+ # But since the decoder uses a causal mask, those positions are masked anyways.
1638
+ # Thus we can create a single static attention_mask here, which is more efficient for compilation
1639
+ extended_attention_mask = jnp.ones((batch_size, max_length), dtype="i4")
1640
+ if decoder_attention_mask is not None:
1641
+ position_ids = decoder_attention_mask.cumsum(-1) - 1
1642
+ extended_attention_mask = lax.dynamic_update_slice(extended_attention_mask, decoder_attention_mask, (0, 0))
1643
+ else:
1644
+ position_ids = jnp.broadcast_to(jnp.arange(seq_length, dtype="i4")[None, :], (batch_size, seq_length))
1645
+
1646
+ return {
1647
+ "past_key_values": past_key_values,
1648
+ "encoder_outputs": encoder_outputs,
1649
+ "encoder_attention_mask": attention_mask,
1650
+ "decoder_attention_mask": extended_attention_mask,
1651
+ "decoder_position_ids": position_ids,
1652
+ }
1653
+
1654
+ def update_inputs_for_generation(self, model_outputs, model_kwargs):
1655
+ model_kwargs["past_key_values"] = model_outputs.past_key_values
1656
+ model_kwargs["decoder_position_ids"] = model_kwargs["decoder_position_ids"][:, -1:] + 1
1657
+ return model_kwargs
1658
+
1659
+
1660
+ FLAX_WHISPER_CONDITIONAL_GENERATION_DOCSTRING = r"""
1661
+ Returns:
1662
+
1663
+ Transcription example:
1664
+
1665
+ ```python
1666
+ >>> from transformers import WhisperProcessor, FlaxWhisperForConditionalGeneration
1667
+ >>> from datasets import load_dataset
1668
+
1669
+ >>> processor = WhisperProcessor.from_pretrained("openai/whisper-tiny.en")
1670
+ >>> model = FlaxWhisperForConditionalGeneration.from_pretrained("openai/whisper-tiny.en", from_pt=True)
1671
+ >>> ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")
1672
+ >>> inputs = processor(ds[0]["audio"]["array"], return_tensors="np")
1673
+ >>> input_features = inputs.input_features
1674
+ >>> generated_ids = model.generate(input_ids=input_features)
1675
+ >>> transcription = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
1676
+ >>> transcription
1677
+ ' Mr. Quilter is the apostle of the middle classes, and we are glad to welcome his gospel.'
1678
+ ```
1679
+ """
1680
+
1681
+ overwrite_call_docstring(
1682
+ FlaxWhisperForConditionalGeneration, WHISPER_INPUTS_DOCSTRING + FLAX_WHISPER_CONDITIONAL_GENERATION_DOCSTRING
1683
+ )
1684
+ append_replace_return_docstrings(
1685
+ FlaxWhisperForConditionalGeneration, output_type=FlaxSeq2SeqLMOutput, config_class=_CONFIG_FOR_DOC
1686
+ )
whisper_jax/partitioner.py ADDED
@@ -0,0 +1,939 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # coding=utf-8
2
+ # Copyright 2023 The T5X Authors and The HuggingFace Inc. team.
3
+ #
4
+ # Licensed under the Apache License, Version 2.0 (the "License");
5
+ # you may not use this file except in compliance with the License.
6
+ # You may obtain a copy of the License at
7
+ #
8
+ # http://www.apache.org/licenses/LICENSE-2.0
9
+ #
10
+ # Unless required by applicable law or agreed to in writing, software
11
+ # distributed under the License is distributed on an "AS IS" BASIS,
12
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
+ # See the License for the specific language governing permissions and
14
+ # limitations under the License.
15
+
16
+ """Utilities for partitioning."""
17
+
18
+ import abc
19
+ import collections
20
+ import dataclasses
21
+ import typing
22
+ from typing import Any, Callable, Optional, Sequence, Tuple, Union
23
+
24
+ import cached_property
25
+ import jax
26
+ import numpy as np
27
+ from absl import logging
28
+ from flax import traverse_util
29
+ from flax.linen import partitioning as flax_partitioning
30
+ from jax import numpy as jnp
31
+ from jax import random
32
+ from jax.experimental import multihost_utils
33
+ from jax.experimental.mesh_utils import create_hybrid_device_mesh
34
+ from jax.experimental.pjit import pjit as jax_pjit
35
+ from jax.sharding import Mesh, PartitionSpec
36
+
37
+
38
+ JaxDevice = Any
39
+ TpuMesh = Tuple[int, int, int, int] # (x, y, z, num_cores).
40
+ OtherMesh = Tuple[int, int]
41
+ HardwareMesh = Union[TpuMesh, OtherMesh]
42
+ PyTreeDef = type(jax.tree_util.tree_structure(None))
43
+ TrainState = Any
44
+ LogicalAxisRules = Sequence[Tuple[str, Optional[str]]]
45
+
46
+ if typing.TYPE_CHECKING: # See b/163639353
47
+ cached_property = property # pylint: disable=invalid-name
48
+ else:
49
+ cached_property = cached_property.cached_property
50
+
51
+
52
+ class AxisNames(tuple):
53
+ """Tuple of strings specifying name for each axis.
54
+
55
+ We create a separate class for this so JAX's pytree utilities can distinguish
56
+ it from a tuple that should be treated as a pytree, instead treating it as a
57
+ leaf.
58
+ """
59
+
60
+ def __new__(cls, *names):
61
+ return tuple.__new__(AxisNames, names)
62
+
63
+ def __repr__(self):
64
+ return "AxisNames%s" % tuple.__repr__(self)
65
+
66
+
67
+ # pjit wrappers for cpu fallback.
68
+ # ----------------------------------------------------------------------------
69
+ # TODO(levskaya): This function is now no different than jax_pjit, but callers
70
+ # currently depend on `backend` argument
71
+ def pjit(
72
+ fun: Callable, # pylint: disable=g-bare-generic
73
+ in_axis_resources,
74
+ out_axis_resources,
75
+ static_argnums: Union[int, Sequence[int]] = (),
76
+ donate_argnums: Union[int, Sequence[int]] = (),
77
+ backend: Optional[str] = None,
78
+ ):
79
+ """Wrapper for pjit."""
80
+ del backend
81
+ return jax_pjit(
82
+ fun, in_axis_resources, out_axis_resources, static_argnums=static_argnums, donate_argnums=donate_argnums
83
+ )
84
+
85
+
86
+ # pjit wrappers for cpu fallback.
87
+ # -----------------------------------------------------------------------------
88
+ # TODO(levskaya): upstream this fallback behavior to jax pjit.
89
+ def pjit_with_cpu_fallback(
90
+ fun: Callable, # pylint: disable=g-bare-generic
91
+ in_axis_resources,
92
+ out_axis_resources,
93
+ static_argnums: Union[int, Sequence[int]] = (),
94
+ donate_argnums: Union[int, Sequence[int]] = (),
95
+ backend: Optional[str] = None,
96
+ ):
97
+ """Wrapper for pjit that calls normal jit on cpu."""
98
+ if jax.devices(backend)[0].platform == "cpu":
99
+ return jax.jit(fun, static_argnums=static_argnums, donate_argnums=donate_argnums)
100
+ else:
101
+ return jax_pjit(
102
+ fun, in_axis_resources, out_axis_resources, static_argnums=static_argnums, donate_argnums=donate_argnums
103
+ )
104
+
105
+
106
+ def with_sharding_constraint(x, axis_resources):
107
+ """Wrapper for pjit with_sharding_constraint, no-op on cpu or outside pjit."""
108
+ if jax.devices()[0].platform == "cpu" or not global_mesh_defined():
109
+ return x
110
+ else:
111
+ return jax.experimental.pjit.with_sharding_constraint(x, axis_resources)
112
+
113
+
114
+ # pjit Mesh creation functions.
115
+ # -----------------------------------------------------------------------------
116
+ def bounds_from_last_device(last_device: JaxDevice) -> HardwareMesh:
117
+ """Get the bound from the given last device."""
118
+ # Must be passed the device at the highest-coordinate corner of the
119
+ # relevant mesh, which is a requirement we know is satisfied by the last
120
+ # device in jax.devices().
121
+ if hasattr(last_device, "coords"):
122
+ x, y, z = last_device.coords
123
+ return x + 1, y + 1, z + 1, last_device.core_on_chip + 1
124
+ else:
125
+ # On non-TPU platforms, the "mesh" is hosts x devices per host in order
126
+ # to take advantage of faster within-host interconnect.
127
+ return jax.host_count(), jax.local_device_count()
128
+
129
+
130
+ def get_coords(device: JaxDevice) -> HardwareMesh:
131
+ """Returns the coordinates of the given device."""
132
+ if hasattr(device, "coords"):
133
+ return (*device.coords, device.core_on_chip)
134
+ return (device.process_index, device.id % jax.local_device_count())
135
+
136
+
137
+ def global_mesh_defined():
138
+ """Checks if global xmap/pjit mesh resource environment is defined."""
139
+ maps_env = jax.experimental.maps.thread_resources.env
140
+ return maps_env.physical_mesh.devices.shape != () # pylint: disable=g-explicit-bool-comparison
141
+
142
+
143
+ def get_mesh(
144
+ model_parallel_submesh: HardwareMesh,
145
+ input_devices: Sequence[JaxDevice] = (),
146
+ input_local_devices: Sequence[JaxDevice] = (),
147
+ tile_by_host_if_needed: bool = True,
148
+ backend: Optional[str] = None,
149
+ ) -> Mesh:
150
+ """Construct an xmap/pjit Mesh for the given model-parallel submesh.
151
+
152
+ The resulting mesh has two resource axes: 'model', with the provided submesh
153
+ shape, and 'data', which covers the rest of the mesh.
154
+
155
+ Args:
156
+ model_parallel_submesh: a HardwareMesh spec, namely (x,y,z,core) on TPU for
157
+ a single model-parallel replica's "tile" in the physical device mesh. The
158
+ first three elements (`x`, `y`, and `z`) should be factors of the pod
159
+ slice; e.g., if you are using df_4x8, then `x` should be a factor of 4
160
+ (one of 1, 2, 4), `y` should be a factor of 8 (one of 1, 2, 4, 8), and `z`
161
+ must be 1, because TPU v3 slices are only 2D. `z` can be >1 for TPU v4
162
+ (and maybe later TPUs) that allow 3D slices. `core` is the number of cores
163
+ to use from each TPU node. As communication is usually fastest inside the
164
+ same node, if you need a tile of more than 1 core, then
165
+ you should first increase `core`: e.g., for TPU v3, (1,1,1,2) is better
166
+ than (2,1,1,1). To pick a good spec, try a few possible values until you
167
+ get high TPU utilization.
168
+ input_devices: the devices to use, will use jax.devices() if this is not
169
+ set.
170
+ input_local_devices: the local devices to use, will use jax.local_devices()
171
+ if this is not set.
172
+ tile_by_host_if_needed: JAX currently requires that the parts of any sharded
173
+ array that are located on one host's local devices form a single
174
+ contiguous slice. A best effort will be made to achieve this without
175
+ "tiling" the device assignment over hosts (which can reduce XLA collective
176
+ performance). If this flag is True, then the device assignment will be
177
+ tiled over hosts if necessary to satisfy this constraint and create a
178
+ buildable mesh; if false, mesh construction will fail instead.
179
+ backend: get devices from the pinned backend, if specified. This is
180
+ useful for explicitly specifying the devices other than relying on
181
+ jax_platform_name.
182
+
183
+ Returns:
184
+ A xmap / pjit Mesh containing the virtual device mesh with data, model axes.
185
+ """
186
+ input_devices = input_devices or jax.devices(backend)
187
+ input_local_devices = input_local_devices or jax.local_devices(0, backend)
188
+ # Sort input_devices based on coords, as backends might not return devices
189
+ # in order.
190
+ last_device = sorted(input_devices, key=get_coords)[-1]
191
+ last_input_local_devices = sorted(input_local_devices, key=get_coords)[-1]
192
+ logging.info(
193
+ "last device coords : %r\nlast local device coords: %r",
194
+ get_coords(last_device),
195
+ get_coords(last_input_local_devices),
196
+ )
197
+ global_hardware_mesh = bounds_from_last_device(last_device)
198
+ mesh_ndim = len(global_hardware_mesh)
199
+ local_hardware_mesh = bounds_from_last_device(last_input_local_devices)
200
+ mesh_err = (
201
+ f"each dimension of the model parallel submesh {model_parallel_submesh} "
202
+ "must be a factor of the corresponding dimension of the global device "
203
+ f"mesh {global_hardware_mesh}"
204
+ )
205
+ assert not any(g % m for g, m in zip(global_hardware_mesh, model_parallel_submesh)), mesh_err
206
+ assert not any(g % l for g, l in zip(global_hardware_mesh, local_hardware_mesh))
207
+ devices = np.empty(global_hardware_mesh, dtype=object)
208
+ for device in input_devices:
209
+ device_coords = get_coords(device)
210
+ devices[device_coords] = device
211
+ tile_by_host = tile_by_host_if_needed
212
+ if len(global_hardware_mesh) == 4:
213
+ # enable contiguous local chunks without host tiling by making Z major
214
+ global_hardware_mesh = typing.cast(Tuple[int, int, int, int], global_hardware_mesh)
215
+ model_parallel_submesh = typing.cast(Tuple[int, int, int, int], model_parallel_submesh)
216
+ gx, gy, gz, gc = global_hardware_mesh
217
+ mx, my, mz, mc = model_parallel_submesh
218
+ if (mx == gx > 1 and my == mz == 1) or (mx == 1 and my == gy > 1 and mz == gz > 1):
219
+ logging.info("ensuring YZ plane has a Z-major device order")
220
+ # YZ should be ZY
221
+ assert mc == gc, (mc, gc)
222
+ global_hardware_mesh = gx, gz, gy, gc
223
+ model_parallel_submesh = mx, mz, my, mc
224
+ devices = devices.swapaxes(1, 2)
225
+ tile_by_host = False
226
+ if (my == gy > 1 and mx == mz == 1) or (my == 1 and mx == gx > 1 and mz == gz > 1):
227
+ logging.info("ensuring XZ plane has a Z-major device order")
228
+ # XZ should be ZX
229
+ assert mc == gc, (mc, gc)
230
+ global_hardware_mesh = gz, gy, gx, gc
231
+ model_parallel_submesh = mz, my, mx, mc
232
+ devices = devices.swapaxes(0, 2)
233
+ tile_by_host = False
234
+ if tile_by_host:
235
+ logging.warning(
236
+ "Tiling device assignment mesh by hosts, which may lead to "
237
+ "reduced XLA collective performance. To avoid this, modify "
238
+ "the model parallel submesh or run with more tasks per host."
239
+ )
240
+ tile_err = (
241
+ "to tile the mesh by hosts, each dimension of the model parallel "
242
+ "submesh must be either a factor or a multiple of the corresponding "
243
+ "dimension of the per-host submesh"
244
+ )
245
+
246
+ def dh_dd_mh_md(g: int, m: int, l: int) -> Tuple[int, int, int, int]:
247
+ """Split a global mesh dimension into four tiling components.
248
+
249
+ Args:
250
+ g: global mesh bounds dimension size
251
+ m: model-parallel submesh bounds dimension size
252
+ l: local submesh bounds dimension size
253
+
254
+ Returns:
255
+ The resulting tuple divides the dimension into the hosts component of
256
+ the data-parallel submesh, the devices component of the data-parallel
257
+ submesh, the hosts component of the model-parallel submesh, and the
258
+ devices component of the model-parallel submesh.
259
+ """
260
+ d = g // m
261
+ if m >= l:
262
+ assert not m % l, tile_err
263
+ return (d, 1, m // l, l)
264
+ else:
265
+ assert not l % m, tile_err
266
+ return (d // (l // m), l // m, 1, m)
267
+
268
+ # e.g. [(x_data_hosts, x_data_devs, x_model_hosts, x_model_devs), ...]
269
+ dh_dd_mh_md_tups = map(dh_dd_mh_md, global_hardware_mesh, model_parallel_submesh, local_hardware_mesh)
270
+ # reshape to e.g. (x_dh, x_dd, x_mh, x_md, y_dh, ...)
271
+ devices = devices.reshape(*(s for t in dh_dd_mh_md_tups for s in t)) # pylint: disable=g-complex-comprehension
272
+ # TODO(jekbradbury): reorder local subgroups for ring locality
273
+ # Transpose to [data_host], [data_device], [model_host], [model_device]
274
+ # block ordering e.g. (x_dh, y_dh, ..., x_dd, y_dd, ...)
275
+ devices = devices.transpose(
276
+ *(4 * i for i in range(mesh_ndim)),
277
+ *(4 * i + 1 for i in range(mesh_ndim)),
278
+ *(4 * i + 2 for i in range(mesh_ndim)),
279
+ *(4 * i + 3 for i in range(mesh_ndim)),
280
+ )
281
+ else:
282
+ # e.g. [(x_data, x_model), (y_data, y_model), ...]
283
+ model_data_tups = [(g // m, m) for g, m in zip(global_hardware_mesh, model_parallel_submesh)]
284
+ # reshape to e.g. (x_data, x_model, y_data, y_model...)
285
+ devices = devices.reshape(*(s for t in model_data_tups for s in t)) # pylint: disable=g-complex-comprehension
286
+ # TODO(jekbradbury): reorder small subgroups for ring locality
287
+ # transpose to e.g. (x_data, y_data, ..., x_model, ...)
288
+ devices = devices.transpose(*(2 * i for i in range(mesh_ndim)), *(2 * i + 1 for i in range(mesh_ndim)))
289
+ # reshape to (data, model)
290
+ devices = devices.reshape(-1, np.prod(model_parallel_submesh))
291
+ global_mesh = Mesh(devices, ["data", "model"])
292
+ logging.info("global_mesh axis_names: %s", global_mesh.axis_names)
293
+ logging.info("global_mesh devices: %s", global_mesh.devices)
294
+ logging.info("global_mesh devices shape: %s", global_mesh.devices.shape)
295
+ return global_mesh
296
+
297
+
298
+ def get_cpu_mesh() -> Mesh:
299
+ """Trivial mesh for CPU Testing."""
300
+ devices = np.empty((jax.host_count(), jax.local_device_count()), dtype=object)
301
+ for device in jax.devices():
302
+ devices[device.process_index, device.id % jax.local_device_count()] = device
303
+ return Mesh(devices, ["data", "model"])
304
+
305
+
306
+ def get_gpu_mesh(num_partitions: int) -> Mesh:
307
+ """Mesh for GPUs that preferentially places 'model' on NVLink."""
308
+ nvlink_size = jax.local_device_count()
309
+ dcn_size = jax.process_count()
310
+ nvlink_mp = min(num_partitions, nvlink_size)
311
+ nvlink_dp, extra1 = divmod(nvlink_size, nvlink_mp)
312
+ dcn_mp, extra2 = divmod(num_partitions, nvlink_mp)
313
+ assert not (extra1 or extra2), (
314
+ "number of partitions on GPU must be a factor" " or multiple of the number of local devices"
315
+ )
316
+ dcn_dp = dcn_size // dcn_mp
317
+
318
+ devices = create_hybrid_device_mesh(
319
+ mesh_shape=[nvlink_dp, nvlink_mp], dcn_mesh_shape=[dcn_dp, dcn_mp], process_is_granule=True
320
+ )
321
+
322
+ global_mesh = Mesh(devices, ["data", "model"])
323
+ logging.info("global_mesh axis_names: %s", global_mesh.axis_names)
324
+ logging.info("global_mesh devices: %s", global_mesh.devices)
325
+ return global_mesh
326
+
327
+
328
+ def default_mesh(
329
+ num_partitions: int, model_parallel_submesh: Optional[HardwareMesh] = None, backend: Optional[str] = None
330
+ ) -> Mesh:
331
+ """Attempt to return a default mesh for simple cases.
332
+
333
+ Args:
334
+ num_partitions: number of partitions to use, will be ignored if
335
+ model_parallel_submesh is provided.
336
+ model_parallel_submesh: 4-tuple that specifies the x,y,z,c submesh to use as
337
+ the model-parallel device tile.
338
+ backend: get devices from the pinned backend, if specified. This is useful
339
+ for explicitly specifying the devices other than relying on
340
+ jax_platform_name.
341
+
342
+ Returns:
343
+ xmap/pjit 2D Mesh with 'data', 'model' mesh axes.
344
+ """
345
+ last_device = jax.devices(backend)[-1]
346
+ platform = last_device.platform
347
+ device_kind = last_device.device_kind
348
+ bounds = bounds_from_last_device(last_device)
349
+
350
+ if model_parallel_submesh:
351
+ return get_mesh(model_parallel_submesh, backend=backend)
352
+
353
+ if platform == "cpu":
354
+ return get_cpu_mesh()
355
+ elif platform == "gpu":
356
+ return get_gpu_mesh(num_partitions)
357
+
358
+ mps = None
359
+ if device_kind in ("TPU v2", "TPU v3"):
360
+ if num_partitions == 1:
361
+ mps = (1, 1, 1, 1)
362
+ elif num_partitions == 2:
363
+ mps = (1, 1, 1, 2)
364
+ elif num_partitions == 4:
365
+ mps = (2, 1, 1, 2)
366
+ elif num_partitions == 8:
367
+ mps = (2, 2, 1, 2)
368
+ elif num_partitions == 16:
369
+ mps = (4, 2, 1, 2)
370
+ # assume the use of megacore on TPU v4
371
+ elif (device_kind == "TPU v4" or device_kind == "TPU v4 lite") and bounds[3] == 1:
372
+ if num_partitions == 1:
373
+ mps = (1, 1, 1, 1)
374
+ elif num_partitions == 2:
375
+ mps = (1, 2, 1, 1)
376
+ elif num_partitions == 4:
377
+ if bounds[0] >= 4:
378
+ mps = (4, 1, 1, 1)
379
+ else:
380
+ mps = (2, 2, 1, 1)
381
+ elif num_partitions == 8:
382
+ if bounds[2] >= 8:
383
+ mps = (1, 1, 8, 1)
384
+ else:
385
+ mps = (4, 2, 1, 1)
386
+ elif num_partitions == 16:
387
+ if bounds[2] >= 16:
388
+ mps = (1, 1, 16, 1)
389
+ elif bounds[0] >= 8:
390
+ mps = (8, 2, 1, 1)
391
+ elif bounds[0] >= 4:
392
+ mps = (4, 4, 1, 1)
393
+ else:
394
+ mps = (2, 2, 4, 1)
395
+
396
+ if mps is None:
397
+ raise ValueError(
398
+ "No default mesh for this configuration: specify " "config.model_parallel_submesh explicitly."
399
+ )
400
+ return get_mesh(mps, backend=backend)
401
+
402
+
403
+ # Data chunking helper.
404
+ # -----------------------------------------------------------------------------
405
+ @dataclasses.dataclass
406
+ class LocalChunkInfo:
407
+ # The logical slice of an array located on this host's local devices.
408
+ slice: Tuple[slice, ...]
409
+ # A unique index for this host/local chunk among chunks with the same slice.
410
+ replica_id: int
411
+
412
+
413
+ class LocalChunker:
414
+ """Utility class to aid chunking of sharded arrays in multihost settings."""
415
+
416
+ def __init__(self, global_mesh: Mesh):
417
+ self.global_mesh = global_mesh
418
+ local_mesh = global_mesh.local_mesh
419
+ first_local_device = local_mesh.devices.reshape(-1)[0]
420
+ host_location = collections.OrderedDict(
421
+ zip(global_mesh.shape.keys(), list(zip(*np.nonzero(global_mesh.devices == first_local_device)))[0])
422
+ )
423
+ self.num_chunks = collections.OrderedDict()
424
+ self.chunk_ids = collections.OrderedDict()
425
+ self.mesh_axes = list(global_mesh.shape.keys())
426
+ for mesh_axis in self.mesh_axes:
427
+ num_devices_per_chunk = local_mesh.shape[mesh_axis]
428
+ self.num_chunks[mesh_axis] = global_mesh.shape[mesh_axis] // num_devices_per_chunk
429
+ self.chunk_ids[mesh_axis] = host_location[mesh_axis] // num_devices_per_chunk
430
+
431
+ def get_local_chunk_info(
432
+ self, global_shape: Tuple[int, ...], mesh_axes: Sequence[Optional[str]]
433
+ ) -> LocalChunkInfo:
434
+ """Get the local chunk info for a given array shape and sharded axes.
435
+
436
+ Args:
437
+ global_shape: the global, unsharded shape of the array to chunk.
438
+ mesh_axes: a sequence of names (or None) of equal rank to `global_shape`
439
+ that specifies which mesh dimensions the array is sharded along.
440
+
441
+ Returns:
442
+ LocalChunkInfo containing the logical slices of the array found on this
443
+ host's local devices, as well as the replica index for this chunk among
444
+ chunks with the same slice. The latter is used to determine which
445
+ host should write this chunk during checkpointing.
446
+ """
447
+ local_slice = [slice(None) for dim in global_shape]
448
+ sharded_mesh_axes = set()
449
+ for i, (mesh_axis, size) in enumerate(zip(mesh_axes, global_shape)):
450
+ if not mesh_axis:
451
+ continue
452
+ sharded_mesh_axes.add(mesh_axis)
453
+ if not isinstance(mesh_axis, str):
454
+ raise NotImplementedError("TODO(jekbradbury)")
455
+ chunk_id = self.chunk_ids[mesh_axis]
456
+ chunk_size = size // self.num_chunks[mesh_axis]
457
+ local_slice[i] = slice(chunk_id * chunk_size, (chunk_id + 1) * chunk_size)
458
+
459
+ replicated_mesh_axes = [mesh_axis for mesh_axis in self.mesh_axes if mesh_axis not in sharded_mesh_axes]
460
+ replica_id = 0
461
+ for mesh_axis in replicated_mesh_axes:
462
+ chunk_id = self.chunk_ids[mesh_axis]
463
+ replica_id = replica_id * self.num_chunks[mesh_axis] + chunk_id
464
+
465
+ return LocalChunkInfo(tuple(local_slice), replica_id)
466
+
467
+
468
+ def standard_logical_axis_rules(
469
+ activation_partitioning_dims: int = 1,
470
+ parameter_partitioning_dims: int = 1,
471
+ additional_rules: Optional[LogicalAxisRules] = None,
472
+ ) -> LogicalAxisRules:
473
+ """Default sharding rules for T5X model in terms of logical axis names.
474
+
475
+ Args:
476
+ activation_partitioning_dims: enables 2-D activation sharding when set to 2.
477
+ parameter_partitioning_dims: enables 2-D parameter sharding when set to 2.
478
+ additional_rules: additional rules (a sequence of tuples) that will be
479
+ appended to the standard rules.
480
+
481
+ Returns:
482
+ Sequence of logical axis rules
483
+ """
484
+ logging.info(
485
+ "`activation_partitioning_dims` = %d, `parameter_partitioning_dims` = %d",
486
+ activation_partitioning_dims,
487
+ parameter_partitioning_dims,
488
+ )
489
+
490
+ if activation_partitioning_dims == 1 and parameter_partitioning_dims == 1:
491
+ rules = [
492
+ ("batch", "data"),
493
+ ("vocab", "model"),
494
+ ("embed", None),
495
+ ("mlp", "model"),
496
+ ("heads", "model"),
497
+ ("kv", None),
498
+ ("joined_kv", "model"), # joined heads+kv dim in 2D attn param layouts
499
+ ]
500
+ elif activation_partitioning_dims == 2 and parameter_partitioning_dims == 1:
501
+ rules = [
502
+ ("batch", "data"),
503
+ ("vocab", "model"),
504
+ ("mlp", "model"),
505
+ ("heads", "model"),
506
+ ("kv", None),
507
+ ("joined_kv", "model"),
508
+ ("embed", "model"),
509
+ ]
510
+ elif activation_partitioning_dims == 1 and parameter_partitioning_dims == 2:
511
+ rules = [
512
+ ("batch", "data"),
513
+ ("vocab", "model"),
514
+ ("mlp", "model"),
515
+ ("heads", "model"),
516
+ ("kv", None),
517
+ ("joined_kv", "model"),
518
+ ("embed", "data"),
519
+ ]
520
+ elif activation_partitioning_dims == 2 and parameter_partitioning_dims == 2:
521
+ rules = [
522
+ ("batch", "data"),
523
+ ("vocab", "model"),
524
+ ("mlp", "model"),
525
+ ("heads", "model"),
526
+ ("kv", None),
527
+ ("joined_kv", "model"),
528
+ ("embed", "model"),
529
+ ("embed", "data"),
530
+ ]
531
+ else:
532
+ raise ValueError(
533
+ f"`activation_partitioning_dims` = {activation_partitioning_dims} "
534
+ f"`parameter_partitioning_dims` = {parameter_partitioning_dims} "
535
+ "is not supported."
536
+ )
537
+
538
+ # Add the common rules for the replicated logical axes names.
539
+ replicated_rules = [
540
+ ("relpos_buckets", None),
541
+ ("abspos_buckets", None),
542
+ ("length", None),
543
+ ("layers", None),
544
+ ("stack", None),
545
+ ("mlp_activations", None),
546
+ ]
547
+ rules.extend(replicated_rules)
548
+
549
+ if additional_rules:
550
+ rules.extend(additional_rules)
551
+
552
+ return rules
553
+
554
+
555
+ # NB: This needs to be top-level for the jax compilation cache.
556
+ def _id_fn(x, ix):
557
+ """Identity function for copying parameters to the devices, sharded."""
558
+ # A pure identity such as `lambda x, *: x` can get optimized away, so we
559
+ # include a random.split as a cheap function that cannot be optimized away.
560
+ y = random.split(random.PRNGKey(jnp.array(ix, dtype=jnp.uint32)))
561
+ return x, y
562
+
563
+
564
+ @dataclasses.dataclass
565
+ class DataLayout:
566
+ """Represents data layout for the partitioned model."""
567
+
568
+ batch_size: int
569
+ shard_id: int
570
+ num_shards: int
571
+ is_first_host_in_replica_set: bool
572
+
573
+
574
+ PartitionedCallable = Callable[..., Any]
575
+ CompiledPartitionedCallable = Callable[..., Any]
576
+
577
+
578
+ class BasePartitioner(metaclass=abc.ABCMeta):
579
+ """Interface for partitioning computations across hardware devices."""
580
+
581
+ def __init__(
582
+ self,
583
+ num_partitions: Optional[int] = None,
584
+ model_parallel_submesh: Optional[HardwareMesh] = None,
585
+ params_on_devices: bool = True,
586
+ backend: Optional[str] = None,
587
+ ):
588
+ """Configures the partitioner.
589
+
590
+ Args:
591
+ num_partitions: the number of partitions to use. Ignored if
592
+ `model_parallel_submesh` is provided.
593
+ model_parallel_submesh: 4-tuple that specifies the x,y,z,c submesh to use
594
+ as the model-parallel device tile. This submesh is used for the larger
595
+ of the two parameter dimensions, and, if 2-D activation sharding is
596
+ enabled, for the model dimension of activations. The rest of the mesh is
597
+ used for data parallelism and, if 2-D parameter sharding is enabled, the
598
+ other parameter dimension.
599
+ params_on_devices: whether to keep the params on devices, if False -
600
+ params stay in the host memory. Note that some partitioners might ignore
601
+ this setting, for example if they don't support storing all params on
602
+ device memory.
603
+ backend: get devices from the pinned backend, if specified. This is useful
604
+ for explicitly specifying the devices other than relying on
605
+ jax_platform_name.
606
+ """
607
+
608
+ if not num_partitions and not model_parallel_submesh:
609
+ raise ValueError("At least one of `num_partitions` or " "`model_parallel_submesh` must be set.")
610
+
611
+ if model_parallel_submesh is not None and len(model_parallel_submesh) != 4:
612
+ logging.error(
613
+ (
614
+ "`model_parallel_submesh` must be either None or a 4-tuple. Got"
615
+ " `model_parallel_submesh`=%s. A ValueError will be raised"
616
+ " beginning March 1, 2022."
617
+ ),
618
+ model_parallel_submesh,
619
+ )
620
+
621
+ if bool(num_partitions) and bool(model_parallel_submesh):
622
+ logging.error(
623
+ "At most one of `num_partitions` or `model_parallel_submesh` can be "
624
+ "set. Got `num_partitions=%s` and `model_parallel_submesh`=%s. A "
625
+ "ValueError will be raised beginning March 21, 2022.",
626
+ num_partitions,
627
+ model_parallel_submesh,
628
+ )
629
+
630
+ self._num_partitions = num_partitions
631
+ self._model_parallel_submesh = model_parallel_submesh
632
+ self._params_on_devices = params_on_devices
633
+ self._data_axis = "data"
634
+ self._backend = backend
635
+
636
+ @property
637
+ def mesh(self) -> Mesh:
638
+ raise NotImplementedError
639
+
640
+ @property
641
+ def data_partition_spec(self) -> PartitionSpec:
642
+ return PartitionSpec(self._data_axis)
643
+
644
+ def get_data_layout(self, batch_size: Optional[int] = None, host_index: Optional[int] = None) -> DataLayout:
645
+ """Returns filled `DataLayout` based on the partitioned model layout.
646
+
647
+ Args:
648
+ batch_size: if set, indicates the requested batch size. The exception will
649
+ be raised if this batch size is not compatible with the layout. If not
650
+ set, the batch size is inferred from the layout.
651
+ host_index: indicates the host index to use for the calculations, if not
652
+ set - use JAX-provided one. Should be in [0, num_hosts) interval and the
653
+ order should match the order of corresponding CPU devices in
654
+ `jax.devices()`.
655
+
656
+ Returns:
657
+ Filled `DataLayout` structure.
658
+ """
659
+ if host_index is not None:
660
+ raise NotImplementedError("Explicit host_index is not yet implemented.")
661
+ if self._data_axis is None:
662
+ return DataLayout(
663
+ batch_size=batch_size,
664
+ shard_id=0,
665
+ num_shards=1,
666
+ is_first_host_in_replica_set=(jax.process_index() == 0),
667
+ )
668
+ mesh_size = self._local_chunker.global_mesh.shape[self._data_axis]
669
+ batch_size = batch_size or mesh_size
670
+ if batch_size % mesh_size:
671
+ raise ValueError(
672
+ f"Batch size ({batch_size}) must be divisible by corresponding " f"mesh size ({mesh_size})."
673
+ )
674
+ num_shards = self._local_chunker.num_chunks[self._data_axis]
675
+ if batch_size % num_shards:
676
+ raise ValueError(f"Batch size ({batch_size}) must be divisible by number of " f"replicas ({num_shards}).")
677
+ replica_id = self._local_chunker.get_local_chunk_info((batch_size,), [self._data_axis]).replica_id
678
+ return DataLayout(
679
+ batch_size=int(batch_size),
680
+ shard_id=int(self._local_chunker.chunk_ids[self._data_axis]),
681
+ num_shards=int(num_shards),
682
+ is_first_host_in_replica_set=(replica_id == 0),
683
+ )
684
+
685
+ def get_local_chunk_info(
686
+ self, global_shape: Tuple[int, ...], mesh_axes: Sequence[Optional[str]]
687
+ ) -> LocalChunkInfo:
688
+ """Returns the local chunk info for a given array shape and sharded axes."""
689
+ return self._local_chunker.get_local_chunk_info(global_shape, mesh_axes)
690
+
691
+ @property
692
+ def params_on_devices(self):
693
+ return self._params_on_devices
694
+
695
+ def move_params_to_devices(self, train_state: TrainState, train_state_axes: TrainState) -> TrainState:
696
+ """Moves the optimizer parameters to devices."""
697
+ p_id_fn = self.partition(
698
+ _id_fn,
699
+ in_axis_resources=(train_state_axes, None),
700
+ out_axis_resources=(train_state_axes, None),
701
+ donate_argnums=(0,),
702
+ )
703
+ if jax.config.jax_array and jax.process_count() > 1:
704
+ train_state = multihost_utils.host_local_array_to_global_array(train_state, self.mesh, train_state_axes)
705
+ train_state, _ = p_id_fn(train_state, jnp.ones((), dtype=jnp.uint32))
706
+ return train_state
707
+
708
+ @property
709
+ @abc.abstractmethod
710
+ def _local_chunker(self):
711
+ """Returns the chunker that matches the parameters of this partitioner."""
712
+ raise NotImplementedError
713
+
714
+ def get_logical_axes(self, train_state: TrainState) -> TrainState:
715
+ """Returns a copy of TrainState with Optional[AxisNames] as leaves."""
716
+ # By default, return None for the logical axes.
717
+ return train_state.restore_state(jax.tree_map(lambda x: None, train_state.state_dict()))
718
+
719
+ def get_mesh_axes(self, train_state: TrainState) -> TrainState:
720
+ """Returns a copy of TrainState with Optional[PartitionSpecs] as leaves."""
721
+ raise NotImplementedError
722
+
723
+ @abc.abstractmethod
724
+ def partition(
725
+ self,
726
+ fn: Callable, # pylint: disable=g-bare-generic
727
+ in_axis_resources,
728
+ out_axis_resources,
729
+ static_argnums: Union[int, Sequence[int]] = (),
730
+ donate_argnums: Union[int, Sequence[int]] = (),
731
+ ) -> PartitionedCallable:
732
+ """Partitions the computation using partitioner-specific implementation.
733
+
734
+ Args:
735
+ fn: the function to partition.
736
+ in_axis_resources: Pytree of structure matching that of arguments to `fn`,
737
+ with all actual arguments replaced by resource assignment
738
+ specifications. It is also valid to specify a pytree prefix (e.g. one
739
+ value in place of a whole subtree), in which case the leaves get
740
+ broadcast to all values in that subtree.
741
+ The valid resource assignment specifications are:
742
+ `None`: in which case the value will be replicated on all devices
743
+ `PartitionSpec`: a tuple of length at most equal to the rank of the
744
+ partitioned value. Each element can be a `None`, a mesh axis or a
745
+ tuple of mesh axes, and specifies the set of resources assigned to
746
+ partition the value's dimension matching its position in the spec.
747
+ out_axis_resources: Like `in_axis_resources`, but specifies resource
748
+ assignment for function outputs.
749
+ static_argnums: an optional int or collection of ints that specify which
750
+ positional arguments to treat as static (compile-time constant) in the
751
+ partitioned function.
752
+ donate_argnums: an optional int or collection of ints that specify which
753
+ argument buffers are "donated" to the computation. It is safe to donate
754
+ argument buffers if you no longer need them once the computation has
755
+ finished.
756
+
757
+ Returns:
758
+ A partitioned version of the input function.
759
+ """
760
+ raise NotImplementedError
761
+
762
+ @abc.abstractmethod
763
+ def compile(self, partitioned_fn: PartitionedCallable, *args) -> CompiledPartitionedCallable:
764
+ """Compiles and returns the partitioned function, or the original.
765
+
766
+ Args:
767
+ partitioned_fn: The partitioned function.
768
+ *args: Sample arguments to the partitioned function matching the input
769
+ shapes that will be passed to the compiled function.
770
+
771
+ Returns:
772
+ The compiled function, or the original if this partitioner does not
773
+ support compilation.
774
+ """
775
+ raise NotImplementedError
776
+
777
+
778
+ class PjittedFnWithContext(PartitionedCallable):
779
+ """Wraps pjitted function to apply the appropriate contexts."""
780
+
781
+ def __init__(self, pjitted_fn, partition_mesh: Mesh, logical_axis_rules: flax_partitioning.LogicalRules = ()):
782
+ self._pjitted_fn = pjitted_fn
783
+ self._mesh = partition_mesh
784
+ self._logical_axis_rules = logical_axis_rules
785
+
786
+ def __call__(self, *args):
787
+ with Mesh(self._mesh.devices, self._mesh.axis_names), flax_partitioning.axis_rules(self._logical_axis_rules):
788
+ return self._pjitted_fn(*args)
789
+
790
+ def lower(self, *args):
791
+ with Mesh(self._mesh.devices, self._mesh.axis_names), flax_partitioning.axis_rules(self._logical_axis_rules):
792
+ return self._pjitted_fn.lower(*args)
793
+
794
+
795
+ class BasePjitPartitioner(BasePartitioner):
796
+ """Partitioner that uses T5X version of jax.pjit."""
797
+
798
+ @cached_property
799
+ def _local_chunker(self) -> LocalChunker:
800
+ return LocalChunker(self.mesh)
801
+
802
+ @cached_property
803
+ def mesh(self) -> Mesh:
804
+ return default_mesh(self._num_partitions, self._model_parallel_submesh, self._backend)
805
+
806
+ def partition(
807
+ self,
808
+ fn: Callable, # pylint: disable=g-bare-generic
809
+ in_axis_resources,
810
+ out_axis_resources,
811
+ static_argnums: Union[int, Sequence[int]] = (),
812
+ donate_argnums: Union[int, Sequence[int]] = (),
813
+ ) -> PjittedFnWithContext:
814
+ pjitted = pjit(
815
+ fn,
816
+ in_axis_resources=in_axis_resources,
817
+ out_axis_resources=out_axis_resources,
818
+ static_argnums=static_argnums,
819
+ donate_argnums=donate_argnums,
820
+ backend=self._backend,
821
+ )
822
+
823
+ return PjittedFnWithContext(pjitted, self.mesh)
824
+
825
+ def compile(self, partitioned_fn: PjittedFnWithContext, *args) -> CompiledPartitionedCallable:
826
+ return partitioned_fn.lower(*args).compile()
827
+
828
+
829
+ class PjitPartitioner(BasePjitPartitioner):
830
+ """Partitioner that uses named axes and jax.pjit."""
831
+
832
+ def __init__(
833
+ self,
834
+ num_partitions: Optional[int] = None,
835
+ model_parallel_submesh: Optional[HardwareMesh] = None,
836
+ params_on_devices: bool = True,
837
+ backend: Optional[str] = None,
838
+ logical_axis_rules: Optional[LogicalAxisRules] = None,
839
+ use_cpu_pjit: Optional[bool] = False,
840
+ ):
841
+ """PjitPartitioner constructor.
842
+
843
+ See https://github.com/google-research/text-to-text-transfer-transformer/blob/main/README.mdx/usage/partitioning for details.
844
+
845
+ Args:
846
+ num_partitions: an integer that specifies the size of the model parallel
847
+ submesh to be automatically selected for the current topology. See
848
+ `model_parallel_submesh` for details on how this submesh is used.
849
+ Mutually exlusive with `model_parallel_submesh`.
850
+ model_parallel_submesh: is a 4-tuple that specifies the `(x, y, z, c)`
851
+ submesh model-parallel device tile, an axis of accelerator parallelism
852
+ orthogonal to data parallelism. Array axes in a model's parameters or
853
+ activations can be sharded over this submesh using axis rules (see
854
+ `logical_axis_rules`) that map them to 'model'. The effective number of
855
+ model sub-partitions is equal to `np.prod(model_parallel_submesh)` and
856
+ must evenly divide the total number of devices (i.e.,
857
+ `jax.device_count() % np.prod(model_parallel_submesh) == 0`). The rest
858
+ of the TPU mesh is the data parallel submesh, providing
859
+ `jax.device_count() // np.prod(model_parallel_submesh)` partitions. It
860
+ is used for data (batch) parallelism and to shard other array axes that
861
+ are mapped to 'data'. This argument is mutually exclusive with
862
+ `num_partitions`.
863
+ params_on_devices: whether to keep the params on devices, if False -
864
+ params stay in the host memory. Note that some partitioners might ignore
865
+ this setting, for example if they don't support storing all params on
866
+ device memory.
867
+ backend: get devices from the pinned backend, if specified. This is
868
+ useful for explicitly specifying the devices other than relying on
869
+ jax_platform_name.
870
+ logical_axis_rules: a priority-ordered sequence of KV tuples that maps
871
+ logical axis names to either `None` (not sharded), 'model' (to shard
872
+ across the model-parallel submesh), or 'data' (to shard across the
873
+ data-parallel submesh).
874
+ use_cpu_pjit: enables wrapper function for pjit which just jits the
875
+ function if using CPU backend.
876
+ """
877
+ super().__init__(
878
+ num_partitions=num_partitions,
879
+ model_parallel_submesh=model_parallel_submesh,
880
+ params_on_devices=params_on_devices,
881
+ backend=backend,
882
+ )
883
+ if logical_axis_rules is None:
884
+ logical_axis_rules = standard_logical_axis_rules()
885
+ self._logical_axis_rules = tuple(logical_axis_rules)
886
+ (self._data_axis,) = flax_partitioning.logical_to_mesh_axes(["batch"], logical_axis_rules)
887
+ self._use_cpu_pjit = use_cpu_pjit
888
+
889
+ def partition(
890
+ self,
891
+ fn: Callable, # pylint: disable=g-bare-generic
892
+ in_axis_resources,
893
+ out_axis_resources,
894
+ static_argnums: Union[int, Sequence[int]] = (),
895
+ donate_argnums: Union[int, Sequence[int]] = (),
896
+ ) -> PjittedFnWithContext:
897
+ """Partitions the function using jax.pjit."""
898
+ if self._use_cpu_pjit:
899
+ pjit_fn = pjit_with_cpu_fallback
900
+ else:
901
+ pjit_fn = pjit
902
+ pjitted = pjit_fn(
903
+ fn,
904
+ in_axis_resources=in_axis_resources,
905
+ out_axis_resources=out_axis_resources,
906
+ static_argnums=static_argnums,
907
+ donate_argnums=donate_argnums,
908
+ backend=self._backend,
909
+ )
910
+
911
+ return PjittedFnWithContext(pjitted, self.mesh, self._logical_axis_rules)
912
+
913
+ @property
914
+ def logical_axis_rules(self):
915
+ """Returns the logical axis rules."""
916
+ return self._logical_axis_rules
917
+
918
+ def get_logical_axes(self, train_state: TrainState) -> TrainState:
919
+ """Returns a copy of TrainState with Optional[AxisNames] as leaves."""
920
+ return train_state.as_logical_axes()
921
+
922
+ def get_mesh_axes(self, train_state: TrainState) -> TrainState:
923
+ """Returns a copy of TrainState with Optional[PartitionSpecs] as leaves."""
924
+ logical_axes = self.get_logical_axes(train_state)
925
+
926
+ def _logical_to_mesh_axes(param_name, logical_axes):
927
+ if logical_axes is None:
928
+ return None
929
+ elif logical_axes is traverse_util.empty_node:
930
+ return traverse_util.empty_node
931
+ try:
932
+ return flax_partitioning.logical_to_mesh_axes(logical_axes, self._logical_axis_rules)
933
+ except ValueError as e:
934
+ raise ValueError(f"Failed to map logical axes for {param_name}") from e
935
+
936
+ flat_logical_axes = traverse_util.flatten_dict(logical_axes.state_dict(), keep_empty_nodes=True, sep="/")
937
+ flat_mesh_axes = {k: _logical_to_mesh_axes(k, v) for k, v in flat_logical_axes.items()}
938
+
939
+ return logical_axes.restore_state(traverse_util.unflatten_dict(flat_mesh_axes, sep="/"))
whisper_jax/pipeline.py ADDED
@@ -0,0 +1,506 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # coding=utf-8
2
+ # Copyright 2023 The HuggingFace Inc. team.
3
+ #
4
+ # Licensed under the Apache License, Version 2.0 (the "License");
5
+ # you may not use this file except in compliance with the License.
6
+ # You may obtain a copy of the License at
7
+ #
8
+ # http://www.apache.org/licenses/LICENSE-2.0
9
+ #
10
+ # Unless required by applicable law or agreed to in writing, software
11
+ # distributed under the License is distributed on an "AS IS" BASIS,
12
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
+ # See the License for the specific language governing permissions and
14
+ # limitations under the License.
15
+
16
+
17
+ import math
18
+
19
+ import jax
20
+ import jax.numpy as jnp
21
+ import numpy as np
22
+ import requests
23
+ from flax import jax_utils
24
+ from flax.core.frozen_dict import freeze
25
+ from flax.training.common_utils import shard
26
+ from jax.sharding import PartitionSpec as P
27
+ from transformers import WhisperProcessor, is_tokenizers_available, WhisperFeatureExtractor, WhisperTokenizerFast
28
+ from transformers.models.whisper.tokenization_whisper import TO_LANGUAGE_CODE, WhisperTokenizer
29
+ from transformers.pipelines.audio_utils import ffmpeg_read
30
+ from transformers.utils import logging
31
+
32
+ from .modeling_flax_whisper import FlaxWhisperForConditionalGeneration
33
+ from .partitioner import PjitPartitioner
34
+ from .train_state import InferenceState
35
+
36
+
37
+ logger = logging.get_logger(__name__)
38
+
39
+ # 2D parameter and activation partitioning for DP
40
+ logical_axis_rules_dp = (
41
+ ("batch", "data"),
42
+ ("mlp", None),
43
+ ("heads", None),
44
+ ("vocab", None),
45
+ ("embed", None),
46
+ ("embed", None),
47
+ ("joined_kv", None),
48
+ ("kv", None),
49
+ ("length", None),
50
+ ("num_mel", None),
51
+ ("channels", None),
52
+ )
53
+
54
+
55
+ class FlaxWhisperPipline:
56
+ def __init__(
57
+ self,
58
+ checkpoint="openai/whisper-large-v2",
59
+ dtype=jnp.float32,
60
+ batch_size=None,
61
+ max_length=None,
62
+ ):
63
+ """
64
+ Args
65
+ checkpoint (`str`, *optional*, defaults to `"openai/whisper-large-v2"):
66
+ The Whisper checkpoint to use with the pipeline. Must be an available checkpoint on the Hugging Face Hub
67
+ with Flax weights.
68
+ dtype (`jax.numpy.dtype`, *optional*, defaults to `jax.numpy.float32`):
69
+ The data type of the computation. Can be one of `jax.numpy.float32`, `jax.numpy.float16` (on GPUs) and
70
+ `jax.numpy.bfloat16` (on TPUs). This can be used to enable half-precision inference on GPUs or TPUs.
71
+ If specified all the computation will be performed with the given `dtype`. **Note that this only
72
+ specifies the dtype of the computation and does not influence the dtype of model parameters.**
73
+ batch_size (`int`, *optional*, defaults to the minimum per-device batch size, i.e. `jax.local_device_count()`):
74
+ The batch size to be used in chunking transcription. Beneficial for transcribing long audio files. Passing
75
+ a batch size in the `__init__` method will be superseded by any batch size passed to the `__call__` method.
76
+ max_length (`int`, *optional*):
77
+ The maximum numbers of tokens to generate. Defaults to `model.config.max_length`.
78
+ """
79
+ self.checkpoint = checkpoint
80
+ self.dtype = dtype
81
+
82
+ self.processor = WhisperProcessor.from_pretrained(self.checkpoint)
83
+ self.feature_extractor = self.processor.feature_extractor
84
+ # potentially load fast tokenizer if available
85
+ tokenizer_cls = WhisperTokenizerFast if is_tokenizers_available() else WhisperTokenizer
86
+ self.tokenizer = tokenizer_cls.from_pretrained(checkpoint)
87
+
88
+ self.model, self.params = FlaxWhisperForConditionalGeneration.from_pretrained(
89
+ self.checkpoint,
90
+ _do_init=False,
91
+ dtype=self.dtype,
92
+ )
93
+
94
+ self.max_length = max_length if max_length is not None else self.model.generation_config.max_length
95
+ self.min_batch_size = jax.local_device_count()
96
+ self.batch_size = (
97
+ batch_size if batch_size is not None else self.min_batch_size
98
+ ) # we need a minimum of 1 batch per-device
99
+
100
+ def generate(params, input_features, forced_decoder_ids, return_timestamps):
101
+ output_ids = self.model.pipeline_generate(
102
+ input_features,
103
+ params=params,
104
+ forced_decoder_ids=forced_decoder_ids,
105
+ return_timestamps=return_timestamps,
106
+ max_length=self.max_length,
107
+ )
108
+ return output_ids
109
+
110
+ # use pmap for DP by default - this is compatible on a Colab TPU v2
111
+ self.params = jax_utils.replicate(self.params)
112
+ self.p_generate = jax.pmap(
113
+ generate, "input_features", in_axes=(0, 0, None), out_axes=0, static_broadcasted_argnums=(3,)
114
+ )
115
+ self.is_sharded = False
116
+
117
+ def shard_params(self, num_mp_partitions=1, logical_axis_rules=logical_axis_rules_dp):
118
+ def init_fn():
119
+ input_shape = (1, self.model.config.num_mel_bins, 2 * self.model.config.max_source_positions)
120
+
121
+ input_features = jnp.zeros(input_shape, dtype="f4")
122
+ input_features = input_features.at[(..., -1)].set(self.model.config.eos_token_id)
123
+
124
+ decoder_input_ids = jnp.zeros((input_shape[0], 1), dtype="i4")
125
+ decoder_attention_mask = jnp.ones_like(decoder_input_ids)
126
+
127
+ batch_size, sequence_length = decoder_input_ids.shape
128
+ decoder_position_ids = jnp.broadcast_to(
129
+ jnp.arange(sequence_length)[None, :], (batch_size, sequence_length)
130
+ )
131
+
132
+ rng = jax.random.PRNGKey(0)
133
+ init_params = self.model.module.init(
134
+ rng,
135
+ input_features=input_features,
136
+ decoder_input_ids=decoder_input_ids,
137
+ decoder_attention_mask=decoder_attention_mask,
138
+ decoder_position_ids=decoder_position_ids,
139
+ return_dict=False,
140
+ )
141
+ return init_params
142
+
143
+ # Axis names metadata
144
+ param_axes = jax.eval_shape(init_fn)["params_axes"]
145
+
146
+ # Create InferenceState, since the partitioner expects it
147
+ state = InferenceState(
148
+ step=jnp.array(0),
149
+ params=freeze(self.model.params_shape_tree),
150
+ params_axes=freeze(param_axes),
151
+ flax_mutables=None,
152
+ flax_mutables_axes=param_axes,
153
+ )
154
+
155
+ partitioner = PjitPartitioner(num_partitions=num_mp_partitions, logical_axis_rules=logical_axis_rules)
156
+
157
+ mesh_axes = partitioner.get_mesh_axes(state)
158
+ params_spec = mesh_axes.params
159
+
160
+ p_shard_params = partitioner.partition(self.model.to_bf16, (params_spec,), params_spec)
161
+
162
+ # This will auto-magically run in mesh context
163
+ self.params = p_shard_params(freeze(jax_utils.unreplicate(self.params)))
164
+ self.is_sharded = True
165
+
166
+ def generate(params, input_features, forced_decoder_ids, return_timestamps):
167
+ output_ids = self.model.pipeline_generate(
168
+ input_features,
169
+ params=params,
170
+ forced_decoder_ids=forced_decoder_ids,
171
+ return_timestamps=return_timestamps,
172
+ max_length=self.max_length,
173
+ )
174
+ return output_ids
175
+
176
+ # Use pjit for generate only once we've sharded the params
177
+ self.p_generate = partitioner.partition(
178
+ generate,
179
+ in_axis_resources=(params_spec, P("data"), None),
180
+ out_axis_resources=P("data"),
181
+ static_argnums=(3,),
182
+ )
183
+
184
+ def generate(self, input_features, language=None, task=None, return_timestamps=False):
185
+ forced_decoder_ids = self.get_forced_decoder_ids(
186
+ language=language, task=task, return_timestamps=return_timestamps
187
+ )
188
+ if not self.is_sharded:
189
+ # if we're using pmap we need to manually replicate the input data across devices and gather the output tokens
190
+ output_ids = self.p_generate(
191
+ freeze(self.params), shard(input_features), forced_decoder_ids, return_timestamps
192
+ ).sequences
193
+ output_ids = jax.device_get(output_ids.reshape(-1, self.max_length))
194
+ else:
195
+ # pjit handles replication / gathering for us auto-magically
196
+ output_ids = self.p_generate(
197
+ freeze(self.params), input_features, forced_decoder_ids, return_timestamps
198
+ ).sequences
199
+ return output_ids
200
+
201
+ def get_forced_decoder_ids(self, generation_config=None, task=None, language=None, return_timestamps=False):
202
+ if generation_config is None:
203
+ generation_config = self.model.generation_config
204
+
205
+ if hasattr(generation_config, "is_multilingual"):
206
+ is_multilingual = generation_config.is_multilingual
207
+ else:
208
+ is_multilingual = None
209
+
210
+ forced_decoder_ids = []
211
+
212
+ if is_multilingual:
213
+ if language is not None:
214
+ language = language.lower()
215
+ if language in generation_config.lang_to_id.keys():
216
+ language_token = language
217
+ elif language in TO_LANGUAGE_CODE.values():
218
+ language_token = f"<|{language}|>"
219
+ elif language in TO_LANGUAGE_CODE.keys():
220
+ language_token = f"<|{TO_LANGUAGE_CODE[language]}|>"
221
+ else:
222
+ if len(language) == 2:
223
+ # ISO 639-1 language code
224
+ acceptable_languages = list(TO_LANGUAGE_CODE.values())
225
+ elif "<" in language or "|" in language or ">" in language:
226
+ # generation config language code
227
+ acceptable_languages = list(generation_config.lang_to_id.keys())
228
+ else:
229
+ # language passed as a string
230
+ acceptable_languages = list(TO_LANGUAGE_CODE.keys())
231
+ raise ValueError(
232
+ f"Unsupported language: {language}. Language should be one of:" f" {acceptable_languages}."
233
+ )
234
+ forced_decoder_ids.append((1, generation_config.lang_to_id[language_token]))
235
+
236
+ if task is not None:
237
+ forced_decoder_ids.append((2, generation_config.task_to_id[task]))
238
+ else:
239
+ forced_decoder_ids.append((2, generation_config.task_to_id["transcribe"]))
240
+
241
+ if not return_timestamps:
242
+ if forced_decoder_ids and forced_decoder_ids[-1][0] != generation_config.no_timestamps_token_id:
243
+ idx = forced_decoder_ids[-1][0] + 1 if forced_decoder_ids else 1
244
+ forced_decoder_ids.append((idx, generation_config.no_timestamps_token_id))
245
+
246
+ return forced_decoder_ids
247
+
248
+ def chunk_iter_with_batch(self, inputs, chunk_len, stride_left, stride_right, batch_size):
249
+ inputs_len = inputs.shape[0]
250
+ step = chunk_len - stride_left - stride_right
251
+
252
+ all_chunk_start_idx = np.arange(0, inputs_len, step)
253
+ num_samples = len(all_chunk_start_idx)
254
+
255
+ num_batches = math.ceil(num_samples / batch_size)
256
+ batch_idx = np.array_split(np.arange(num_samples), num_batches)
257
+
258
+ for idx in batch_idx:
259
+ chunk_start_idx = all_chunk_start_idx[idx]
260
+ chunk_end_idx = chunk_start_idx + chunk_len
261
+
262
+ chunks = [inputs[chunk_start:chunk_end] for chunk_start, chunk_end in zip(chunk_start_idx, chunk_end_idx)]
263
+ processed = self.feature_extractor(
264
+ chunks, sampling_rate=self.feature_extractor.sampling_rate, return_tensors="np"
265
+ )
266
+
267
+ _stride_left = np.where(chunk_start_idx == 0, 0, stride_left)
268
+ is_last = np.where(stride_right > 0, chunk_end_idx > inputs_len, chunk_end_idx >= inputs_len)
269
+ _stride_right = np.where(is_last, 0, stride_right)
270
+
271
+ chunk_lens = [chunk.shape[0] for chunk in chunks]
272
+ strides = [
273
+ (chunk_l, _stride_l, _stride_r)
274
+ for chunk_l, _stride_l, _stride_r in zip(chunk_lens, _stride_left, _stride_right)
275
+ ]
276
+
277
+ yield {"stride": strides, **processed}
278
+
279
+ def preprocess_batch(self, inputs, chunk_length_s=30.0, stride_length_s=None, batch_size=None):
280
+ if isinstance(inputs, np.ndarray):
281
+ logger.warning(
282
+ "Numpy array passed as input - no sampling rate checks will be performed."
283
+ "It is strongly recommended to pass the input as a dictionary with an 'array' key "
284
+ "containing the numpy array representing the audio, and a 'sampling_rate' key "
285
+ "containing the sampling rate associated with the audio array."
286
+ "Failing to do so can result in silent errors that might be hard to debug."
287
+ )
288
+
289
+ if isinstance(inputs, str):
290
+ if inputs.startswith("http://") or inputs.startswith("https://"):
291
+ # We need to actually check for a real protocol, otherwise it's impossible to use a local file
292
+ # like http_huggingface_co.png
293
+ inputs = requests.get(inputs).content
294
+ else:
295
+ with open(inputs, "rb") as f:
296
+ inputs = f.read()
297
+
298
+ if isinstance(inputs, bytes):
299
+ inputs = ffmpeg_read(inputs, self.feature_extractor.sampling_rate)
300
+
301
+ stride = None
302
+ if isinstance(inputs, dict):
303
+ stride = inputs.get("stride", None)
304
+ # Accepting `"array"` which is the key defined in `datasets` for
305
+ # better integration
306
+ if not ("sampling_rate" in inputs and "array" in inputs):
307
+ raise ValueError(
308
+ "When passing a dictionary to FlaxWhisperPipline, the dict needs to contain an 'array' key "
309
+ "containing the numpy array representing the audio, and a 'sampling_rate' key "
310
+ "containing the sampling rate associated with the audio array."
311
+ )
312
+
313
+ in_sampling_rate = inputs.get("sampling_rate")
314
+ inputs = inputs.get("array", None)
315
+
316
+ if in_sampling_rate != self.feature_extractor.sampling_rate:
317
+ try:
318
+ import librosa
319
+ except ImportError as err:
320
+ raise ImportError(
321
+ "To support resampling audio files, please install 'librosa' and 'soundfile'."
322
+ ) from err
323
+
324
+ inputs = librosa.resample(
325
+ inputs, orig_sr=in_sampling_rate, target_sr=self.feature_extractor.sampling_rate
326
+ )
327
+ ratio = self.feature_extractor.sampling_rate / in_sampling_rate
328
+ else:
329
+ ratio = 1
330
+
331
+ if not isinstance(inputs, np.ndarray):
332
+ raise ValueError(f"We expect a numpy ndarray as input, got `{type(inputs)}`")
333
+ if len(inputs.shape) != 1:
334
+ raise ValueError("We expect a single channel audio input for AutomaticSpeechRecognitionPipeline")
335
+
336
+ if stride is not None:
337
+ if stride[0] + stride[1] > inputs.shape[0]:
338
+ raise ValueError("Stride is too large for input")
339
+
340
+ # Stride needs to get the chunk length here, it's going to get
341
+ # swallowed by the `feature_extractor` later, and then batching
342
+ # can add extra data in the inputs, so we need to keep track
343
+ # of the original length in the stride so we can cut properly.
344
+ stride = (inputs.shape[0], int(round(stride[0] * ratio)), int(round(stride[1] * ratio)))
345
+
346
+ if chunk_length_s:
347
+ if stride_length_s is None:
348
+ stride_length_s = chunk_length_s / 6
349
+
350
+ if isinstance(stride_length_s, (int, float)):
351
+ stride_length_s = [stride_length_s, stride_length_s]
352
+
353
+ chunk_len = round(chunk_length_s * self.feature_extractor.sampling_rate)
354
+ stride_left = round(stride_length_s[0] * self.feature_extractor.sampling_rate)
355
+ stride_right = round(stride_length_s[1] * self.feature_extractor.sampling_rate)
356
+
357
+ if chunk_len < stride_left + stride_right:
358
+ raise ValueError("Chunk length must be superior to stride length")
359
+
360
+ for item in self.chunk_iter_with_batch(
361
+ inputs,
362
+ chunk_len,
363
+ stride_left,
364
+ stride_right,
365
+ batch_size,
366
+ ):
367
+ yield item
368
+ else:
369
+ processed = self.feature_extractor(
370
+ inputs, sampling_rate=self.feature_extractor.sampling_rate, return_tensors="np"
371
+ )
372
+ if stride is not None:
373
+ processed["stride"] = stride
374
+ yield processed
375
+
376
+ def postprocess(self, model_outputs, return_timestamps=None, return_language=None):
377
+ # unpack the outputs from list(dict(list)) to list(dict)
378
+ model_outputs = [dict(zip(output, t)) for output in model_outputs for t in zip(*output.values())]
379
+
380
+ time_precision = self.feature_extractor.chunk_length / self.model.config.max_source_positions
381
+ # Send the chunking back to seconds, it's easier to handle in whisper
382
+ sampling_rate = self.feature_extractor.sampling_rate
383
+ for output in model_outputs:
384
+ if "stride" in output:
385
+ chunk_len, stride_left, stride_right = output["stride"]
386
+ # Go back in seconds
387
+ chunk_len /= sampling_rate
388
+ stride_left /= sampling_rate
389
+ stride_right /= sampling_rate
390
+ output["stride"] = chunk_len, stride_left, stride_right
391
+
392
+ text, optional = self.tokenizer._decode_asr(
393
+ model_outputs,
394
+ return_timestamps=return_timestamps,
395
+ return_language=return_language,
396
+ time_precision=time_precision,
397
+ )
398
+ return {"text": text, **optional}
399
+
400
+ def forward(self, model_inputs, batch_size=None, language=None, task=None, return_timestamps=False):
401
+ # We need to keep track of some additional input arguments for post-processing so need to forward these on after running generation
402
+ input_features = model_inputs.pop("input_features")
403
+ input_batch_size = input_features.shape[0]
404
+
405
+ if input_batch_size != batch_size:
406
+ padding = np.zeros([batch_size - input_batch_size, *input_features.shape[1:]], input_features.dtype)
407
+ input_features = np.concatenate([input_features, padding])
408
+
409
+ pred_ids = self.generate(input_features, language=language, task=task, return_timestamps=return_timestamps)[
410
+ :input_batch_size
411
+ ]
412
+
413
+ # tokenizer's decode method expects an extra dim - we insert it here for convenience
414
+ out = {"tokens": pred_ids[:, None, :]}
415
+
416
+ stride = model_inputs.pop("stride", None)
417
+ if stride is not None:
418
+ out["stride"] = stride
419
+
420
+ return out
421
+
422
+ def __call__(
423
+ self,
424
+ inputs,
425
+ chunk_length_s=30.0,
426
+ stride_length_s=None,
427
+ batch_size=None,
428
+ language=None,
429
+ task=None,
430
+ return_timestamps=None,
431
+ generate_kwargs=None,
432
+ ):
433
+ """
434
+ Transcribe an audio input sequence to a text transcription, optionally with timestamps.
435
+
436
+ Args:
437
+ inputs (`np.ndarray` or `bytes` or `str` or `dict`):
438
+ The inputs is either:
439
+ - `str` that is the filename of the audio file, the file will be read at the correct sampling rate
440
+ to get the waveform using *ffmpeg*. This requires *ffmpeg* to be installed on the system.
441
+ - `bytes` is the byte content of an audio file and is interpreted by *ffmpeg* in the
442
+ same way.
443
+ - (`np.ndarray` of shape (n, ) of type `np.float32` or `np.float64`)
444
+ Raw audio assumed to be at the correct sampling rate (16kHz). Note that no further sampling
445
+ rate check will be done.
446
+ - `dict` form can be used to pass raw audio sampled at arbitrary `sampling_rate` and let this
447
+ pipeline do the resampling. The dict must be in the format `{"sampling_rate": int, "array":
448
+ np.array}`. Optionally an additional argument `"stride": (left: int, right: int)` can be used to
449
+ ask the pipeline to treat the first `left` samples and last `right` samples to be ignored in
450
+ decoding (but used at inference to provide more context to the model). In general, this additional
451
+ stride argument is not required.
452
+ chunk_length_s (`float`, *optional*, defaults to 30.0):
453
+ The input length for each chunk. If `chunk_length_s = 0` then chunking is disabled. By default, the chunk
454
+ length is set 30.0s, equal to Whisper's context window.
455
+ stride_length_s (`float`, *optional*, defaults to `chunk_length_s / 6`):
456
+ The length of stride on the left and right of each chunk. Used only with `chunk_length_s > 0`. This enables
457
+ the model to *see* more context and infer letters better than without this context but the pipeline
458
+ discards the stride bits at the end to make the final reconstitution as perfect as possible.
459
+
460
+ <Tip>
461
+
462
+ For more information on how to effectively use `stride_length_s`, refer to the [ASR chunking
463
+ blog post](https://huggingface.co/blog/asr-chunking).
464
+
465
+ </Tip>
466
+ batch_size (`int`, *optional*, defaults to the minimum per-device batch size, i.e. `jax.local_device_count()`):
467
+ The batch size to be used in chunking transcription. Beneficial for transcribing long audio files. Passing
468
+ a batch size in the `__call__` method will supersede any batch size passed to the `__init__`.
469
+ task (`str`, *optional*):
470
+ Task to use for generation, either `"transcribe"` or `"translate"`. Defaults to `"transcribe"`.
471
+ language (`str`, *optional*):
472
+ Language token to use for generation, can be either in the form of `"<|en|>"`, `"en"` or `"english"`.
473
+ Defaults to `None`, meaning the language is automatically inferred from the audio input.
474
+ return_timestamps (*optional*, `bool`):
475
+ Whether to return timestamps in the prediction. Defaults to False. If set to true, the pipeline
476
+ will return two keys in the output dictionary: `"text"` containing the text transcription, and `"chunks"`
477
+ containing the transcription segments chunked by their utterance-level timestamps.
478
+
479
+ Return:
480
+ `Dict`: A dictionary with the following keys:
481
+ - **text** (`str` ) -- The recognised text.
482
+ - **chunks** (*optional(, `List[Dict]`)
483
+ When using `return_timestamps`, the `chunks` will become a list containing all the various text
484
+ chunks identified by the model, *e.g.* `[{"text": "hi ", "timestamps": (0.5,0.9), {"text":
485
+ "there", "timestamps": (1.0, 1.5)}]`. The original full text can roughly be recovered by doing
486
+ `"".join(chunk["text"] for chunk in output["chunks"])`.
487
+ """
488
+ batch_size = batch_size if batch_size is not None else self.batch_size
489
+ if batch_size % self.min_batch_size != 0:
490
+ raise ValueError(
491
+ f"Batch size must be a multiple of the number of JAX devices, but got batch size {batch_size} and num devices {self.min_batch_size}."
492
+ )
493
+
494
+ dataloader = self.preprocess_batch(
495
+ inputs, chunk_length_s=chunk_length_s, stride_length_s=stride_length_s, batch_size=batch_size
496
+ )
497
+ model_outputs = []
498
+ # iterate over our chunked audio samples
499
+ for batch in dataloader:
500
+ model_outputs.append(
501
+ self.forward(
502
+ batch, batch_size=batch_size, language=language, task=task, return_timestamps=return_timestamps
503
+ )
504
+ )
505
+ post_processed = self.postprocess(model_outputs, return_timestamps=return_timestamps)
506
+ return post_processed
whisper_jax/train_state.py ADDED
@@ -0,0 +1,130 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # coding=utf-8
2
+ # Copyright 2023 The T5X Authors and The HuggingFace Inc. team.
3
+ #
4
+ # Licensed under the Apache License, Version 2.0 (the "License");
5
+ # you may not use this file except in compliance with the License.
6
+ # You may obtain a copy of the License at
7
+ #
8
+ # http://www.apache.org/licenses/LICENSE-2.0
9
+ #
10
+ # Unless required by applicable law or agreed to in writing, software
11
+ # distributed under the License is distributed on an "AS IS" BASIS,
12
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
+ # See the License for the specific language governing permissions and
14
+ # limitations under the License.
15
+
16
+ """Utilities for partitioning."""
17
+
18
+ from typing import Any, Mapping, MutableMapping, Optional, Tuple
19
+
20
+ import flax.core
21
+ import flax.serialization
22
+ import flax.struct
23
+ import jax.numpy as jnp
24
+ from flax import traverse_util
25
+ from flax.core import scope as flax_scope
26
+ from flax.linen import partitioning as flax_partitioning
27
+
28
+
29
+ EMPTY_DICT = flax.core.freeze({})
30
+ FrozenDict = flax_scope.FrozenDict
31
+ FrozenVariableDict = flax_scope.FrozenVariableDict
32
+ MutableVariableDict = flax_scope.MutableVariableDict
33
+ VariableDict = flax_scope.VariableDict
34
+
35
+
36
+ def _validate_params_axes(params_axes, params):
37
+ axis_names = flax_partitioning.get_axis_names(params_axes)
38
+ missing_params_axes = set(traverse_util.flatten_dict(params, sep="/")) - set(
39
+ traverse_util.flatten_dict(axis_names, sep="/")
40
+ )
41
+ if missing_params_axes:
42
+ raise ValueError(f"Missing axis names for parameters: {missing_params_axes}")
43
+
44
+
45
+ def _split_variables_and_axes(variables_and_axes: FrozenVariableDict) -> Tuple[FrozenVariableDict, FrozenVariableDict]:
46
+ """Splits `variables_and_axes` into two separate dicts with the same keys."""
47
+ # For each `key`, `key_axes` (if any) are its axes in `variables_and_axes`.
48
+ variables = {}
49
+ axes = {}
50
+ for k, v in variables_and_axes.items():
51
+ if k.endswith("_axes"):
52
+ axes[k[:-5]] = v # k without "_axes".
53
+ _validate_params_axes(v, variables_and_axes[k[:-5]]) # k without "_axes".
54
+ else:
55
+ variables[k] = v
56
+ return flax.core.freeze(variables), flax.core.freeze(axes)
57
+
58
+
59
+ class InferenceState(flax.struct.PyTreeNode):
60
+ """State compatible with FlaxOptimTrainState without optimizer state."""
61
+
62
+ step: jnp.ndarray
63
+ params: flax_scope.FrozenVariableDict
64
+ params_axes: Optional[flax_scope.FrozenVariableDict] = None
65
+ flax_mutables: flax_scope.FrozenDict = EMPTY_DICT
66
+ flax_mutables_axes: Optional[flax_scope.FrozenVariableDict] = None
67
+
68
+ @classmethod
69
+ def create(cls, model_variables: FrozenVariableDict) -> "InferenceState":
70
+ other_variables, params = model_variables.pop("params")
71
+ if "params_axes" in other_variables:
72
+ other_variables, params_axes = other_variables.pop("params_axes")
73
+ _validate_params_axes(params_axes, params)
74
+ else:
75
+ params_axes = None
76
+
77
+ # Split other_variables into mutables and their corresponding axes.
78
+ flax_mutables, flax_mutables_axes = _split_variables_and_axes(other_variables)
79
+ flax_mutables_axes = flax_mutables_axes or None
80
+ return InferenceState(
81
+ step=jnp.array(0),
82
+ params=params,
83
+ params_axes=params_axes,
84
+ flax_mutables=flax_mutables,
85
+ flax_mutables_axes=flax_mutables_axes,
86
+ )
87
+
88
+ @property
89
+ def param_states(self) -> FrozenVariableDict:
90
+ """The optimizer states of the parameters as a PyTree."""
91
+ raise NotImplementedError("InferenceState has no optimizer states.")
92
+
93
+ def apply_gradient(self, *args, **kwargs) -> "InferenceState":
94
+ raise NotImplementedError("InferenceState does not support `apply_gradient`.")
95
+
96
+ def state_dict(self) -> MutableMapping[str, Any]:
97
+ state_dict = {"target": flax.core.unfreeze(self.params), "state": {"step": self.step}}
98
+ if self.flax_mutables:
99
+ state_dict["flax_mutables"] = flax.core.unfreeze(self.flax_mutables)
100
+ return state_dict
101
+
102
+ def replace_step(self, step: jnp.ndarray) -> "InferenceState":
103
+ return self.replace(step=step)
104
+
105
+ def replace_params(self, params: FrozenVariableDict) -> "InferenceState":
106
+ return self.replace(params=params)
107
+
108
+ def replace_flax_mutables(self, flax_mutables: FrozenDict) -> "InferenceState":
109
+ return self.replace(flax_mutables=flax_mutables)
110
+
111
+ def restore_state(self, state_dict: Mapping[str, Any]) -> "InferenceState":
112
+ return self.replace(
113
+ params=flax.core.freeze(state_dict["target"]),
114
+ step=state_dict["state"]["step"],
115
+ flax_mutables=flax.core.freeze(state_dict["flax_mutables"])
116
+ if "flax_mutables" in state_dict
117
+ else EMPTY_DICT,
118
+ )
119
+
120
+ def as_logical_axes(self) -> "InferenceState":
121
+ # Set step to None so that when the logical axes are processed by the
122
+ # flax.partitioning.logical_to_mesh_axes function, it will be skipped
123
+ # because jax.tree_map will short circut and never call the function on the
124
+ # step.
125
+ flax_mutables_axes = self.flax_mutables_axes or EMPTY_DICT
126
+ return InferenceState(
127
+ step=None,
128
+ params=flax_partitioning.get_axis_names(self.params_axes),
129
+ flax_mutables=flax_partitioning.get_axis_names(flax_mutables_axes),
130
+ )