<jupyter_start><jupyter_text><jupyter_code># Install Flax and JAX
!pip install --upgrade -q "jax[cuda11_cudnn805]" -f https://storage.googleapis.com/jax-releases/jax_releases.html
!pip install --upgrade -q git+https://github.com/google/flax.git
!pip install --upgrade -q git+https://github.com/deepmind/dm-haiku  # Haiku is here just for comparison purposes
import jax
from jax import lax, random, numpy as jnp

# NN lib built on top of JAX developed by Google Research (Brain team)
# Flax was "designed for flexibility" hence the name (Flexibility + JAX -> Flax)
import flax
from flax.core import freeze, unfreeze
from flax import linen as nn  # nn notation also used in PyTorch and in Flax's older API
from flax.training import train_state  # a useful dataclass to keep train state

# DeepMind's NN JAX lib - just for comparison purposes, we're not learning Haiku here
import haiku as hk 

# JAX optimizers - a separate lib developed by DeepMind
import optax

# Flax doesn't have its own data loading functions - we'll be using PyTorch dataloaders
from torchvision.datasets import MNIST
from torch.utils.data import DataLoader

# Python libs
import functools  # useful utilities for functional programs
from typing import Any, Callable, Sequence, Optional

# Other important 3rd party libs
import numpy as np
import matplotlib.pyplot as plt
# init a new base model
model = nn.Dense(features=5)
nn.Dense.__bases__
# init random PRNG
k1, k2 = jax.random.split(jax.random.PRNGKey(2022))
x = jax.random.normal(k1, (10,))  # random (10, 1) dataset

y, params = model.init_with_output(k2, x)  # or use init() to remove output y
jax.tree_map(lambda x: x.shape, params)  # kernel is weight, biases are init and inferred shape of input

# apply feedforward
y = model.apply(params, x)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.08250276854928018, 'fraction_numerical': 0.01107419712070875, 'mean_word_length': 4.594427244582043, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 6, 'https://': 3, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '308162', 'license': 'BSD-3-Clause', 'n_tokens_mistral': 608, 'n_tokens_neox': 567, 'n_words': 227, 'path': 'intro_to_flax.ipynb', 'repo_name': 'kbrezinski/JAX-Practice', 'size': '1806'}