<commit_before>from __future__ import absolute_import
from copy import copy
import numpy as numpy_original
from autograd.core import (Node, primitive, zeros_like,
                           differentiable_ops, nondifferentiable_ops)
from . import numpy_wrapper as anp

@primitive
def take(A, idx):
    return A[idx]
def make_grad_take(ans, A, idx):
    return lambda g : untake(g, idx, A)
take.defgrad(make_grad_take)

@primitive
def untake(x, idx, template):
    return SparseArray(template, idx, x)
untake.defgrad(lambda ans, x, idx, template : lambda g : take(g, idx))
untake.defgrad_is_zero(argnums=(1, 2))

class ArrayNode(Node):
    __slots__ = []
    __getitem__ = take
    __array_priority__ = 100.0

    def __init__(self, value):
        if type(Node) is numpy_original.ndarray:
            value = value.view(anp.ndarray)
        self.value = value
        self.tapes = {}

    # Constants w.r.t float data just pass though
    shape = property(lambda self: self.value.shape)
    ndim  = property(lambda self: self.value.ndim)
    size  = property(lambda self: self.value.size)
    T = property(lambda self: anp.transpose(self))

    @staticmethod
    def zeros_like(value):
        return anp.zeros(value.shape)

Node.type_mappings[anp.ndarray] = ArrayNode
Node.type_mappings[numpy_original.ndarray] = ArrayNode

# Binary ops already wrapped by autograd.numpy.ndarray
inherited_methods = ['dot'] + differentiable_ops +  nondifferentiable_ops
for method_name in inherited_methods:
    setattr(ArrayNode, method_name, getattr(anp.ndarray, method_name).__func__)

# These numpy.ndarray methods are just refs to an equivalent numpy function
nondiff_methods = ['all', 'any', 'argmax', 'argmin', 'argpartition',
                   'argsort', 'nonzero', 'searchsorted', 'round']
diff_methods = ['clip', 'compress', 'cumprod', 'cumsum', 'diagonal',
                'max', 'mean', 'min', 'prod', 'ptp', 'ravel', 'repeat',
                'reshape', 'squeeze', 'std', 'sum', 'swapaxes', 'take',
                'trace', 'transpose', 'var']<commit_msg>Removed array type casting in ArrayNode (shouldn't be needed any more)<commit_after>from __future__ import absolute_import
from copy import copy
import numpy as numpy_original
from autograd.core import (Node, primitive, zeros_like,
                           differentiable_ops, nondifferentiable_ops)
from . import numpy_wrapper as anp

@primitive
def take(A, idx):
    return A[idx]
def make_grad_take(ans, A, idx):
    return lambda g : untake(g, idx, A)
take.defgrad(make_grad_take)

@primitive
def untake(x, idx, template):
    return SparseArray(template, idx, x)
untake.defgrad(lambda ans, x, idx, template : lambda g : take(g, idx))
untake.defgrad_is_zero(argnums=(1, 2))

class ArrayNode(Node):
    __slots__ = []
    __getitem__ = take
    __array_priority__ = 100.0

    # Constants w.r.t float data just pass though
    shape = property(lambda self: self.value.shape)
    ndim  = property(lambda self: self.value.ndim)
    size  = property(lambda self: self.value.size)
    T = property(lambda self: anp.transpose(self))

    @staticmethod
    def zeros_like(value):
        return anp.zeros(value.shape)

Node.type_mappings[anp.ndarray] = ArrayNode
Node.type_mappings[numpy_original.ndarray] = ArrayNode

# Binary ops already wrapped by autograd.numpy.ndarray
inherited_methods = ['dot'] + differentiable_ops +  nondifferentiable_ops
for method_name in inherited_methods:
    setattr(ArrayNode, method_name, getattr(anp.ndarray, method_name).__func__)

# These numpy.ndarray methods are just refs to an equivalent numpy function
nondiff_methods = ['all', 'any', 'argmax', 'argmin', 'argpartition',
                   'argsort', 'nonzero', 'searchsorted', 'round']
diff_methods = ['clip', 'compress', 'cumprod', 'cumsum', 'diagonal',
                'max', 'mean', 'min', 'prod', 'ptp', 'ravel', 'repeat',
                'reshape', 'squeeze', 'std', 'sum', 'swapaxes', 'take',
                'trace', 'transpose', 'var']
{'fraction_non_alphanumeric': 0.12089290193127665, 'fraction_numerical': 0.0030097817908201654, 'mean_word_length': 3.8106151990349817, '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': '3924506', 'n_tokens_mistral': 1317, 'n_tokens_neox': 1271, 'n_words': 345}