<commit_before>from autograd import make_vjp

import autograd.numpy as np
import autograd.numpy.random as npr

dot_0 = lambda A, B, g: make_vjp(np.dot)(A, B)[0](g)
dot_1 = lambda A, B, g: make_vjp(np.dot, argnum=1)(A, B)[0](g)

dot_0_0 = lambda A, B, g: make_vjp(dot_0)(A, B, g)[0](A)
dot_0_1 = lambda A, B, g: make_vjp(dot_0)(A, B, g)[0](A)
dot_0_2 = lambda A, B, g: make_vjp(dot_0)(A, B, g)[0](A)

dot_1_0 = lambda A, B, g: make_vjp(dot_1)(A, B, g)[0](B)
dot_1_1 = lambda A, B, g: make_vjp(dot_1)(A, B, g)[0](B)
dot_1_2 = lambda A, B, g: make_vjp(dot_1)(A, B, g)[0](B)

A = npr.randn(2, 3, 4, 5)
B = npr.randn(2, 3, 5, 4)
g = npr.randn(2, 3, 4, 2, 3, 4)

def time_dot_0():
    dot_0(A, B, g)

def time_dot_1():
    dot_1(A, B, g)

def time_dot_0_0():
    dot_0_0(A, B, g)

def time_dot_0_1():
    dot_0_1(A, B, g)

def time_dot_0_2():
    dot_0_2(A, B, g)<commit_msg>Add time of dot primitive for comparison<commit_after>from autograd import make_vjp

import autograd.numpy as np
import autograd.numpy.random as npr

dot_0 = lambda A, B, g: make_vjp(np.dot)(A, B)[0](g)
dot_1 = lambda A, B, g: make_vjp(np.dot, argnum=1)(A, B)[0](g)

dot_0_0 = lambda A, B, g: make_vjp(dot_0)(A, B, g)[0](A)
dot_0_1 = lambda A, B, g: make_vjp(dot_0)(A, B, g)[0](A)
dot_0_2 = lambda A, B, g: make_vjp(dot_0)(A, B, g)[0](A)

dot_1_0 = lambda A, B, g: make_vjp(dot_1)(A, B, g)[0](B)
dot_1_1 = lambda A, B, g: make_vjp(dot_1)(A, B, g)[0](B)
dot_1_2 = lambda A, B, g: make_vjp(dot_1)(A, B, g)[0](B)

A = npr.randn(2, 3, 4, 5)
B = npr.randn(2, 3, 5, 4)
g = npr.randn(2, 3, 4, 2, 3, 4)

def time_dot():
    np.dot(A, B)

def time_dot_0():
    dot_0(A, B, g)

def time_dot_1():
    dot_1(A, B, g)

def time_dot_0_0():
    dot_0_0(A, B, g)

def time_dot_0_1():
    dot_0_1(A, B, g)

def time_dot_0_2():
    dot_0_2(A, B, g)
{'fraction_non_alphanumeric': 0.20177876598110062, 'fraction_numerical': 0.06559199555308505, 'mean_word_length': 4.187319884726225, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, '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}}
{'dir': 'git-commits-cleaned', 'id': '2464611', 'n_tokens_mistral': 1020, 'n_tokens_neox': 978, 'n_words': 243}