Posts by Year

MLP Language Model with JAX

6 minute read

Bigram Language Model with JAX

13 minute read

Introduction to basic probability.

2 minute read

Let us start with the notion of a sample space. A sample space describes the set of all possible outcomes when performing an experiment.

Permutations and Combinations.

3 minute read

An ordered set of elements in [n]={1,2,...,n} is called a permutation.

Introduction to INT8 quantization in JAX

3 minute read

Introduction In this blog post, we will provide an intuitive introduction to a technique called quantization. The topic has always seemed a bit mysterious to...

High performance Inference on TPUs using Maxtext

3 minute read

Attention module in JAX

9 minute read

The attention module is the key ingredient of what makes up a transformer layer. In this blogpost we will show how to implement it from scratch in JAX alongs...

Multi chip performance in JAX

4 minute read

The larger the models we use get the more it becomes necessary to be able to perform training of machine learning models over multiple chips. In this blog po...