Flax Reproduction

---
name: flax-reproduction
description: "Routes Distil-Whisper JAX/Flax reproduction, long-form evaluation,
  and conversion workflows."
disable-model-invocation: true
metadata:
  disco-role: operating
license: MIT
---

# Flax Reproduction

## Purpose

Use this sub-skill when the user wants the original JAX/Flax Distil-Whisper package, the Flax pipeline, long-form transcription, checkpoint conversion, or the TPU/GPU-oriented reproduction flow under `training/flax/`.

## Include here

- The `distil_whisper` package under `training/flax/distil_whisper`.
- Flax student initialization.
- Flax evaluation and long-form transcription.
- Checkpoint conversion from a Flax training state to Hugging Face weights.
- JAX/Flax workflow notes, including scan weights, TPU-oriented settings, and the Flax pipeline.

## Exclude or route elsewhere

- Direct Transformers inference belongs in `inference`.
- PyTorch pseudo-labelling, training, and evaluation belong in `pytorch-training`.
- Benchmark-only or hardcoded local debug scripts are not the default route unless the user explicitly asks about them.

## Read next

- `references/workflows.md` for Flax initialization, evaluation, long-form, and conversion recipes.
- `references/troubleshooting.md` for JAX, SciPy, NumPy, and distributed-init failures.
- `../../references/model-overview.md` for checkpoint selection context.
- `../../scripts/check-env.py` before importing the package or launching a script.
- `scripts/create_student_model.py` for a bundled initialization helper.
- `scripts/convert_train_state_to_hf.py` for the bundled conversion/export helper.

## How to route

- "How do I reproduce the original Flax workflow?" -> start here.
- "How do I use FlaxWhisperPipeline?" -> start here.
- "How do I convert a JAX training state to HF weights?" -> start here.
- "How do I run long-form transcription in the Flax stack?" -> start here.