# {py:mod}`dockyard_rl.models.jax.moe.router`

```{py:module} dockyard_rl.models.jax.moe.router
```

```{autodoc2-docstring} dockyard_rl.models.jax.moe.router
:allowtitles:
```

## Module Contents

### Classes

````{list-table}
:class: autosummary longtable
:align: left

* - {py:obj}`TokenChoiceTopKRouter <dockyard_rl.models.jax.moe.router.TokenChoiceTopKRouter>`
  - ```{autodoc2-docstring} dockyard_rl.models.jax.moe.router.TokenChoiceTopKRouter
    :summary:
    ```
````

### Data

````{list-table}
:class: autosummary longtable
:align: left

* - {py:obj}`Array <dockyard_rl.models.jax.moe.router.Array>`
  - ```{autodoc2-docstring} dockyard_rl.models.jax.moe.router.Array
    :summary:
    ```
* - {py:obj}`ScoreFunc <dockyard_rl.models.jax.moe.router.ScoreFunc>`
  - ```{autodoc2-docstring} dockyard_rl.models.jax.moe.router.ScoreFunc
    :summary:
    ```
````

### API

````{py:data} Array
:canonical: dockyard_rl.models.jax.moe.router.Array
:value: >
   None

```{autodoc2-docstring} dockyard_rl.models.jax.moe.router.Array
```

````

````{py:data} ScoreFunc
:canonical: dockyard_rl.models.jax.moe.router.ScoreFunc
:value: >
   None

```{autodoc2-docstring} dockyard_rl.models.jax.moe.router.ScoreFunc
```

````

````{py:class} TokenChoiceTopKRouter(dim: int, num_experts: int, *, top_k: int = 1, score_func: dockyard_rl.models.jax.moe.router.ScoreFunc = 'sigmoid', route_norm: bool = False, route_scale: float = 1.0, num_expert_groups: typing.Optional[int] = None, num_limited_groups: typing.Optional[int] = None, rngs: flax.nnx.Rngs, param_dtype: jax.numpy.dtype = jnp.float32)
:canonical: dockyard_rl.models.jax.moe.router.TokenChoiceTopKRouter

Bases: {py:obj}`flax.nnx.Module`

```{autodoc2-docstring} dockyard_rl.models.jax.moe.router.TokenChoiceTopKRouter
```

```{rubric} Initialization
```

```{autodoc2-docstring} dockyard_rl.models.jax.moe.router.TokenChoiceTopKRouter.__init__
```

````