<commit_before>from __future__ import absolute_import, print_function
import sys
import warnings
import copy
import operator as op
import types
import math
import re
import numpy as np
from functools import partial
from future.utils import iteritems, raise_from, raise_
from collections import defaultdict

def grad(fun, argnum=0):
    """
    Returns a function which computes the gradient of `fun` with respect to
    positional argument number `argnum`. The returned function takes the same
    arguments as `fun`, but returns the gradient instead. The function `fun`
    should be scalar-valued. The gradient has the same type as the argument."""
    @attach_name_and_doc(fun, argnum, 'Gradient')
    def gradfun(*args,**kwargs):
        return backward_pass(*forward_pass(fun,args,kwargs,argnum))
    return gradfun

def forward_pass(fun, args, kwargs, argnum=0):
    tape = CalculationTape()
    arg_wrt = args[argnum]
    start_node = new_node(safe_type(getval(arg_wrt)), [tape])
    args = list(args)
    args[argnum] = merge_tapes(start_node, arg_wrt)
    try: end_node = fun(*args, **kwargs)
    except Exception as e: add_extra_error_message(e)
    return start_node, end_node, tape

def backward_pass(start_node, end_node, tape):
    if not isinstance(end_node, Node) or tape not in end_node.tapes:
        warnings.warn("Output seems independent of input. Returning zero gradient.")
        return zeros_like(start_node)
    if type(end_node) is not FloatNode:
        try:
            end_node = FloatNode.cast(end_node, 1.0)
        except TypeError:
            raise TypeError(
                "Output type {} can't be cast to float. "
                "Function grad requires a scalar-valued function. "
                "Try jacobian or elementwise_grad.".format(type(end_node.value)))

    for node in tape:
        node.outgrads = []
    end_node.tapes[tape].outgrads = [1.0]

    tape.complete = True
    tape = copy.copy(tape)
    while tape:
        node = tape.pop()
        if node.outgrads:
            cur_outgrad = node.sum_outgrads()
            assert type(new_node(getval(cur_outgrad))) == node.node_type, \
                "Types are {0} and {1}".format(type(new_node(getval(cur_outgrad))), node.node_type)
            for gradfun, parent in node.parent_grad_ops:
                og = cast_to_node_type(gradfun(cur_outgrad), parent.node_type, parent.node_value)
                parent.outgrads.append(og)
    return cur_outgrad

def attach_name_and_doc(fun, argnum, opname):
    namestr = "{op}_{fun}_wrt_argnum_{argnum}".format(
        op=opname.lower(), fun=fun.__name__, argnum=argnum)
    docstr = "{op} of function {fun} with respect to argument number {argnum}. " \
        "Has the same arguments as {fun} but the return value has type of" \
        "argument {argnum}".format(op=opname, fun=fun.__name__, argnum=argnum)

    def wrap(gradfun):
        try:
            gradfun.__name__ = namestr<commit_msg>fix memory leak by not copying tape in backward_pass<commit_after>from __future__ import absolute_import, print_function
import sys
import warnings
import copy
import operator as op
import types
import math
import re
import numpy as np
from functools import partial
from future.utils import iteritems, raise_from, raise_
from collections import defaultdict

def grad(fun, argnum=0):
    """
    Returns a function which computes the gradient of `fun` with respect to
    positional argument number `argnum`. The returned function takes the same
    arguments as `fun`, but returns the gradient instead. The function `fun`
    should be scalar-valued. The gradient has the same type as the argument."""
    @attach_name_and_doc(fun, argnum, 'Gradient')
    def gradfun(*args,**kwargs):
        return backward_pass(*forward_pass(fun,args,kwargs,argnum))
    return gradfun

def forward_pass(fun, args, kwargs, argnum=0):
    tape = CalculationTape()
    arg_wrt = args[argnum]
    start_node = new_node(safe_type(getval(arg_wrt)), [tape])
    args = list(args)
    args[argnum] = merge_tapes(start_node, arg_wrt)
    try: end_node = fun(*args, **kwargs)
    except Exception as e: add_extra_error_message(e)
    return start_node, end_node, tape

def backward_pass(start_node, end_node, tape):
    if not isinstance(end_node, Node) or tape not in end_node.tapes:
        warnings.warn("Output seems independent of input. Returning zero gradient.")
        return zeros_like(start_node)
    if type(end_node) is not FloatNode:
        try:
            end_node = FloatNode.cast(end_node, 1.0)
        except TypeError:
            raise TypeError(
                "Output type {} can't be cast to float. "
                "Function grad requires a scalar-valued function. "
                "Try jacobian or elementwise_grad.".format(type(end_node.value)))

    for node in tape:
        node.outgrads = []
    end_node.tapes[tape].outgrads = [1.0]

    tape.complete = True
    for node in tape[::-1]:
        if node.outgrads:
            cur_outgrad = node.sum_outgrads()
            assert type(new_node(getval(cur_outgrad))) == node.node_type, \
                "Types are {0} and {1}".format(type(new_node(getval(cur_outgrad))), node.node_type)
            for gradfun, parent in node.parent_grad_ops:
                og = cast_to_node_type(gradfun(cur_outgrad), parent.node_type, parent.node_value)
                parent.outgrads.append(og)
    return cur_outgrad

def attach_name_and_doc(fun, argnum, opname):
    namestr = "{op}_{fun}_wrt_argnum_{argnum}".format(
        op=opname.lower(), fun=fun.__name__, argnum=argnum)
    docstr = "{op} of function {fun} with respect to argument number {argnum}. " \
        "Has the same arguments as {fun} but the return value has type of" \
        "argument {argnum}".format(op=opname, fun=fun.__name__, argnum=argnum)

    def wrap(gradfun):
        try:
            gradfun.__name__ = namestr
{'fraction_non_alphanumeric': 0.09623217922606925, 'fraction_numerical': 0.002885268160217244, 'mean_word_length': 3.1676096181046676, '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': '3343539', 'n_tokens_mistral': 1829, 'n_tokens_neox': 1804, 'n_words': 527}