from __future__ import absolute_import as _abs

import os

import sys

if sys.version_info[0] == 3:

import builtins as __builtin__

else:

import __builtin__

curr_path = os.path.dirname(os.path.abspath(os.path.expanduser(__file__)))

if hasattr(__builtin__, "NNVM_BASE_PATH"):

assert __builtin__.NNVM_BASE_PATH == curr_path

else:

__builtin__.NNVM_BASE_PATH = curr_path

if hasattr(__builtin__, "NNVM_LIBRARY_NAME"):

assert __builtin__.NNVM_LIBRARY_NAME == curr_path

else:

__builtin__.NNVM_LIBRARY_NAME = "libtinyflow"

import ctypes as _ctypes

from nnvm.name import NameManager

from nnvm._base import c_str, check_call, _LIB

from nnvm import symbol, graph

from nnvm import _symbol_internal

__all__ = ["float32", "placeholder", "Variable", "group",

"initialize_all_variables", "gradients"]

# data type table

float32 = 0

# global list of all variable initializers

_all_variable_inits = []

def Variable(init=None, name=None):

name = NameManager.current.get(name, 'variable')

v = symbol.Variable(name)

if init is not None:

if not isinstance(init, symbol.Symbol):

raise TypeError("Expect initialization expression to be Symbol")

_all_variable_inits.append(symbol.assign(v, init))

return v

def initialize_all_variables():

global _all_variable_inits

init_op = group(*_all_variable_inits)

_all_variable_inits = []

return init_op

def placeholder(dtype, shape=None, name=None):

v = symbol.placeholder(name=name, dtype=dtype)

return v

def group(*inputs):

x = _symbol_internal._nop()

x._add_control_deps(symbol.Group(inputs))

return x

def gradients(ys, xs, grad_ys=None):

if isinstance(ys, list):

ys = symbol.Group(ys)

g = graph.create(ys)

g._set_symbol_list_attr('grad_ys', ys)

g._set_symbol_list_attr('grad_xs', xs)

ny = len(ys.list_output_names())

if grad_ys is None:

grad_ys = [symbol.ones_like(ys[i]) for i in range(ny)]

g._set_symbol_list_attr('grad_ys_out_grad', grad_ys)

sym = g.apply('Gradient').symbol

nx = len(xs) if isinstance(xs, list) else len(xs.list_output_names())

ret = [sym[i] for i in range(nx)]

return ret