Craft & Insight
A space for thoughts and insights on technology, engineering, and creation.
Recent Articles
View AllDiscrete Language Diffusion: Global Refinement on JAX/TPU
A deep dive into implementing Discrete Masked Language Diffusion using JAX and Flax NNX. We explore transition schedules, bidirectional architectures, and performance optimization for TPU v6e.
LLMJAXTPU
2026-04-30
Read More
How Do GPUs Perform Machine Learning Computations?
Explore the principles of hardware acceleration from Python code to GPU transistors through JAX and CUDA.
GPUMachine LearningCUDA
2025-12-05
Read More
Pre-training Decoder-based Tiny LLM with JAX and TPU
We dissect the entire process from raw text data being read from disk, tokenized, and reborn as meaningful sentences through hardware called TPU. Let's implement the design of the latest Llama model directly with JAX and transform from a user of the model to a designer of the model.
LLMJAXTPU
2025-12-02
Read More