raise NotImplementedError
def flatten(self, value):
return array_ops.reshape(value, tensor.Tensor(-1))
def unflatten(self, value):
return array_ops.reshape(value, tensor.Tensor(self.shape))
def mut_add(self, x, y):
"""Add wrapper safe for IndexedSlices and LazyZero."""
if isinstance(ag_core.getval(x), tensor.LazyZero):
return y
if isinstance(ag_core.getval(y), tensor.LazyZero):
return x
if isinstance(x, ops.IndexedSlices):
x = _indexed_slices_to_tensor(x)
if isinstance(y, ops.IndexedSlices):
y = _indexed_slices_to_tensor(y)
if x is None:
return y
if y is None:
return x
return math_ops.add(x, y)
ag_core.register_vspace(TensorVSpace, tensor.Tensor)
ag_core.register_vspace(TensorVSpace, ops.Tensor)
ag_core.register_vspace(TensorVSpace, ops.IndexedSlices)
ag_core.register_vspace(TensorVSpace, tensor.LazyZero)
ag_core.register_node(TensorNode, tensor.LazyZero)
def zeros(self):
return ImplicitTape()
class ImplicitTapeNode(ag_core.Node):
"""Node to wrap ImplicitTape in."""
def __eq__(self, other):
return self is other
def __hash__(self):
return id(self)
ag_core.register_node(ImplicitTapeNode, ImplicitTape)
ag_core.register_vspace(ImplicitTapeVSpace, ImplicitTape)
# TODO(apassos) try to not do this.
class NoneVSpace(ag_core.VSpace):
"""VSpace for python None."""
def __init__(self, _):
self.size = 0
def zeros(self):
return 0
ag_core.register_vspace(NoneVSpace, type(None))
def flatten(self, value, covector=False):
if covector:
return anp.ravel(anp.stack([anp.real(value), - anp.imag(value)]))
else:
return anp.ravel(anp.stack([anp.real(value), anp.imag(value)]))
def unflatten(self, value, covector=False):
reshaped = anp.reshape(value, (2,) + self.shape)
if covector:
return anp.array(reshaped[0] - 1j * reshaped[1])
else:
return anp.array(reshaped[0] + 1j * reshaped[1])
register_node(ArrayNode, np.ndarray)
register_vspace(lambda x: ComplexArrayVSpace(x)
if np.iscomplexobj(x)
else ArrayVSpace(x), np.ndarray)
array_types = set([anp.ndarray, ArrayNode])
for type_ in [float, anp.float64, anp.float32, anp.float16]:
register_node(ArrayNode, type_)
register_vspace(ArrayVSpace, type_)
for type_ in [complex, anp.complex64, anp.complex128]:
register_node(ArrayNode, type_)
register_vspace(ComplexArrayVSpace, type_)
# 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',
return np.concatenate(
[vs.flatten(v, covector) for vs, v in zip(self.shape, value)])
else:
return np.zeros((0, ))
def unflatten(self, value, covector=False):
result = []
start = 0
for vs in self.shape:
N = vs.size
result.append(vs.unflatten(value[start:start + N], covector))
start += N
return self.sequence_type(result)
register_vspace(SequenceVSpace, list)
register_vspace(SequenceVSpace, tuple)
class DictNode(Node):
__slots__ = []
def __getitem__(self, idx):
return dict_take(self, idx)
def __len__(self):
return len(self.value)
def __iter__(self):
return self.value.__iter__()
def zeros(self):
return ImplicitTape()
class ImplicitTapeNode(ag_core.Node):
"""Node to wrap ImplicitTape in."""
def __eq__(self, other):
return self is other
def __hash__(self):
return id(self)
ag_core.register_node(ImplicitTapeNode, ImplicitTape)
ag_core.register_vspace(ImplicitTapeVSpace, ImplicitTape)
# TODO(apassos) try to not do this.
class NoneVSpace(ag_core.VSpace):
"""VSpace for python None."""
def __init__(self, _):
self.size = 0
ag_core.register_vspace(NoneVSpace, type(None))
class _TapeStack(threading.local):
if covector:
return np.ravel(np.stack([np.real(value), -np.imag(value)]))
else:
return np.ravel(np.stack([np.real(value), np.imag(value)]))
def unflatten(self, value, covector=False):
reshaped = np.reshape(value, (2, ) + self.shape)
if covector:
return np.array(reshaped[0] - 1j * reshaped[1])
else:
return np.array(reshaped[0] + 1j * reshaped[1])
register_node(ArrayNode, np.ndarray)
register_vspace(
lambda x: ComplexArrayVSpace(x)
if np.iscomplexobj(x) else ArrayVSpace(x), np.ndarray)
array_types = set([anp.ndarray, ArrayNode])
for type_ in [float, anp.float64, anp.float32, anp.float16]:
register_node(ArrayNode, type_)
register_vspace(ArrayVSpace, type_)
for type_ in [complex, anp.complex64, anp.complex128]:
register_node(ArrayNode, type_)
register_vspace(ComplexArrayVSpace, type_)
# These numpy.ndarray methods are just refs to an equivalent numpy function
nondiff_methods = [
'all', 'any', 'argmax', 'argmin', 'argpartition', 'argsort', 'nonzero',
'searchsorted', 'round'
raise NotImplementedError
def flatten(self, value):
return array_ops.reshape(value, tensor.Tensor(-1))
def unflatten(self, value):
return array_ops.reshape(value, tensor.Tensor(self.shape))
def mut_add(self, x, y):
"""Add wrapper safe for IndexedSlices and LazyZero."""
if isinstance(ag_core.getval(x), tensor.LazyZero):
return y
if isinstance(ag_core.getval(y), tensor.LazyZero):
return x
if isinstance(x, ops.IndexedSlices):
x = _indexed_slices_to_tensor(x)
if isinstance(y, ops.IndexedSlices):
y = _indexed_slices_to_tensor(y)
if x is None:
return y
if y is None:
return x
return math_ops.add(x, y)
ag_core.register_vspace(TensorVSpace, tensor.Tensor)
ag_core.register_vspace(TensorVSpace, ops.Tensor)
ag_core.register_vspace(TensorVSpace, ops.IndexedSlices)
ag_core.register_vspace(TensorVSpace, tensor.LazyZero)
ag_core.register_node(TensorNode, tensor.LazyZero)
if self.shape:
return np.concatenate(
[vs.flatten(v, covector) for vs, v in zip(self.shape, value)])
else:
return np.zeros((0,))
def unflatten(self, value, covector=False):
result = []
start = 0
for vs in self.shape:
N = vs.size
result.append(vs.unflatten(value[start:start + N], covector))
start += N
return self.sequence_type(result)
register_vspace(SequenceVSpace, list)
register_vspace(SequenceVSpace, tuple)
class DictNode(Node):
__slots__ = []
def __getitem__(self, idx): return dict_take(self, idx)
def __len__(self): return len(self.value)
def __iter__(self): return self.value.__iter__()
def items(self): return list(self.iteritems())
def keys(self): return list(self.iterkeys())
def values(self): return list(self.itervalues())
def iteritems(self): return ((k, self[k]) for k in self)
def iterkeys(self): return iter(self)
def itervalues(self): return (self[k] for k in self)
register_node(DictNode, dict)
