https://github.com/pymc-devs/pymc3
Tip revision: d97fda83a25f4692dbdd25988df9f224b5a81c3e authored by Chris Fonnesbeck on 25 January 2018, 00:56:22 UTC
Merge pull request #2814 from pymc-devs/3.3release
Merge pull request #2814 from pymc-devs/3.3release
Tip revision: d97fda8
theanof.py
import numpy as np
import theano
from theano import theano, scalar, tensor as tt
from theano.configparser import change_flags
from theano.gof import Op
from theano.gof.graph import inputs
from theano.sandbox.rng_mrg import MRG_RandomStreams
from .blocking import ArrayOrdering
from .data import GeneratorAdapter
from .vartypes import typefilter, continuous_types
__all__ = ['gradient',
'hessian',
'hessian_diag',
'inputvars',
'cont_inputs',
'floatX',
'smartfloatX',
'jacobian',
'CallableTensor',
'join_nonshared_inputs',
'make_shared_replacements',
'generator',
'set_tt_rng',
'tt_rng']
def inputvars(a):
"""
Get the inputs into a theano variables
Parameters
----------
a : theano variable
Returns
-------
r : list of tensor variables that are inputs
"""
return [v for v in inputs(makeiter(a)) if isinstance(v, tt.TensorVariable)]
def cont_inputs(f):
"""
Get the continuous inputs into a theano variables
Parameters
----------
a : theano variable
Returns
-------
r : list of tensor variables that are continuous inputs
"""
return typefilter(inputvars(f), continuous_types)
def floatX(X):
"""
Convert a theano tensor or numpy array to theano.config.floatX type.
"""
try:
return X.astype(theano.config.floatX)
except AttributeError:
# Scalar passed
return np.asarray(X, dtype=theano.config.floatX)
def smartfloatX(x):
"""
Convert non int types to floatX
"""
if str(x.dtype).startswith('float'):
x = floatX(x)
return x
"""
Theano derivative functions
"""
def gradient1(f, v):
"""flat gradient of f wrt v"""
return tt.flatten(tt.grad(f, v, disconnected_inputs='warn'))
empty_gradient = tt.zeros(0, dtype='float32')
def gradient(f, vars=None):
if vars is None:
vars = cont_inputs(f)
if vars:
return tt.concatenate([gradient1(f, v) for v in vars], axis=0)
else:
return empty_gradient
def jacobian1(f, v):
"""jacobian of f wrt v"""
f = tt.flatten(f)
idx = tt.arange(f.shape[0], dtype='int32')
def grad_i(i):
return gradient1(f[i], v)
return theano.map(grad_i, idx)[0]
def jacobian(f, vars=None):
if vars is None:
vars = cont_inputs(f)
if vars:
return tt.concatenate([jacobian1(f, v) for v in vars], axis=1)
else:
return empty_gradient
def jacobian_diag(f, x):
idx = tt.arange(f.shape[0], dtype='int32')
def grad_ii(i):
return theano.grad(f[i], x)[i]
return theano.scan(grad_ii, sequences=[idx],
n_steps=f.shape[0],
name='jacobian_diag')[0]
@change_flags(compute_test_value='ignore')
def hessian(f, vars=None):
return -jacobian(gradient(f, vars), vars)
@change_flags(compute_test_value='ignore')
def hessian_diag1(f, v):
g = gradient1(f, v)
idx = tt.arange(g.shape[0], dtype='int32')
def hess_ii(i):
return gradient1(g[i], v)[i]
return theano.map(hess_ii, idx)[0]
@change_flags(compute_test_value='ignore')
def hessian_diag(f, vars=None):
if vars is None:
vars = cont_inputs(f)
if vars:
return -tt.concatenate([hessian_diag1(f, v) for v in vars], axis=0)
else:
return empty_gradient
def makeiter(a):
if isinstance(a, (tuple, list)):
return a
else:
return [a]
class IdentityOp(scalar.UnaryScalarOp):
@staticmethod
def st_impl(x):
return x
def impl(self, x):
return x
def grad(self, inp, grads):
return grads
def c_code(self, node, name, inp, out, sub):
return "{z} = {x};".format(x=inp[0], z=out[0])
def __eq__(self, other):
return type(self) == type(other)
def __hash__(self):
return hash(type(self))
def make_shared_replacements(vars, model):
"""
Makes shared replacements for all *other* variables than the ones passed.
This way functions can be called many times without setting unchanging variables. Allows us
to use func.trust_input by removing the need for DictToArrayBijection and kwargs.
Parameters
----------
vars : list of variables not to make shared
model : model
Returns
-------
Dict of variable -> new shared variable
"""
othervars = set(model.vars) - set(vars)
return {var: theano.shared(var.tag.test_value, var.name + '_shared') for var in othervars}
def join_nonshared_inputs(xs, vars, shared, make_shared=False):
"""
Takes a list of theano Variables and joins their non shared inputs into a single input.
Parameters
----------
xs : list of theano tensors
vars : list of variables to join
Returns
-------
tensors, inarray
tensors : list of same tensors but with inarray as input
inarray : vector of inputs
"""
if not vars:
raise ValueError('Empty list of variables.')
joined = tt.concatenate([var.ravel() for var in vars])
if not make_shared:
tensor_type = joined.type
inarray = tensor_type('inarray')
else:
inarray = theano.shared(joined.tag.test_value, 'inarray')
ordering = ArrayOrdering(vars)
inarray.tag.test_value = joined.tag.test_value
get_var = {var.name: var for var in vars}
replace = {
get_var[var]: reshape_t(inarray[slc], shp).astype(dtyp)
for var, slc, shp, dtyp in ordering.vmap}
replace.update(shared)
xs_special = [theano.clone(x, replace, strict=False) for x in xs]
return xs_special, inarray
def reshape_t(x, shape):
"""Work around fact that x.reshape(()) doesn't work"""
if shape != ():
return x.reshape(shape)
else:
return x[0]
class CallableTensor(object):
"""Turns a symbolic variable with one input into a function that returns symbolic arguments
with the one variable replaced with the input.
"""
def __init__(self, tensor):
self.tensor = tensor
def __call__(self, input):
""" Replaces the single input of symbolic variable to be the passed argument.
Parameters
----------
input : TensorVariable
"""
oldinput, = inputvars(self.tensor)
return theano.clone(self.tensor, {oldinput: input}, strict=False)
scalar_identity = IdentityOp(scalar.upgrade_to_float, name='scalar_identity')
identity = tt.Elemwise(scalar_identity, name='identity')
class GeneratorOp(Op):
"""
Generator Op is designed for storing python generators inside theano graph.
__call__ creates TensorVariable
It has 2 new methods
- var.set_gen(gen) : sets new generator
- var.set_default(value) : sets new default value (None erases default value)
If generator is exhausted, variable will produce default value if it is not None,
else raises `StopIteration` exception that can be caught on runtime.
Parameters
----------
gen : generator that implements __next__ (py3) or next (py2) method
and yields np.arrays with same types
default : np.array with the same type as generator produces
"""
__props__ = ('generator',)
def __init__(self, gen, default=None):
super(GeneratorOp, self).__init__()
if not isinstance(gen, GeneratorAdapter):
gen = GeneratorAdapter(gen)
self.generator = gen
self.set_default(default)
def make_node(self, *inputs):
gen_var = self.generator.make_variable(self)
return theano.Apply(self, [], [gen_var])
def perform(self, node, inputs, output_storage, params=None):
if self.default is not None:
output_storage[0][0] = next(self.generator, self.default)
else:
output_storage[0][0] = next(self.generator)
def do_constant_folding(self, node):
return False
__call__ = change_flags(compute_test_value='off')(Op.__call__)
def set_gen(self, gen):
if not isinstance(gen, GeneratorAdapter):
gen = GeneratorAdapter(gen)
if not gen.tensortype == self.generator.tensortype:
raise ValueError('New generator should yield the same type')
self.generator = gen
def set_default(self, value):
if value is None:
self.default = None
else:
value = np.asarray(value, self.generator.tensortype.dtype)
t1 = (False,) * value.ndim
t2 = self.generator.tensortype.broadcastable
if not t1 == t2:
raise ValueError('Default value should have the '
'same type as generator')
self.default = value
def generator(gen, default=None):
"""
Generator variable with possibility to set default value and new generator.
If generator is exhausted variable will produce default value if it is not None,
else raises `StopIteration` exception that can be caught on runtime.
Parameters
----------
gen : generator that implements __next__ (py3) or next (py2) method
and yields np.arrays with same types
default : np.array with the same type as generator produces
Returns
-------
TensorVariable
It has 2 new methods
- var.set_gen(gen) : sets new generator
- var.set_default(value) : sets new default value (None erases default value)
"""
return GeneratorOp(gen, default)()
_tt_rng = MRG_RandomStreams()
def tt_rng(random_seed=None):
"""
Get the package-level random number generator or new with specified seed.
Parameters
----------
random_seed : int
If not None
returns *new* theano random generator without replacing package global one
Returns
-------
`theano.sandbox.rng_mrg.MRG_RandomStreams` instance
`theano.sandbox.rng_mrg.MRG_RandomStreams`
instance passed to the most recent call of `set_tt_rng`
"""
if random_seed is None:
return _tt_rng
else:
ret = MRG_RandomStreams(random_seed)
return ret
def set_tt_rng(new_rng):
"""
Set the package-level random number generator.
Parameters
----------
new_rng : `theano.sandbox.rng_mrg.MRG_RandomStreams` instance
The random number generator to use.
"""
# pylint: disable=global-statement
global _tt_rng
# pylint: enable=global-statement
if isinstance(new_rng, int):
new_rng = MRG_RandomStreams(new_rng)
_tt_rng = new_rng
def floatX_array(x):
return floatX(np.array(x))
def set_theano_conf(values):
"""Change the theano configuration and return old values.
This is similar to `theano.configparser.change_flags`, but it
returns the original values in a pickleable form.
"""
variables = {}
unknown = set(values.keys())
for variable in theano.configparser._config_var_list:
if variable.fullname in values:
variables[variable.fullname] = variable
unknown.remove(variable.fullname)
if len(unknown) > 0:
raise ValueError("Unknown theano config settings: %s" % unknown)
old = {}
for name, variable in variables.items():
old_value = variable.__get__(True, None)
try:
variable.__set__(None, values[name])
except Exception:
for key, old_value in old.items():
variables[key].__set__(None, old_value)
raise
old[name] = old_value
return old
def ix_(*args):
"""
Theano np.ix_ analog
See numpy.lib.index_tricks.ix_ for reference
"""
out = []
nd = len(args)
for k, new in enumerate(args):
if new is None:
out.append(slice(None))
new = tt.as_tensor(new)
if new.ndim != 1:
raise ValueError("Cross index must be 1 dimensional")
new = new.reshape((1,)*k + (new.size,) + (1,)*(nd-k-1))
out.append(new)
return tuple(out)
def largest_common_dtype(tensors):
dtypes = set(str(t.dtype) if hasattr(t, 'dtype')
else smartfloatX(np.asarray(t)).dtype
for t in tensors)
return np.stack([np.ones((), dtype=dtype) for dtype in dtypes]).dtype