<jupyter_start><jupyter_text><jupyter_code>import torchvision
from torchvision import datasets
from torchvision import transforms

import numpy as np
import jax
import jax.numpy as jnp
import itertools

from bokeh.io import output_notebook, show
from bokeh.layouts import gridplot
from bokeh.plotting import figure<jupyter_output><empty_output><jupyter_text>This is required for Bokeh to work in notebooks.<jupyter_code>output_notebook()<jupyter_output><empty_output><jupyter_text>According to [NIST](https://www.nist.gov/itl/products-and-services/emnist-dataset),> The EMNIST dataset is a set of handwritten character digits derived from the [NIST Special Database 19](https://www.nist.gov/srd/nist-special-database-19) and converted to a 28x28 pixel image format and dataset structure that directly matches the [MNIST dataset](http://yann.lecun.com/exdb/mnist/). Further information on the dataset contents and conversion process can be found in the paper available at https://arxiv.org/abs/1702.05373v1.Since we are going to work with JAX, let's transform the PyTorch Tensors to JAX DeviceArrays.<jupyter_code>transform=transforms.Compose([lambda img: torchvision.transforms.functional.rotate(img, 90),
                              transforms.ToTensor(),
                              jnp.array])

training_data = datasets.EMNIST(
    root="~/data",
    split="byclass",
    download=True,
    transform=transform
)<jupyter_output>Downloading https://www.itl.nist.gov/iaui/vip/cs_links/EMNIST/gzip.zip to /root/data/EMNIST/raw/gzip.zip<jupyter_text>Here are some examples from the dataset<jupyter_code>def plot_one(item):
    image, raw_label = item
    label = training_data.classes[raw_label]

    p = figure(title=f"label = {label}",
               tooltips=[("x", "$x"), ("y", "$y"), ("value", "@image")],
               match_aspect=True)
    p.x_range.range_padding = p.y_range.range_padding = 0

    # must give a vector of image data for image parameter
    subplot = p.image(image=[np.array(image.squeeze())],
                      x=0, y=0, dw=1, dh=1, level="image")
    
    p.title.align = "center"
    p.axis.visible = False
    p.grid.grid_line_width = 0.5

    return p

# Take the first 25 images
subplots = list(map(plot_one, itertools.islice(training_data, 25)))

grid = gridplot(subplots, ncols=5,
                toolbar_location=None,
                plot_width=150, plot_height=150)

show(grid)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.10893512851897184, 'fraction_numerical': 0.015503875968992248, 'mean_word_length': 3.874751491053678, 'pattern_counts': {'":': 0, '<': 16, '<?xml version=': 0, '>': 17, 'https://': 4, 'lorem ipsum': 0, 'www.': 3, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '4', 'dir': 'jupyter-structured-clean-dedup', 'id': '289320', 'license': 'MIT', 'n_tokens_mistral': 830, 'n_tokens_neox': 786, 'n_words': 193, 'path': 'notebooks/emnist_viz_jax.ipynb', 'repo_name': 'patel-zeel/probml-notebooks', 'size': '2451'}