<jupyter_start><jupyter_text>Ablation StudyThis collects and shows the data from the ablation study on vectorization and compilation through JAX and Custom TFP.<jupyter_code>import pickle
import pandas as pd
import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt
pd.set_option('display.max_rows', 200)<jupyter_output><empty_output><jupyter_text>JAX<jupyter_code>expts = ['logreg', 'ffnn', 'mnist', 'embed', 'cifar10', 'lstm']
expt_dict = {}
for expt in expts:
    expt_dict[expt] = {}
    pickle_name = f'./raw/jaxdp_{expt}_bs_128_priv_True'
    try:
        with open(pickle_name+'.pkl', 'rb') as f:
            expt_dict[expt]['base'] = np.median(pickle.load(f)['timings'])
    except: print(f'Failed {expt} xla')
    try:
        with open(pickle_name+'_novmap.pkl', 'rb') as f:
            expt_dict[expt]['nv'] = np.median(pickle.load(f)['timings'])
    except: print(f'Failed {expt} xla')
    try:
        with open(pickle_name+'_nojit.pkl', 'rb') as f:
            expt_dict[expt]['nj'] = np.median(pickle.load(f)['timings'])
    except: print(f'Failed {expt} xla')
    try:
        with open(pickle_name+'_nojit_novmap.pkl', 'rb') as f:
            expt_dict[expt]['nvj'] = np.median(pickle.load(f)['timings'])
    except: print(f'Failed {expt} xla')
expt_dict<jupyter_output><empty_output><jupyter_text>Custom TFP<jupyter_code>expts = ['logreg', 'ffnn', 'mnist', 'embed', 'lstm', 'cifar10']
tf_expt_dict = {}
for expt in expts:
    tf_expt_dict[expt] = {}
    pickle_name = f'./raw/tf2dp_{expt}_bs_128_priv_True_xla_'
    try:
        with open(pickle_name+'False.pkl', 'rb') as f:
            tf_expt_dict[expt]['base'] = np.median(pickle.load(f)['timings'])
    except: print(f'Failed {expt} base')
    try:
        with open(pickle_name+'True.pkl', 'rb') as f:
            tf_expt_dict[expt]['xla'] = np.median(pickle.load(f)['timings'])
    except: print(f'Failed {expt} xla')
    try:
        with open(pickle_name+'False_novmap.pkl', 'rb') as f:
            tf_expt_dict[expt]['nv'] = np.median(pickle.load(f)['timings'])
    except: print(f'Failed {expt} nv')
    try:
        with open(pickle_name+'False_nojit_novmap.pkl', 'rb') as f:
            tf_expt_dict[expt]['nvj'] = np.median(pickle.load(f)['timings'])
    except: print(f'Failed {expt} nvj')
    try:
        with open(pickle_name+'True_novmap.pkl', 'rb') as f:
            tf_expt_dict[expt]['xla_nv'] = np.median(pickle.load(f)['timings'])
    except: print(f'Failed {expt} xla_nv')
tf_expt_dict<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.16805390408244153, 'fraction_numerical': 0.005548949663099485, 'mean_word_length': 4.007936507936508, 'pattern_counts': {'":': 0, '<': 13, '<?xml version=': 0, '>': 13, 'https://': 0, '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': '3', 'dir': 'jupyter-structured-clean-dedup', 'id': '370442', 'license': 'MIT', 'n_tokens_mistral': 1032, 'n_tokens_neox': 1012, 'n_words': 178, 'path': 'results/ablation.ipynb', 'repo_name': 'lxuechen/fast-dpsgd', 'size': '2523'}