# Path: autodiff/context.py
# class Context(object):
#     def __init__(self, device=None, borrowable=(), force_floatX=False):
#         """
#         borrowable : tuple of objects
#             If an object in this tuple is encountered while tracing the
#             function, then its symbolic representation will alias that object's
#             memory location. This means that *inplace* operations on the Python
#             (likely NumPy) object will affect the symbolic function.
# 
#         force_floatX : bool
#             If True, floats and float NumPy ndarrays will be cast to the dtype
#             specified at theano.config.floatX when forming symbolic shared
#             variables, if they do not have it already. Objects in `borrowable`
#             are never cast.
# 
#         """
#         self.svars = {}
#         self.nogc = []  # ids that must not be reused
#         # XXX: rethink to avoid actually holding on to all these intermediates.
#         self.device = device
#         self.borrowable_ids = [id(b) for b in borrowable]
#         self.force_floatX = force_floatX
#         self.constants = set()
# 
#     def __iter__(self):
#         return self.svars.__iter__()
# 
#     def shadow(self, rval, sval, force=True):
#         assert hasattr(sval, 'type')  # assert sval is Theano variable
#         if force:
#             self.svars[id(rval)] = sval
#         else:
#             self.svars.setdefault(id(rval), sval)
# 
#         # -- shadow vars have to match dtype and ndim
#         if isinstance(rval, np.ndarray):
#             if str(rval.dtype) == 'bool':
#                 assert sval.dtype == 'int8', (rval.dtype, sval.dtype)
#             elif not self.force_floatX:
#                 assert str(rval.dtype) == sval.dtype, (rval, sval)
#             assert rval.ndim == sval.ndim, (rval, sval)
# 
#         # -- assert postcondition
#         assert sval is self.getvar(rval)
#         self.nogc.append(rval)
# 
#     def call(self, fn, args=(), kwargs={}):
#         vm = FrameVM(self, fn)
#         return vm.call(args, kwargs)
# 
#     def shared(self, obj, name=None, borrow=None):
#         if borrow is None:
#             borrow = (id(obj) in self.borrowable_ids)
#         if self.force_floatX and not borrow:
#             if (isinstance(obj, np.ndarray)
#                and 'float' in str(obj.dtype)
#                and str(obj.dtype) != theano.config.floatX):
#                 obj = obj.astype(theano.config.floatX)
# 
#         # not all objects have shared constructors with a borrow keyword
#         # for example theano.shared(np.float32(1)) works but
#         # theano.shared(np.float32(1), borrow=[False|True]) fails
#         if self.device == 'cpu':
#             try:
#                 return theano.tensor._shared(obj, borrow=borrow)
#             except:
#                 return theano.tensor._shared(obj)
#         else:
#             try:
#                 return theano.shared(obj, borrow=borrow)
#             except:
#                 return theano.shared(obj)
# 
#     def getvar(self, var):
#         return self.svars.get(id(var), var)
# 
#     def reset(self):
#         self.constants.clear()


import unittest
import numpy as np
import theano
from autodiff.context import Context
    def test_recalculate(self):
        x = np.zeros(3)
        c = Context()
        assert compute_stuff(x) == 12
        y = c.call(compute_stuff, (x,))
        assert y == 12

        f = recalculate_fn(c, y, x)
        assert f(x) == 12
        assert f(x + 1) != 12


    def test_grad(self):
        x = np.zeros(3)

        c = Context()
        y = c.call(compute_stuff, (x,))
        assert id(x) in c.svars
        assert id(y) in c.svars
        dy_dx_fn = grad_fn(c, y, x)

        assert np.all(dy_dx_fn(x + 0) == 8)
        assert np.all(dy_dx_fn(x + 1) == 16)
        assert np.all(dy_dx_fn(x + 2) == 24)


    def test_loop(self):
        # test that non-data-dependent loops are unrolled properly

        x = np.zeros(3)

# Path: autodiff/context.py
# class Context(object):
#     def __init__(self, device=None, borrowable=(), force_floatX=False):
#         """
#         borrowable : tuple of objects
#             If an object in this tuple is encountered while tracing the
#             function, then its symbolic representation will alias that object's
#             memory location. This means that *inplace* operations on the Python
#             (likely NumPy) object will affect the symbolic function.
# 
#         force_floatX : bool
#             If True, floats and float NumPy ndarrays will be cast to the dtype
#             specified at theano.config.floatX when forming symbolic shared
#             variables, if they do not have it already. Objects in `borrowable`
#             are never cast.
# 
#         """
#         self.svars = {}
#         self.nogc = []  # ids that must not be reused
#         # XXX: rethink to avoid actually holding on to all these intermediates.
#         self.device = device
#         self.borrowable_ids = [id(b) for b in borrowable]
#         self.force_floatX = force_floatX
#         self.constants = set()
# 
#     def __iter__(self):
#         return self.svars.__iter__()
# 
#     def shadow(self, rval, sval, force=True):
#         assert hasattr(sval, 'type')  # assert sval is Theano variable
#         if force:
#             self.svars[id(rval)] = sval
#         else:
#             self.svars.setdefault(id(rval), sval)
# 
#         # -- shadow vars have to match dtype and ndim
#         if isinstance(rval, np.ndarray):
#             if str(rval.dtype) == 'bool':
#                 assert sval.dtype == 'int8', (rval.dtype, sval.dtype)
#             elif not self.force_floatX:
#                 assert str(rval.dtype) == sval.dtype, (rval, sval)
#             assert rval.ndim == sval.ndim, (rval, sval)
# 
#         # -- assert postcondition
#         assert sval is self.getvar(rval)
#         self.nogc.append(rval)
# 
#     def call(self, fn, args=(), kwargs={}):
#         vm = FrameVM(self, fn)
#         return vm.call(args, kwargs)
# 
#     def shared(self, obj, name=None, borrow=None):
#         if borrow is None:
#             borrow = (id(obj) in self.borrowable_ids)
#         if self.force_floatX and not borrow:
#             if (isinstance(obj, np.ndarray)
#                and 'float' in str(obj.dtype)
#                and str(obj.dtype) != theano.config.floatX):
#                 obj = obj.astype(theano.config.floatX)
# 
#         # not all objects have shared constructors with a borrow keyword
#         # for example theano.shared(np.float32(1)) works but
#         # theano.shared(np.float32(1), borrow=[False|True]) fails
#         if self.device == 'cpu':
#             try:
#                 return theano.tensor._shared(obj, borrow=borrow)
#             except:
#                 return theano.tensor._shared(obj)
#         else:
#             try:
#                 return theano.shared(obj, borrow=borrow)
#             except:
#                 return theano.shared(obj)
# 
#     def getvar(self, var):
#         return self.svars.get(id(var), var)
# 
#     def reset(self):
#         self.constants.clear()


# Path: autodiff/tests/test_context.py
import unittest
import numpy as np
import theano
from autodiff.context import Context

    def test_recalculate(self):
        x = np.zeros(3)
        c = Context()
        assert compute_stuff(x) == 12
        y = c.call(compute_stuff, (x,))
        assert y == 12

        f = recalculate_fn(c, y, x)
        assert f(x) == 12
        assert f(x + 1) != 12


    def test_grad(self):
        x = np.zeros(3)

        c = Context()
        y = c.call(compute_stuff, (x,))
        assert id(x) in c.svars
        assert id(y) in c.svars
        dy_dx_fn = grad_fn(c, y, x)

        assert np.all(dy_dx_fn(x + 0) == 8)
        assert np.all(dy_dx_fn(x + 1) == 16)
        assert np.all(dy_dx_fn(x + 2) == 24)


    def test_loop(self):
        # test that non-data-dependent loops are unrolled properly

        x = np.zeros(3)

        c = Context()