# `vllm_mlx.rerank_forward`

MLX forward pass for BERT-family sequence classification models.

[View the complete module source at #L1-L265](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/rerank_forward.py#L1-L265).

## API details

Each callable below includes its exact signature, type annotations, inputs, defaults, return contract, documented exceptions, implementation source, and parsed docstring sections when the source provides them.

::: vllm_mlx.rerank_forward
    options:
      members:
        - classifier_forward
        - _position_ids_for_config
        - _classification_head_forward
        - _detect_prefix
        - _layer_norm
        - _encoder_layer
        - _gelu
        - _gelu_new
        - _relu
        - _silu
        - _apply_hidden_activation
      filters: []
      show_if_no_docstring: true

## Complete contract reference

Expand any definition for its exact inputs, annotations, defaults, return contract, directly raised exceptions, source-grounded behavior, and immutable line link. This section includes private and nested definitions that ordinary API generators omit.

<details class="api-contract" id="contract-vllm_mlx.rerank_forward.classifier_forward" markdown="1">
<summary><code>vllm_mlx.rerank_forward.classifier_forward</code> · function</summary>

```python
vllm_mlx.rerank_forward.classifier_forward(input_ids: mx.array, attention_mask: mx.array, weights: dict[str, mx.array], config: dict) -> mx.array
```

Run a BERT-family classifier forward pass on MLX.

**Parameters**

| Name | Type | Required | Default | Description |
| --- | --- | --- | --- | --- |
| `input_ids` | `mx.array` | `yes` | `none` | (batch, seq_len) token IDs. |
| `attention_mask` | `mx.array` | `yes` | `none` | (batch, seq_len) attention mask (1=attend, 0=pad). |
| `weights` | `dict[str, mx.array]` | `yes` | `none` | Dict mapping weight name -> mx.array. |
| `config` | `dict` | `yes` | `none` | Model config dict (from config.json). |

**Returns**

- Type: `mx.array`
- Direct return expressions: `logits`

**Exceptions and behavior**

Function `classifier_forward` calls `config.get`, `_detect_prefix`, `_position_ids_for_config`, `mx.zeros_like`; returns `logits`.
No direct `raise` statement appears in this definition.

[View source #L16-L85](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/rerank_forward.py#L16-L85).

</details>

<details class="api-contract" id="contract-vllm_mlx.rerank_forward._position_ids_for_config" markdown="1">
<summary><code>vllm_mlx.rerank_forward._position_ids_for_config</code> · function</summary>

```python
vllm_mlx.rerank_forward._position_ids_for_config(config: dict, input_ids: mx.array, attention_mask: mx.array | None) -> mx.array
```

Build BERT or RoBERTa-family absolute position IDs.

**Parameters**

| Name | Type | Required | Default | Description |
| --- | --- | --- | --- | --- |
| `config` | `dict` | `yes` | `none` | Required positional or keyword input. |
| `input_ids` | `mx.array` | `yes` | `none` | Required positional or keyword input. |
| `attention_mask` | `mx.array \| None` | `yes` | `none` | Required positional or keyword input. |

**Returns**

- Type: `mx.array`
- Direct return expressions: `mx.arange(seq_len)[None, :]`; `mx.arange(padding_idx + 1, seq_len + padding_idx + 1)[None, :]`; `positions.astype(mx.int32)`

**Exceptions and behavior**

Function `_position_ids_for_config` calls `str(config.get('model_type', '')).lower`, `str`, `config.get`, `mx.arange`; has 3 explicit return paths.
No direct `raise` statement appears in this definition.

[View source #L88-L105](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/rerank_forward.py#L88-L105).

</details>

<details class="api-contract" id="contract-vllm_mlx.rerank_forward._classification_head_forward" markdown="1">
<summary><code>vllm_mlx.rerank_forward._classification_head_forward</code> · function</summary>

```python
vllm_mlx.rerank_forward._classification_head_forward(pooled: mx.array, weights: dict[str, mx.array]) -> mx.array
```

Run BERT flat or XLM-RoBERTa two-layer sequence-classification head.

**Parameters**

| Name | Type | Required | Default | Description |
| --- | --- | --- | --- | --- |
| `pooled` | `mx.array` | `yes` | `none` | Required positional or keyword input. |
| `weights` | `dict[str, mx.array]` | `yes` | `none` | Required positional or keyword input. |

**Returns**

- Type: `mx.array`
- Direct return expressions: `hidden @ weights['classifier.out_proj.weight'].T + weights['classifier.out_proj.bias']`; `pooled @ weights['classifier.weight'].T + weights['classifier.bias']`

**Exceptions and behavior**

Function `_classification_head_forward` calls `mx.tanh`; has 2 explicit return paths.
No direct `raise` statement appears in this definition.

[View source #L108-L122](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/rerank_forward.py#L108-L122).

</details>

<details class="api-contract" id="contract-vllm_mlx.rerank_forward._detect_prefix" markdown="1">
<summary><code>vllm_mlx.rerank_forward._detect_prefix</code> · function</summary>

```python
vllm_mlx.rerank_forward._detect_prefix(weights: dict) -> str
```

Detect the model weight prefix (bert, roberta, xlm-roberta).

**Parameters**

| Name | Type | Required | Default | Description |
| --- | --- | --- | --- | --- |
| `weights` | `dict` | `yes` | `none` | Required positional or keyword input. |

**Returns**

- Type: `str`
- Direct return expressions: `'bert'`; `'roberta'`; `'xlm-roberta'`

**Exceptions and behavior**

Function `_detect_prefix` calls `key.startswith`; has 3 explicit return paths.
No direct `raise` statement appears in this definition.

[View source #L125-L135](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/rerank_forward.py#L125-L135).

</details>

<details class="api-contract" id="contract-vllm_mlx.rerank_forward._layer_norm" markdown="1">
<summary><code>vllm_mlx.rerank_forward._layer_norm</code> · function</summary>

```python
vllm_mlx.rerank_forward._layer_norm(x: mx.array, weight: mx.array, bias: mx.array, eps: float) -> mx.array
```

Apply layer normalization.

**Parameters**

| Name | Type | Required | Default | Description |
| --- | --- | --- | --- | --- |
| `x` | `mx.array` | `yes` | `none` | Required positional or keyword input. |
| `weight` | `mx.array` | `yes` | `none` | Required positional or keyword input. |
| `bias` | `mx.array` | `yes` | `none` | Required positional or keyword input. |
| `eps` | `float` | `yes` | `none` | Required positional or keyword input. |

**Returns**

- Type: `mx.array`
- Direct return expressions: `weight * (x - mean) / mx.sqrt(var + eps) + bias`

**Exceptions and behavior**

Function `_layer_norm` calls `mx.mean`, `mx.var`, `mx.sqrt`; returns `weight * (x - mean) / mx.sqrt(var + eps) + bias`.
No direct `raise` statement appears in this definition.

[View source #L138-L142](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/rerank_forward.py#L138-L142).

</details>

<details class="api-contract" id="contract-vllm_mlx.rerank_forward._encoder_layer" markdown="1">
<summary><code>vllm_mlx.rerank_forward._encoder_layer</code> · function</summary>

```python
vllm_mlx.rerank_forward._encoder_layer(hidden: mx.array, ext_mask: mx.array | None, weights: dict, prefix: str, num_heads: int, head_dim: int, eps: float, config: dict) -> mx.array
```

Run one BERT encoder layer (self-attention + FFN).

**Parameters**

| Name | Type | Required | Default | Description |
| --- | --- | --- | --- | --- |
| `hidden` | `mx.array` | `yes` | `none` | Required positional or keyword input. |
| `ext_mask` | `mx.array \| None` | `yes` | `none` | Required positional or keyword input. |
| `weights` | `dict` | `yes` | `none` | Required positional or keyword input. |
| `prefix` | `str` | `yes` | `none` | Required positional or keyword input. |
| `num_heads` | `int` | `yes` | `none` | Required positional or keyword input. |
| `head_dim` | `int` | `yes` | `none` | Required positional or keyword input. |
| `eps` | `float` | `yes` | `none` | Required positional or keyword input. |
| `config` | `dict` | `yes` | `none` | Required positional or keyword input. |

**Returns**

- Type: `mx.array`
- Direct return expressions: `hidden`

**Exceptions and behavior**

Function `_encoder_layer` calls `(hidden @ q_w.T + q_b).reshape(batch_size, seq_len, num_heads, head_dim).transpose`, `(hidden @ q_w.T + q_b).reshape`, `(hidden @ k_w.T + k_b).reshape(batch_size, seq_len, num_heads, head_dim).transpose`, `(hidden @ k_w.T + k_b).reshape`; returns `hidden`.
No direct `raise` statement appears in this definition.

[View source #L145-L217](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/rerank_forward.py#L145-L217).

</details>

<details class="api-contract" id="contract-vllm_mlx.rerank_forward._gelu" markdown="1">
<summary><code>vllm_mlx.rerank_forward._gelu</code> · function</summary>

```python
vllm_mlx.rerank_forward._gelu(x: mx.array) -> mx.array
```

GELU activation (exact form).

**Parameters**

| Name | Type | Required | Default | Description |
| --- | --- | --- | --- | --- |
| `x` | `mx.array` | `yes` | `none` | Required positional or keyword input. |

**Returns**

- Type: `mx.array`
- Direct return expressions: `nn.gelu(x)`

**Exceptions and behavior**

Function `_gelu` calls `nn.gelu`; returns `nn.gelu(x)`.
No direct `raise` statement appears in this definition.

[View source #L220-L222](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/rerank_forward.py#L220-L222).

</details>

<details class="api-contract" id="contract-vllm_mlx.rerank_forward._gelu_new" markdown="1">
<summary><code>vllm_mlx.rerank_forward._gelu_new</code> · function</summary>

```python
vllm_mlx.rerank_forward._gelu_new(x: mx.array) -> mx.array
```

BERT GELU approximation used by transformers gelu_new.

**Parameters**

| Name | Type | Required | Default | Description |
| --- | --- | --- | --- | --- |
| `x` | `mx.array` | `yes` | `none` | Required positional or keyword input. |

**Returns**

- Type: `mx.array`
- Direct return expressions: `0.5 * x * (1.0 + mx.tanh(0.7978845608028654 * (x + 0.044715 * x ** 3)))`

**Exceptions and behavior**

Function `_gelu_new` calls `mx.tanh`; returns `0.5 * x * (1.0 + mx.tanh(0.7978845608028654 * (x + 0.044715 * x ** 3)))`.
No direct `raise` statement appears in this definition.

[View source #L225-L227](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/rerank_forward.py#L225-L227).

</details>

<details class="api-contract" id="contract-vllm_mlx.rerank_forward._relu" markdown="1">
<summary><code>vllm_mlx.rerank_forward._relu</code> · function</summary>

```python
vllm_mlx.rerank_forward._relu(x: mx.array) -> mx.array
```

ReLU activation.

**Parameters**

| Name | Type | Required | Default | Description |
| --- | --- | --- | --- | --- |
| `x` | `mx.array` | `yes` | `none` | Required positional or keyword input. |

**Returns**

- Type: `mx.array`
- Direct return expressions: `mx.maximum(x, 0)`

**Exceptions and behavior**

Function `_relu` calls `mx.maximum`; returns `mx.maximum(x, 0)`.
No direct `raise` statement appears in this definition.

[View source #L230-L232](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/rerank_forward.py#L230-L232).

</details>

<details class="api-contract" id="contract-vllm_mlx.rerank_forward._silu" markdown="1">
<summary><code>vllm_mlx.rerank_forward._silu</code> · function</summary>

```python
vllm_mlx.rerank_forward._silu(x: mx.array) -> mx.array
```

SiLU/swish activation.

**Parameters**

| Name | Type | Required | Default | Description |
| --- | --- | --- | --- | --- |
| `x` | `mx.array` | `yes` | `none` | Required positional or keyword input. |

**Returns**

- Type: `mx.array`
- Direct return expressions: `x * mx.sigmoid(x)`

**Exceptions and behavior**

Function `_silu` calls `mx.sigmoid`; returns `x * mx.sigmoid(x)`.
No direct `raise` statement appears in this definition.

[View source #L235-L237](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/rerank_forward.py#L235-L237).

</details>

<details class="api-contract" id="contract-vllm_mlx.rerank_forward._apply_hidden_activation" markdown="1">
<summary><code>vllm_mlx.rerank_forward._apply_hidden_activation</code> · function</summary>

```python
vllm_mlx.rerank_forward._apply_hidden_activation(x: mx.array, config: dict) -> mx.array
```

Apply the configured encoder hidden activation.

**Parameters**

| Name | Type | Required | Default | Description |
| --- | --- | --- | --- | --- |
| `x` | `mx.array` | `yes` | `none` | Required positional or keyword input. |
| `config` | `dict` | `yes` | `none` | Required positional or keyword input. |

**Returns**

- Type: `mx.array`
- Direct return expressions: `_gelu(x)`; `_gelu_new(x)`; `_relu(x)`; `_silu(x)`

**Exceptions and behavior**

Function `_apply_hidden_activation` calls `config.get`, `isinstance`, `hidden_act.get`, `str(hidden_act).lower`; can raise `ValueError`; has 4 explicit return paths.
Directly raised exceptions: `ValueError`.

[View source #L240-L265](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/rerank_forward.py#L240-L265).

</details>

## Complete symbol map

This map also includes private definitions and nested helpers. The signature column exposes every explicit input even when an internal helper has no dedicated parameter prose.

| Symbol | Kind | Signature and inputs | What it does | Source |
| --- | --- | --- | --- | --- |
| [`classifier_forward`](#contract-vllm_mlx.rerank_forward.classifier_forward) | function | `classifier_forward(input_ids: mx.array, attention_mask: mx.array, weights: dict[str, mx.array], config: dict) -> mx.array` | Run a BERT-family classifier forward pass on MLX. | [#L16-L85](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/rerank_forward.py#L16-L85) |
| [`_position_ids_for_config`](#contract-vllm_mlx.rerank_forward._position_ids_for_config) | function | `_position_ids_for_config(config: dict, input_ids: mx.array, attention_mask: mx.array \| None) -> mx.array` | Build BERT or RoBERTa-family absolute position IDs. | [#L88-L105](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/rerank_forward.py#L88-L105) |
| [`_classification_head_forward`](#contract-vllm_mlx.rerank_forward._classification_head_forward) | function | `_classification_head_forward(pooled: mx.array, weights: dict[str, mx.array]) -> mx.array` | Run BERT flat or XLM-RoBERTa two-layer sequence-classification head. | [#L108-L122](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/rerank_forward.py#L108-L122) |
| [`_detect_prefix`](#contract-vllm_mlx.rerank_forward._detect_prefix) | function | `_detect_prefix(weights: dict) -> str` | Detect the model weight prefix (bert, roberta, xlm-roberta). | [#L125-L135](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/rerank_forward.py#L125-L135) |
| [`_layer_norm`](#contract-vllm_mlx.rerank_forward._layer_norm) | function | `_layer_norm(x: mx.array, weight: mx.array, bias: mx.array, eps: float) -> mx.array` | Apply layer normalization. | [#L138-L142](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/rerank_forward.py#L138-L142) |
| [`_encoder_layer`](#contract-vllm_mlx.rerank_forward._encoder_layer) | function | `_encoder_layer(hidden: mx.array, ext_mask: mx.array \| None, weights: dict, prefix: str, num_heads: int, head_dim: int, eps: float, config: dict) -> mx.array` | Run one BERT encoder layer (self-attention + FFN). | [#L145-L217](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/rerank_forward.py#L145-L217) |
| [`_gelu`](#contract-vllm_mlx.rerank_forward._gelu) | function | `_gelu(x: mx.array) -> mx.array` | GELU activation (exact form). | [#L220-L222](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/rerank_forward.py#L220-L222) |
| [`_gelu_new`](#contract-vllm_mlx.rerank_forward._gelu_new) | function | `_gelu_new(x: mx.array) -> mx.array` | BERT GELU approximation used by transformers gelu_new. | [#L225-L227](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/rerank_forward.py#L225-L227) |
| [`_relu`](#contract-vllm_mlx.rerank_forward._relu) | function | `_relu(x: mx.array) -> mx.array` | ReLU activation. | [#L230-L232](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/rerank_forward.py#L230-L232) |
| [`_silu`](#contract-vllm_mlx.rerank_forward._silu) | function | `_silu(x: mx.array) -> mx.array` | SiLU/swish activation. | [#L235-L237](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/rerank_forward.py#L235-L237) |
| [`_apply_hidden_activation`](#contract-vllm_mlx.rerank_forward._apply_hidden_activation) | function | `_apply_hidden_activation(x: mx.array, config: dict) -> mx.array` | Apply the configured encoder hidden activation. | [#L240-L265](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/rerank_forward.py#L240-L265) |