tpu-inference 0.11.1.dev202511150811__py3-none-any.whl

tpu_inference/layers/jax/moe/gpt_oss_moe.py

@@ -0,0 +1,185 @@
+from dataclasses import InitVar, dataclass
+
+import jax
+import jax.numpy as jnp
+from flax import nnx
+from flax.typing import Sharding
+from jaxtyping import Float
+
+from tpu_inference.layers.jax.base import create_param
+from tpu_inference.layers.jax.layers import FlaxUtils
+from tpu_inference.layers.jax.moe.moe import Router
+
+modeling_flax_utils = FlaxUtils()
+
+
+@dataclass(kw_only=True)
+class GptOssRouter(Router):
+    """Router module for Mixture-of-Experts (MoE) layers.
+
+    This module determines which experts each token should be routed.
+
+    """
+    e_sharding: Sharding = ()
+
+    def __post_init__(self, rngs: nnx.Rngs):
+        """
+        Initializes the parent's kernel and adds the new bias parameter.
+        """
+        super().__post_init__(rngs)
+
+        self.bias_E = create_param(rngs,
+                                   shape=(self.num_experts, ),
+                                   dtype=self.dtype,
+                                   sharding=self.e_sharding,
+                                   random_init=self.random_init)
+
+    def __call__(self, x_TD: Float):
+        """
+        Overrides the parent's forward pass to include the bias.
+        """
+        x_TD = jnp.asarray(x_TD, self.dtype)
+        x_TD = nnx.with_sharding_constraint(x_TD, self.activation_ffw_td)
+
+        router_logits_TE = jnp.einsum('TD,DE -> TE', x_TD,
+                                      self.kernel_DE.value)
+
+        router_logits_TE += self.bias_E.value
+
+        weights_TX, selected_experts_TX = jax.lax.top_k(
+            router_logits_TE, self.num_experts_per_tok)
+
+        normalized_weights_TX = jax.nn.softmax(weights_TX.astype(self.dtype),
+                                               axis=-1)
+
+        return normalized_weights_TX, selected_experts_TX
+
+
+def _swiglu(x: Float, alpha: Float, limit: Float) -> Float:
+    """Implements the specific SwiGLU from the golden implementation."""
+    x_glu, x_linear = x[..., ::2], x[..., 1::2]
+
+    x_glu = jnp.clip(x_glu, a_max=limit)
+    x_linear = jnp.clip(x_linear, a_min=-limit, a_max=limit)
+
+    gated_activation = x_glu * jax.nn.sigmoid(alpha * x_glu)
+
+    return gated_activation * (x_linear + 1)
+
+
+@dataclass(kw_only=True)
+class CombineExperts(nnx.Module):
+    """Module for combining expert outputs with weighted sum."""
+    dtype: jnp.dtype
+
+    def __call__(self, down_proj_TED: Float, weights_TX: Float,
+                 indices_TX: jax.Array) -> Float:
+        """Combines expert outputs using weighted sum.
+
+        Args:
+            down_proj_TED: Expert outputs, shape (tokens, experts, hidden_dim)
+            weights_TX: Router weights, shape (tokens, experts_per_token)
+            indices_TX: Selected expert indices, shape (tokens, experts_per_token)
+
+        Returns:
+            Combined output, shape (tokens, hidden_dim)
+        """
+        with jax.named_scope("combine_experts"):
+            indices_for_gather = indices_TX[..., None]
+            gathered_down_proj_TED = jnp.take_along_axis(down_proj_TED,
+                                                         indices_for_gather,
+                                                         axis=1)
+            output_TD = jnp.einsum('TXD,TX -> TD', gathered_down_proj_TED,
+                                   weights_TX)
+
+        return output_TD.astype(self.dtype)
+
+
+@dataclass(kw_only=True)
+class GptOssMoE(nnx.Module):
+    """
+    JAX implementation of the GPT-OSS Mixture-of-Experts MLP block.
+    """
+    dtype: jnp.dtype
+    hidden_size: int
+    intermediate_size_moe: int
+    num_local_experts: int
+    router: GptOssRouter
+    rngs: InitVar[nnx.Rngs]
+
+    swiglu_limit: float = 7.0
+    swiglu_alpha: float = 1.702
+
+    # Sharding specifications
+    activation_ffw_td: Sharding
+    edf_sharding: Sharding
+    efd_sharding: Sharding
+    ed_sharding: Sharding
+
+    random_init: bool = False
+
+    def __call__(self, x_TD: Float) -> Float:
+        """Performs the forward pass for the GPT-OSS MoE layer."""
+        x_TD = jnp.asarray(x_TD, self.dtype)
+        x_TD = nnx.with_sharding_constraint(x_TD, self.activation_ffw_td)
+
+        weights_TX, indices_TX = self.router(x_TD)
+
+        # First MLP layer (up-projection)
+        with jax.named_scope("MLP #1"):
+            up_proj_TEF2 = jnp.einsum('TD,EDF -> TEF', x_TD,
+                                      self.mlp1_weight_EDF2.value)
+            up_proj_TEF2 += self.mlp1_bias_EF2.value
+
+            fuse_TEF = _swiglu(up_proj_TEF2,
+                               alpha=self.swiglu_alpha,
+                               limit=self.swiglu_limit)
+
+        # Second MLP layer (down-projection)
+        with jax.named_scope("MLP #2"):
+            down_proj_TED = jnp.einsum('TEF,EFD -> TED', fuse_TEF,
+                                       self.mlp2_weight_EFD.value)
+            down_proj_TED += self.mlp2_bias_ED.value
+
+        # Weighted sum of expert outputs
+        output_TD = self.combine_experts(down_proj_TED, weights_TX, indices_TX)
+
+        return output_TD
+
+    def __post_init__(self, rngs: nnx.Rngs):
+        """Initializes all weights and biases for the MoE block."""
+        D, F, E = self.hidden_size, self.intermediate_size_moe, self.num_local_experts
+
+        self.combine_experts = CombineExperts(dtype=self.dtype)
+
+        # MLP #1 Weights (Combined Gate and Up-projection) and Bias
+        self.mlp1_weight_EDF2 = create_param(
+            rngs,
+            shape=(E, D, F * 2),
+            dtype=self.dtype,
+            sharding=self.edf_sharding,
+            random_init=self.random_init,
+        )
+        self.mlp1_bias_EF2 = create_param(
+            rngs,
+            shape=(E, F * 2),
+            dtype=self.dtype,
+            sharding=self.ed_sharding,
+            random_init=self.random_init,
+        )
+
+        # MLP #2 Weights (Down-projection) and Bias
+        self.mlp2_weight_EFD = create_param(
+            rngs,
+            shape=(E, F, D),
+            dtype=self.dtype,
+            sharding=self.efd_sharding,
+            random_init=self.random_init,
+        )
+        self.mlp2_bias_ED = create_param(
+            rngs,
+            shape=(E, D),
+            dtype=self.dtype,
+            sharding=self.ed_sharding,
+            random_init=self.random_init,
+        )

tpu_inference/layers/jax/moe/moe.py

@@ -0,0 +1,209 @@
+from dataclasses import InitVar, dataclass
+
+import jax
+import jax.numpy as jnp
+from flax import nnx
+from flax.typing import Sharding
+from jaxtyping import Float
+
+from tpu_inference.layers.jax.base import create_param
+from tpu_inference.layers.jax.layers import FlaxUtils
+
+modeling_flax_utils = FlaxUtils()
+
+
+@dataclass(kw_only=True)
+class Router(nnx.Module):
+    """Router module for Mixture-of-Experts (MoE) layers.
+
+    This module determines which experts each token should be routed to based on the input.
+
+    Attributes:
+    """
+    dtype: jnp.dtype
+    hidden_size: int
+    num_experts: int
+    num_experts_per_tok: int
+    router_act: str
+    rngs: InitVar[nnx.Rngs]
+    activation_ffw_td: Sharding
+    ed_sharding: Sharding
+    random_init: bool = False
+
+    def __call__(self, x_TD: Float):
+        """Routes tokens to experts.
+
+        Args:
+            x_TD: Input array of shape (sequence_length, d_model).
+
+        Returns:
+            A tuple containing:
+                - normalized_weights_TX: Normalized weights for selected experts, shape (sequence_length, num_experts_per_tok).
+                - selected_experts_TX: Indices of selected experts, shape (sequence_length, num_experts_per_tok).
+        """
+        x_TD = jnp.asarray(x_TD, self.dtype)
+        x_TD = nnx.with_sharding_constraint(x_TD, self.activation_ffw_td)
+        router_act = modeling_flax_utils.ACT2FN[self.router_act]
+        router_logits_TE = jnp.einsum('TD,DE -> TE', x_TD,
+                                      self.kernel_DE.value)
+        weights_TX, selected_experts_TX = jax.lax.top_k(
+            router_logits_TE, self.num_experts_per_tok)
+        if self.router_act != "sigmoid":  # sigmoid does not accept axis argument.
+            normalized_weights_TX = router_act(weights_TX.astype(self.dtype),
+                                               axis=-1)
+        else:
+            normalized_weights_TX = router_act(weights_TX.astype(self.dtype))
+        return normalized_weights_TX, selected_experts_TX
+
+    def __post_init__(self, rngs: nnx.Rngs):
+        """Generates the router kernel (weights) for routing."""
+        shape = (self.hidden_size, self.num_experts)
+        self.kernel_DE = create_param(rngs,
+                                      shape=shape,
+                                      dtype=self.dtype,
+                                      sharding=self.ed_sharding,
+                                      random_init=self.random_init)
+
+
+@dataclass(kw_only=True)
+class MoE(nnx.Module):
+    """Mixture-of-Experts (MoE) Routed MLP Layer.
+
+    This module implements a MoE layer with a router and multiple expert MLPs.
+
+    Attributes:
+        router: The Router module.
+    """
+    dtype: jnp.dtype
+    num_local_experts: int
+    apply_expert_weight_before_computation: bool
+    hidden_size: int
+    intermediate_size_moe: int
+    hidden_act: str
+    rngs: InitVar[nnx.Rngs]
+    router: nnx.Module
+    activation_ffw_td: Sharding
+    activation_ffw_ted: Sharding
+    edf_sharding: Sharding
+    efd_sharding: Sharding
+    random_init: bool = False
+
+    def __call__(self, x_TD: Float):
+        """Performs the forward pass of the MoE layer.
+
+        Args:
+            x_TD: Input array of shape (sequence_length, d_model).
+
+        Returns:
+            Output array of shape (sequence_length, d_model) after passing through MoE.
+        """
+        x_TD = jnp.asarray(x_TD, self.dtype)
+        x_TD = nnx.with_sharding_constraint(x_TD, self.activation_ffw_td)
+        weights_TX, indices_TX = self.router(x_TD)
+        one_hot_indices_TXE = jax.nn.one_hot(
+            indices_TX, num_classes=self.num_local_experts, dtype=self.dtype)
+        full_weights_TE = jnp.sum(one_hot_indices_TXE * weights_TX[..., None],
+                                  axis=1)
+
+        # Some models use the routing scores to weight the data instead of
+        # weighting the expert outputs.
+        if self.apply_expert_weight_before_computation:
+            with jax.named_scope("pre_computing_weight"):
+                return self._moe_fwd_preapply_router_weights(
+                    x_TD, full_weights_TE)
+        else:
+            return self._moe_fwd(x_TD, full_weights_TE)
+
+    def __post_init__(self, rngs: nnx.Rngs):
+        """Generates the kernels (weights) for the router and experts (gating, up-projection, and down-projection layers)."""
+
+        D = self.hidden_size
+        F = self.intermediate_size_moe
+        shape_gating = (self.num_local_experts, D, F)
+        shape_up = (self.num_local_experts, D, F)
+        shape_down = (self.num_local_experts, F, D)
+
+        self.kernel_gating_EDF = create_param(rngs,
+                                              shape=shape_gating,
+                                              dtype=self.dtype,
+                                              sharding=self.edf_sharding,
+                                              random_init=self.random_init)
+        self.kernel_up_proj_EDF = create_param(rngs,
+                                               shape=shape_up,
+                                               dtype=self.dtype,
+                                               sharding=self.edf_sharding,
+                                               random_init=self.random_init)
+        self.kernel_down_proj_EFD = create_param(rngs,
+                                                 shape=shape_down,
+                                                 dtype=self.dtype,
+                                                 sharding=self.efd_sharding,
+                                                 random_init=self.random_init)
+
+    def _moe_fwd_preapply_router_weights(self, x_TD: jax.Array, weights_TE):
+        """Performs the forward pass of the MoE experts with router weights pre-applied to the inputs.
+
+        Args:
+            x_TD: Input array for the experts, shape (sequence_length, hidden_size).
+            weights_TE: Router weights, shape (sequence_length, num_experts).
+
+        Returns:
+            Output array of shape (sequence_length, d_model).
+        """
+        # Data needs to be replicated since it will be weighted by the router
+        # scores before being passed to each expert.
+        num_experts = weights_TE.shape[-1]
+        x_TED = jnp.repeat(x_TD[:, None, :], num_experts, 1)
+        weights_TED = weights_TE[..., None]
+        x_TED = jnp.asarray(x_TED, self.dtype)
+
+        with jax.named_scope("activation_expert_weighting"):
+            x_TED = x_TED * weights_TED
+
+        x_TED = nnx.with_sharding_constraint(x_TED, self.activation_ffw_ted)
+        with jax.named_scope("gating"):
+            gating_TEF = jnp.einsum('TED,EDF -> TEF', x_TED,
+                                    self.kernel_gating_EDF.value)
+            activated_gating_TEF = modeling_flax_utils.ACT2FN[self.hidden_act](
+                gating_TEF)
+        with jax.named_scope("up_projection"):
+            up_proj_TEF = jnp.einsum('TED,EDF -> TEF', x_TED,
+                                     self.kernel_up_proj_EDF.value)
+
+        fuse_TEF = activated_gating_TEF * up_proj_TEF
+
+        with jax.named_scope("down_projection"):
+            down_proj_TED = jnp.einsum('TEF,EFD -> TED', fuse_TEF,
+                                       self.kernel_down_proj_EFD.value)
+        with jax.named_scope("sum"):
+            output_TD = down_proj_TED.sum(axis=1)
+        return output_TD.astype(self.dtype)
+
+    def _moe_fwd(self, x_TD: Float, weights):
+        """Performs the basic forward pass of the MoE experts without dropping or megablocks.
+
+        Args:
+            x_TD: Input array for the experts, shape (sequence_length, d_model).
+            weights: Weights for combining expert outputs, shape (sequence_length, num_experts).
+
+        Returns:
+            Output array of shape (sequence_length, d_model).
+        """
+        x_TD = jnp.asarray(x_TD, self.dtype)
+        x_TD = nnx.with_sharding_constraint(x_TD, self.activation_ffw_td)
+        with jax.named_scope("gating"):
+            gating_TEF = jnp.einsum('TD,EDF -> TEF', x_TD,
+                                    self.kernel_gating_EDF.value)
+            activated_gating_TEF = modeling_flax_utils.ACT2FN[self.hidden_act](
+                gating_TEF)
+        with jax.named_scope("up_projection"):
+            up_proj_TEF = jnp.einsum('TD,EDF -> TEF', x_TD,
+                                     self.kernel_up_proj_EDF.value)
+
+        fuse_TEF = activated_gating_TEF * up_proj_TEF
+
+        with jax.named_scope("down_projection"):
+            down_proj_TED = jnp.einsum('TEF,EFD -> TED', fuse_TEF,
+                                       self.kernel_down_proj_EFD.value)
+        with jax.named_scope("sum"):
+            output_TD = jnp.einsum('TED,TE -> TD', down_proj_TED, weights)
+        return output_TD.astype(self.dtype)

tpu_inference/layers/jax/rope.py

@@ -0,0 +1,280 @@
+import math
+from dataclasses import dataclass, field
+from typing import Optional, Tuple
+
+import jax
+from flax import nnx
+from jax import numpy as jnp
+from jax.experimental.layout import Layout, with_layout_constraint
+from jax.sharding import NamedSharding, PartitionSpec
+
+
+@dataclass(kw_only=True)
+class RotaryEmbedding(nnx.Module):
+    """
+    An implementation of the original rotary positional embedding.
+    """
+    rotary_dim: int
+    rope_theta: float
+    original_max_position_embeddings: int
+    dtype: jnp.dtype
+    sin_cos_cache: Optional[jax.Array] = field(init=False, default=None)
+
+    def initialize_cache(self):
+        """Computes and caches the sin/cos embeddings."""
+        if self.sin_cos_cache is None:
+            self.sin_cos_cache = self._compute_sin_cos()
+
+    def _compute_inv_freq(self):
+        fractions_H = jnp.arange(0, self.rotary_dim, 2,
+                                 dtype=jnp.float32) / self.rotary_dim
+        inv_freq_H = 1.0 / (self.rope_theta**fractions_H)
+        return inv_freq_H
+
+    def _compute_sin_cos(self):
+        inv_freq_H = self._compute_inv_freq()
+        t = jnp.arange(self.original_max_position_embeddings,
+                       dtype=jnp.float32)
+
+        freqs = jnp.einsum("...T,k->...Tk",
+                           t,
+                           inv_freq_H,
+                           precision=jax.lax.Precision.HIGHEST)
+        sin, cos = jnp.sin(freqs), jnp.cos(freqs)
+        cache = jnp.concatenate((cos, sin), axis=-1)
+        return cache
+
+    def apply_rope(self, positions: jax.Array, x_TNH: jax.Array):
+        assert x_TNH.ndim == 3
+        assert self.sin_cos_cache is not None, "RoPE cache not initialized."
+        cos_sin_TH = self.sin_cos_cache[positions]
+        # cos, sin: (T, H/2)
+        cos_TH, sin_TH = jnp.split(cos_sin_TH, 2, axis=-1)
+        assert sin_TH.ndim == 2 and cos_TH.ndim == 2
+        # cos, sin: (T, 1, H/2)
+        cos_T1H, sin_T1H = cos_TH[:, None, :], sin_TH[:, None, :]
+        # first_half, second_half: (T, N, H/2)
+        first_half_TNH, second_half_TNH = jnp.split(x_TNH, 2, axis=-1)
+        combined = jnp.concatenate([
+            first_half_TNH * cos_T1H - second_half_TNH * sin_T1H,
+            second_half_TNH * cos_T1H + first_half_TNH * sin_T1H
+        ],
+                                   axis=-1)
+        return combined.astype(self.dtype)
+
+
+@dataclass(kw_only=True)
+class DeepseekScalingRotaryEmbedding(RotaryEmbedding):
+    """
+    Rotary Embedding for deepseek, with scaling and YaRN method.
+    """
+    scaling_factor: float
+    beta_fast: int = 32
+    beta_slow: int = 1
+    mscale_value: float = 1
+    mscale_all_dim: float = 0
+
+    def initialize_cache(self, mesh: jax.sharding.Mesh):
+        """Computes and caches the sin/cos embeddings."""
+        # The second condition is for the Qwix case, where we need to call `initialize_cache` on
+        # the abstract model.  Thus, when we go to call `initialize_cache` on the concrete model,
+        # this method will have been called already, but we need to recompute the cache so that
+        # it's concrete (otherwise, it'll still be a jax.ShapeDtypeStruct).
+        if self.sin_cos_cache is not None and not isinstance(
+                self.sin_cos_cache, jax.ShapeDtypeStruct):
+            return
+        mscale_val = _yarn_get_mscale(
+            self.scaling_factor, self.mscale_value) / _yarn_get_mscale(
+                self.scaling_factor, self.mscale_all_dim)
+        replicated_sharding = NamedSharding(mesh, PartitionSpec())
+        self.mscale = jax.device_put(mscale_val, replicated_sharding)
+        self.sin_cos_cache = self._compute_sin_cos()
+
+    def _compute_inv_freq(self):
+        fractions = jnp.arange(0, self.rotary_dim, 2,
+                               dtype=jnp.float32) / self.rotary_dim
+        inv_freq_extrapolation = 1.0 / (self.rope_theta**fractions)
+        inv_freq_interpolation = 1.0 / (self.scaling_factor *
+                                        self.rope_theta**fractions)
+        low, high = _yarn_find_correction_range(
+            self.beta_fast, self.beta_slow, self.rotary_dim, self.rope_theta,
+            self.original_max_position_embeddings)
+
+        # Get n-d rotational scaling corrected for extrapolation
+        inv_freq_mask = 1 - _yarn_linear_ramp_mask(
+            low, high, self.rotary_dim // 2).astype(jnp.float32)
+        inv_freq = inv_freq_interpolation * (
+            1 - inv_freq_mask) + inv_freq_extrapolation * inv_freq_mask
+        return inv_freq
+
+    @jax.jit
+    def _compute_sin_cos(self):
+        inv_freq_H = self._compute_inv_freq()
+        t = jnp.arange(self.original_max_position_embeddings *
+                       self.scaling_factor,
+                       dtype=jnp.float32)
+        freqs = jnp.einsum("...T,k->...Tk", t, inv_freq_H)
+        sin, cos = jnp.sin(freqs) * self.mscale, jnp.cos(freqs) * self.mscale
+        cache = jnp.concatenate((cos, sin), axis=-1)
+        H = cache.shape[1]
+        target_dim = ((H - 1) // 128 + 1) * 128
+        padding_amount = target_dim - self.rotary_dim
+        pad_width = ((0, 0), (0, padding_amount))
+        cache_padded = jnp.pad(cache, pad_width, mode='constant')
+        desired_layout = Layout(major_to_minor=(1, 0))
+        cache_padded = with_layout_constraint(cache_padded, desired_layout)
+        return cache_padded
+
+    def apply_rope(self, positions: jax.Array, x_TNH: jax.Array):
+        assert x_TNH.ndim == 3
+        assert self.sin_cos_cache is not None, "RoPE cache not initialized."
+        cos_sin_padded = self.sin_cos_cache[positions]
+        cos_sin_TH = cos_sin_padded[:, :self.rotary_dim]
+        # cos, sin: (T, H/2)
+        cos_TH, sin_TH = jnp.split(cos_sin_TH, 2, axis=-1)
+        assert sin_TH.ndim == 2 and cos_TH.ndim == 2
+        # cos, sin: (T, 1, H/2)
+        cos_T1H, sin_T1H = cos_TH[:, None, :], sin_TH[:, None, :]
+        # even, odd: (T, N, H/2)
+        even_TNH, odd_TNH = x_TNH[..., ::2], x_TNH[..., 1::2]
+        combined_TNH = jnp.stack([
+            even_TNH * cos_T1H - odd_TNH * sin_T1H,
+            odd_TNH * cos_T1H + even_TNH * sin_T1H
+        ],
+                                 axis=-1).reshape(x_TNH.shape)
+        return combined_TNH.astype(self.dtype)
+
+
+# Calculates the temperature scaling factor for YaRN to adjust
+# RoPE embedding magnitudes.
+def _yarn_get_mscale(scale, mscale):
+    return jnp.where(scale <= 1, 1.0, 0.1 * mscale * jnp.log(scale) + 1.0)
+
+
+# Inverses dim formula to find dim based on number of rotations.
+def _yarn_find_correction_dim(num_rotations,
+                              dim,
+                              base=10000,
+                              max_position_embeddings=2048):
+    return (dim * math.log(max_position_embeddings /
+                           (num_rotations * 2 * math.pi))) / (2 *
+                                                              math.log(base))
+
+
+# Finds dim range bounds based on rotations.
+def _yarn_find_correction_range(low_rot,
+                                high_rot,
+                                dim,
+                                base=10000,
+                                max_position_embeddings=2048):
+    low = math.floor(
+        _yarn_find_correction_dim(low_rot, dim, base, max_position_embeddings))
+    high = math.ceil(
+        _yarn_find_correction_dim(high_rot, dim, base,
+                                  max_position_embeddings))
+    return max(low, 0), min(high, dim - 1)  # Clamp values just in case
+
+
+# Creates a 1D mask that ramps linearly from 0 to 1 between min and max indices.
+def _yarn_linear_ramp_mask(min, max, dim):
+    if min == max:
+        max += 0.001  # Prevent singularity
+
+    linear_func = (jnp.arange(dim, dtype=jnp.float32) - min) / (max - min)
+    ramp_func = jnp.clip(linear_func, 0, 1)
+    return ramp_func
+
+
+@dataclass(kw_only=True)
+class GptOssRotaryEmbedding(nnx.Module):
+    """
+    JAX implementation of the Rotary Positional Embedding with YaRN scaling.
+    """
+    head_dim: int
+    rope_theta: float
+    dtype: jnp.dtype
+    initial_context_length: int = 4096
+    rope_scaling_factor: float = 1.0
+    rope_ntk_alpha: float = 1.0
+    rope_ntk_beta: float = 32.0
+
+    def _compute_concentration_and_inv_freq(self) -> Tuple[float, jax.Array]:
+        """
+        Computes the inverse frequencies and concentration factor for YaRN.
+        See YaRN paper: https://arxiv.org/abs/2309.00071
+        """
+        freq = self.rope_theta**(
+            jnp.arange(0, self.head_dim, 2, dtype=jnp.float32) / self.head_dim)
+
+        if self.rope_scaling_factor > 1.0:
+            concentration = 0.1 * jnp.log(self.rope_scaling_factor) + 1.0
+
+            d_half = self.head_dim / 2
+            # NTK by parts
+            low = (d_half * jnp.log(self.initial_context_length /
+                                    (self.rope_ntk_beta * 2 * jnp.pi)) /
+                   jnp.log(self.rope_theta))
+            high = (d_half * jnp.log(self.initial_context_length /
+                                     (self.rope_ntk_alpha * 2 * jnp.pi)) /
+                    jnp.log(self.rope_theta))
+
+            interpolation = 1.0 / (self.rope_scaling_factor * freq)
+            extrapolation = 1.0 / freq
+
+            ramp = (jnp.arange(d_half, dtype=jnp.float32) - low) / (high - low)
+            mask = 1 - jnp.clip(ramp, 0, 1)
+
+            inv_freq = interpolation * (1 - mask) + extrapolation * mask
+        else:
+            concentration = 1.0
+            inv_freq = 1.0 / freq
+
+        return concentration, inv_freq
+
+    def _compute_cos_sin(self,
+                         positions: jax.Array) -> Tuple[jax.Array, jax.Array]:
+        """Computes cosine and sine embeddings for given positions."""
+        concentration, inv_freq_H = self._compute_concentration_and_inv_freq()
+
+        # freqs: (T, H/2)
+        freqs = jnp.einsum("T,H->TH",
+                           positions.astype(jnp.float32),
+                           inv_freq_H,
+                           precision=jax.lax.Precision.HIGHEST)
+
+        cos = jnp.cos(freqs) * concentration
+        sin = jnp.sin(freqs) * concentration
+        return cos, sin
+
+    def __call__(self, query_TNH: jax.Array, key_TNH: jax.Array,
+                 positions: jax.Array) -> Tuple[jax.Array, jax.Array]:
+        """
+        Applies rotary embeddings to query and key tensors.
+        Args:
+            query_TNH: Query tensor with shape (num_tokens, num_heads, head_dim)
+            key_TNH: Key tensor with shape (num_tokens, num_kv_heads, head_dim)
+            positions: A 1D array of token positions.
+        """
+        # cos, sin: (T, H/2)
+        cos_TH, sin_TH = self._compute_cos_sin(positions)
+
+        # Reshape for broadcasting: (T, 1, H/2)
+        cos_T1H = cos_TH[:, None, :]
+        sin_T1H = sin_TH[:, None, :]
+
+        def _apply_rotation(x_TNH: jax.Array) -> jax.Array:
+            # Split the last dimension
+            first_half, second_half = jnp.split(x_TNH, 2, axis=-1)
+
+            # Apply rotation
+            rotated_x = jnp.concatenate([
+                first_half * cos_T1H - second_half * sin_T1H,
+                second_half * cos_T1H + first_half * sin_T1H
+            ],
+                                        axis=-1)
+            return rotated_x.astype(self.dtype)
+
+        rotated_query = _apply_rotation(query_TNH)
+        rotated_key = _apply_rotation(key_TNH)
+
+        return rotated_query, rotated_key