def wrap_namespace(old, new): """ Wraps namespace of array library. """ unchanged_types = {float, int, type(None), type} int_types = { _cp.int, _cp.int8, _cp.int16, _cp.int32, _cp.int64, _cp.integer, } function_types = {_cp.ufunc, types.FunctionType, types.BuiltinFunctionType} for name, obj in old.items(): if obj in notrace_functions: new[name] = notrace_primitive(obj) # Note: type(obj) == _cp.ufunc doesn't work! Should use: # # isinstance(obj, _cp.ufunc) # elif (type(obj) in function_types or isinstance(obj, _cp.ufunc) # or isinstance(obj, _cp.core.fusion.reduction) ): new[name] = primitive(obj) elif type(obj) is type and obj in int_types: new[name] = wrap_intdtype(obj) elif type(obj) in unchanged_types: new[name] = obj
def autograd_register(f, s_f): """Register a function and its sensitivity for AutoGrad. Args: f (function): Function to register. s_f (function): Sensitivity of `f`. Returns: function: AutoGrad primitive. """ # Create a primitive for `f`. f_primitive = primitive(f) # Register the sensitivity. def vjp_argnums(nums, y, args, kw_args): def vjp(s_y): grads = as_tuple(s_f(s_y, y, *args, **kw_args)) return tuple([grads[i] for i in nums]) return vjp defvjp_argnums(f_primitive, vjp_argnums) # Return the AutoGrad primitive. return f_primitive
def wrap_namespace(old, new): unchanged_types = {float, int, type(None), type} int_types = {_np.int8, _np.int16, _np.int32, _np.int64, _np.integer} for name, obj in old.items(): if obj in notrace_functions: new[name] = notrace_primitive(obj) elif callable(obj) and type(obj) is not type: new[name] = primitive(obj) elif type(obj) is type and obj in int_types: new[name] = wrap_intdtype(obj) elif type(obj) in unchanged_types: new[name] = obj
def wrap_namespace(old, new): unchanged_types = {float, int, type(None), type} int_types = {_np.int, _np.int8, _np.int16, _np.int32, _np.int64, _np.integer} function_types = {_np.ufunc, types.FunctionType, types.BuiltinFunctionType} for name, obj in old.items(): if obj in notrace_functions: new[name] = notrace_primitive(obj) elif type(obj) in function_types: new[name] = primitive(obj) elif type(obj) is type and obj in int_types: new[name] = wrap_intdtype(obj) elif type(obj) in unchanged_types: new[name] = obj
def decorator(func): """Decorate a function to define its custome gradient(s). Parameters ---------- func : callable Function whose gradients will be assigned by grad_funcs. Returns ------- wrapped_function : callable Function func with gradients specified by grad_funcs. """ wrapped_function = primitive(func) def wrapped_grad_func(i, ans, *args, **kwargs): grads = grad_funcs[i](*args, **kwargs) if isinstance(grads, float): return lambda g: g * grads if grads.ndim == 2: return lambda g: g[..., None] * grads if grads.ndim == 3: return lambda g: g[..., None, None] * grads return lambda g: g * grads if len(grad_funcs) == 1: defvjp( wrapped_function, lambda ans, *args, **kwargs: wrapped_grad_func(0, ans, *args, **kwargs), ) elif len(grad_funcs) == 2: defvjp( wrapped_function, lambda ans, *args, **kwargs: wrapped_grad_func(0, ans, *args, **kwargs), lambda ans, *args, **kwargs: wrapped_grad_func(1, ans, *args, **kwargs), ) elif len(grad_funcs) == 3: defvjp( wrapped_function, lambda ans, *args, **kwargs: wrapped_grad_func(0, ans, *args, **kwargs), lambda ans, *args, **kwargs: wrapped_grad_func(1, ans, *args, **kwargs), lambda ans, *args, **kwargs: wrapped_grad_func(2, ans, *args, **kwargs), ) else: raise NotImplementedError( "custom_gradient is not yet implemented " "for more than 3 gradients." ) return wrapped_function
def wrap_namespace(old, new): unchanged_types = {float, int, type(None), type} int_types = { _np.int, _np.int8, _np.int16, _np.int32, _np.int64, _np.integer } function_types = {_np.ufunc, types.FunctionType, types.BuiltinFunctionType} for name, obj in old.items(): if obj in notrace_functions: # NOTE(brendan): notrace_primitive has to marshal out all of the # values from potentially boxed obj's. new[name] = notrace_primitive(obj) elif type(obj) in function_types: new[name] = primitive(obj) elif type(obj) is type and obj in int_types: new[name] = wrap_intdtype(obj) elif type(obj) in unchanged_types: new[name] = obj
def decorator(func): wrapped_function = primitive(func) def wrapped_grad_func(i, ans, *args, **kwargs): grads = grad_funcs[i](*args, **kwargs) if isinstance(grads, float): return lambda g: g * grads if grads.ndim == 2: return lambda g: g[..., None] * grads if grads.ndim == 3: return lambda g: g[..., None, None] * grads return lambda g: g * grads if len(grad_funcs) == 1: defvjp( wrapped_function, lambda ans, *args, **kwargs: wrapped_grad_func(0, ans, *args, **kwargs), ) elif len(grad_funcs) == 2: defvjp( wrapped_function, lambda ans, *args, **kwargs: wrapped_grad_func(0, ans, *args, **kwargs), lambda ans, *args, **kwargs: wrapped_grad_func(1, ans, *args, **kwargs), ) elif len(grad_funcs) == 3: defvjp( wrapped_function, lambda ans, *args, **kwargs: wrapped_grad_func(0, ans, *args, **kwargs), lambda ans, *args, **kwargs: wrapped_grad_func(1, ans, *args, **kwargs), lambda ans, *args, **kwargs: wrapped_grad_func(2, ans, *args, **kwargs), ) else: raise NotImplementedError( "custom_gradient is not yet implemented " "for more than 3 gradients." ) return wrapped_function
from __future__ import absolute_import import scipy.special import autograd.numpy as np from autograd.extend import primitive, defvjp, defjvp from autograd.numpy.numpy_vjps import unbroadcast_f, repeat_to_match_shape ### Beta function ### beta = primitive(scipy.special.beta) betainc = primitive(scipy.special.betainc) betaln = primitive(scipy.special.betaln) defvjp(beta, lambda ans, a, b: unbroadcast_f(a, lambda g: g * ans * (psi(a) - psi(a + b))), lambda ans, a, b: unbroadcast_f(b, lambda g: g * ans * (psi(b) - psi(a + b)))) defvjp(betainc, lambda ans, a, b, x: unbroadcast_f(x, lambda g: g * np.power(x, a - 1) * np.power(1 - x, b - 1) / beta(a, b)), argnums=[2]) defvjp(betaln, lambda ans, a, b: unbroadcast_f(a, lambda g: g * (psi(a) - psi(a + b))), lambda ans, a, b: unbroadcast_f(b, lambda g: g * (psi(b) - psi(a + b)))) ### Gamma functions ### polygamma = primitive(scipy.special.polygamma) psi = primitive(scipy.special.psi) # psi(x) is just polygamma(0, x) digamma = primitive(scipy.special.digamma) # digamma is another name for psi. gamma = primitive(scipy.special.gamma) gammaln = primitive(scipy.special.gammaln) gammainc = primitive(scipy.special.gammainc) gammaincc = primitive(scipy.special.gammaincc) gammasgn = primitive(scipy.special.gammasgn) rgamma = primitive(scipy.special.rgamma)
# -*- coding: utf-8 -*- from __future__ import division from scipy.stats import norm as _scipy_norm import autograd.numpy as np from autograd.scipy.stats import norm from autograd.extend import primitive, defvjp from autograd.numpy.numpy_vjps import unbroadcast_f # TODO: next release of autograd will have this built in. logsf = primitive(_scipy_norm.logsf) defvjp( logsf, lambda ans, x, loc=0.0, scale=1.0: unbroadcast_f( x, lambda g: -g * np.exp( norm.logpdf(x, loc, scale) - logsf(x, loc, scale))), lambda ans, x, loc=0.0, scale=1.0: unbroadcast_f( loc, lambda g: g * np.exp( norm.logpdf(x, loc, scale) - logsf(x, loc, scale))), lambda ans, x, loc=0.0, scale=1.0: unbroadcast_f( scale, lambda g: g * np.exp( norm.logpdf(x, loc, scale) - logsf(x, loc, scale)) * (x - loc) / scale), )
from __future__ import absolute_import from builtins import range import scipy.integrate import autograd.numpy as np from autograd.extend import primitive, defvjp_argnums from autograd import make_vjp from autograd.misc import flatten from autograd.builtins import tuple odeint = primitive(scipy.integrate.odeint) def grad_odeint(yt, func, y0, t, func_args, **kwargs): # Extended from "Scalable Inference of Ordinary Differential # Equation Models of Biochemical Processes", Sec. 2.4.2 # Fabian Froehlich, Carolin Loos, Jan Hasenauer, 2017 # https://arxiv.org/abs/1711.08079 T, D = np.shape(yt) flat_args, unflatten = flatten(func_args) def flat_func(y, t, flat_args): return func(y, t, *unflatten(flat_args)) def unpack(x): # y, vjp_y, vjp_t, vjp_args return x[0:D], x[D:2 * D], x[2 * D], x[2 * D + 1:] def augmented_dynamics(augmented_state, t, flat_args):
return -transpose(tri(anp.matmul(anp.reshape(v, ans.shape), T(ans)))) return solve_triangular_grad defvjp(solve_triangular, make_grad_solve_triangular, lambda ans, a, b, trans=0, lower=False, **kwargs: lambda g: solve_triangular(a, g, trans=_flip(a, trans), lower=lower), None) ### cholesky solve_trans = lambda L, X: solve_triangular(L, X, lower=True, trans='T') solve_conj = lambda L, X: solve_trans(L, T(solve_trans(L, T(X)))) phi = lambda X: anp.tril(X) / (1. + anp.eye(X.shape[-1])) cholesky = primitive(np.linalg.cholesky) defvjp(cholesky, lambda L, A: lambda g: symm(solve_conj(L, phi(anp.matmul(T(L), g))))) ### operations on cholesky factors solve_tri = partial(solve_triangular, lower=True) solve_posdef_from_cholesky = lambda L, x: solve_tri(L, solve_tri(L, x), trans='T') @primitive def inv_posdef_from_cholesky(L, lower=True): flat_L = np.reshape(L, (-1,) + L.shape[-2:]) return np.reshape(cyla.inv_posdef_from_cholesky(C(flat_L), lower), L.shape) square_grad = lambda X: lambda g: anp.matmul(g, X) + anp.matmul(T(g), X) sym_inv_grad = lambda Xinv: lambda g: -anp.matmul(Xinv, anp.matmul(g, Xinv))
import numpy as np from autograd.extend import primitive, defvjp, vspace ''' Define here various primitives needed for the main code To use with both numpy and autograd backends, define the autograd primitive of a numpy function fnc as fnc_ag, and then define the vjp''' def T(x): return np.swapaxes(x, -1, -2) '''=========== NP.SQRT STABLE AROUND 0 =========== ''' sqrt_ag = primitive(np.sqrt) def vjp_maker_sqrt(ans, x): def vjp(g): return g * 0.5 * (x + 1e-10)**0.5/(x + 1e-10) # return np.where(np.abs(x) > 1e-10, g * 0.5 * x**-0.5, 0.) return vjp defvjp(sqrt_ag, vjp_maker_sqrt) def vjp_maker_meshgridx(ans, x): def vjp(g): return np.sum(g,axis=1) return vjp '''=========== inv =========== ''' inv_ag = primitive(np.linalg.inv) def vjp_maker_inv(ans, x): return lambda g: -np.dot(np.dot(T(ans), g), T(ans))
betainc as _scipy_betainc, ) __all__ = [ "gammainc", # regularized lower incomplete gamma function "gammaincc", # regularized upper incomplete gamma function "gamma", # gamma function "betainc", "betaincln", "beta", ] LOG_EPISILON = 1e-35 MACHINE_EPISLON_POWER = np.finfo(float).eps**(1 / 2) gammainc = primitive(_scipy_gammainc) gammaincc = primitive(_scipy_gammaincc) betainc = primitive(_scipy_betainc) @primitive def gammainccln(a, x): return np.log(np.clip(gammaincc(a, x), LOG_EPISILON, 1 - LOG_EPISILON)) @primitive def gammaincln(a, x): return np.log(np.clip(gammainc(a, x), LOG_EPISILON, 1 - LOG_EPISILON)) @primitive
from __future__ import absolute_import import scipy.stats import autograd.numpy as np from autograd.numpy.numpy_vjps import unbroadcast_f from autograd.extend import primitive, defvjp pdf = primitive(scipy.stats.multivariate_normal.pdf) logpdf = primitive(scipy.stats.multivariate_normal.logpdf) entropy = primitive(scipy.stats.multivariate_normal.entropy) # With thanks to Eric Bresch. # Some formulas are from # "An extended collection of matrix derivative results # for forward and reverse mode algorithmic differentiation" # by Mike Giles # https://people.maths.ox.ac.uk/gilesm/files/NA-08-01.pdf def generalized_outer_product(x): if np.ndim(x) == 1: return np.outer(x, x) return np.matmul(x, np.swapaxes(x, -1, -2)) def covgrad(x, mean, cov, allow_singular=False): if allow_singular: raise NotImplementedError( "The multivariate normal pdf is not " "differentiable w.r.t. a singular covariance matix") J = np.linalg.inv(cov)
if use_autograd.use == 0: # autograd turned off import numpy as np eig = np.linalg.eig inv = np.linalg.inv sqrt = np.sqrt else: # autograd turned on import autograd.numpy as np from autograd import grad from autograd.extend import primitive, defvjp from primitives import vjp_maker_eig, vjp_maker_sqrt, vjp_maker_inv from primitives_fix import grad_eig import numpy as npf eig = primitive(npf.linalg.eig) sqrt = primitive(npf.sqrt) inv = primitive(npf.linalg.inv) #defvjp(eig, vjp_maker_eig) defvjp(eig, grad_eig) defvjp(sqrt, vjp_maker_sqrt) defvjp(inv, vjp_maker_inv) class RCWA_obj: def __init__(self, nG, L1, L2, freq, theta, phi, verbose=1): '''The time harmonic convention is exp(-i omega t), speed of light = 1 The first and last layer must be uniform Two kinds of layers are currently supported: uniform layer, patterned layer from grids. Interface for patterned layer by
def test_no_jvp_def(): fun = primitive(lambda x: 2. * x) deriv(fun)(1.)
def test_no_vjp_def(): fun = primitive(lambda x: 2. * x) grad(fun)(1.)
_diag = lambda a: anp.eye(a.shape[-1]) * a # batched diagonal, similar to matrix_diag in tensorflow def _matrix_diag(a): reps = anp.array(a.shape) reps[:-1] = 1 reps[-1] = a.shape[-1] newshape = list(a.shape) + [a.shape[-1]] return _diag(anp.tile(a, reps).reshape(newshape)) # https://arxiv.org/pdf/1701.00392.pdf Eq(4.77) # Note the formula from Sec3.1 in https://people.maths.ox.ac.uk/gilesm/files/NA-08-01.pdf is incomplete inv = primitive(anp.linalg.inv) def grad_inv(ans, x): return lambda g: -_dot(_dot(T(ans), g), T(ans)) defvjp(inv, grad_inv) eig = primitive(anp.linalg.eig) def grad_eig(ans, x): """Gradient of a general square (complex valued) matrix""" e, u = ans # eigenvalues as 1d array, eigenvectors in columns n = e.shape[-1]
from __future__ import absolute_import import autograd.numpy as np import scipy.stats from autograd.extend import primitive, defvjp from autograd.numpy.numpy_vjps import unbroadcast_f from autograd.scipy.special import beta, psi cdf = primitive(scipy.stats.beta.cdf) logpdf = primitive(scipy.stats.beta.logpdf) pdf = primitive(scipy.stats.beta.pdf) def grad_beta_logpdf_arg0(x, a, b): return (1 + a * (x - 1) + x * (b - 2)) / (x * (x - 1)) def grad_beta_logpdf_arg1(x, a, b): return np.log(x) - psi(a) + psi(a + b) def grad_beta_logpdf_arg2(x, a, b): return np.log1p(-x) - psi(b) + psi(a + b) defvjp(cdf, lambda ans, x, a, b: unbroadcast_f( x, lambda g: g * np.power(x, a - 1) * np.power(1 - x, b - 1) / beta( a, b)), argnums=[0]) defvjp(
from __future__ import absolute_import import scipy.stats import autograd.numpy as np from autograd.scipy.special import digamma from autograd.extend import primitive, defvjp rvs = primitive(scipy.stats.dirichlet.rvs) pdf = primitive(scipy.stats.dirichlet.pdf) logpdf = primitive(scipy.stats.dirichlet.logpdf) defvjp( logpdf, lambda ans, x, alpha: lambda g: g * (alpha - 1) / x, lambda ans, x, alpha: lambda g: g * (digamma(np.sum(alpha)) - digamma(alpha) + np.log(x))) # Same as log pdf, but multiplied by the pdf (ans). defvjp( pdf, lambda ans, x, alpha: lambda g: g * ans * (alpha - 1) / x, lambda ans, x, alpha: lambda g: g * ans * (digamma(np.sum(alpha)) - digamma(alpha) + np.log(x)))
"""Gradients of the univariate t distribution.""" from __future__ import absolute_import import scipy.stats import autograd.numpy as np from autograd.extend import primitive, defvjp from autograd.numpy.numpy_vjps import unbroadcast_f from autograd.scipy.special import psi pdf = primitive(scipy.stats.t.pdf) cdf = primitive(scipy.stats.t.cdf) logpdf = primitive(scipy.stats.t.logpdf) logcdf = primitive(scipy.stats.t.logcdf) def grad_tlogpdf_diff(diff, df): return -diff * (1.0 + df) / (diff**2 + df) def grad_tlogpdf_x(x, df, loc, scale): return grad_tlogpdf_diff((x - loc) / scale, df) / scale def grad_tlogpdf_loc(x, df, loc, scale): return -grad_tlogpdf_diff((x - loc) / scale, df) / scale def grad_tlogpdf_scale(x, df, loc, scale): diff = x - loc return -(df * (scale**2 - diff**2))/(scale * (df * scale**2 + diff**2)) def grad_tlogpdf_df(x, df, loc, scale): y = (x - loc)/scale return 0.5 * ((y**2 * (df+1))/(df * (y**2 + df)) - np.log(y**2 / df + 1) - 1.0/df -psi(df/2.0) + psi((df + 1)/2.0)) defvjp(pdf, lambda ans, x, df, loc=0.0, scale=1.0: unbroadcast_f(x, lambda g: g * ans * grad_tlogpdf_x( x, df, loc, scale)), lambda ans, x, df, loc=0.0, scale=1.0: unbroadcast_f(df, lambda g: g * ans * grad_tlogpdf_df( x, df, loc, scale)), lambda ans, x, df, loc=0.0, scale=1.0:
from __future__ import absolute_import, division import autograd.numpy as np import scipy.stats from autograd.extend import primitive, defvjp from autograd.numpy.numpy_vjps import unbroadcast_f from autograd.scipy.special import gamma cdf = primitive(scipy.stats.chi2.cdf) logpdf = primitive(scipy.stats.chi2.logpdf) pdf = primitive(scipy.stats.chi2.pdf) def grad_chi2_logpdf(x, df): return np.where(df % 1 == 0, (df - x - 2) / (2 * x), 0) defvjp(cdf, lambda ans, x, df: unbroadcast_f(x, lambda g: g * np.power(2., -df/2) * np.exp(-x/2) * np.power(x, df/2 - 1) / gamma(df/2)), argnums=[0]) defvjp(logpdf, lambda ans, x, df: unbroadcast_f(x, lambda g: g * grad_chi2_logpdf(x, df)), argnums=[0]) defvjp(pdf, lambda ans, x, df: unbroadcast_f(x, lambda g: g * ans * grad_chi2_logpdf(x, df)), argnums=[0])
from __future__ import absolute_import import scipy.misc from autograd.extend import primitive, defvjp import autograd.numpy as anp from autograd.numpy.numpy_vjps import repeat_to_match_shape logsumexp = primitive(scipy.misc.logsumexp) def make_grad_logsumexp(ans, x, axis=None, b=1.0, keepdims=False): shape, dtype = anp.shape(x), anp.result_type(x) def vjp(g): g_repeated, _ = repeat_to_match_shape(g, shape, dtype, axis, keepdims) ans_repeated, _ = repeat_to_match_shape(ans, shape, dtype, axis, keepdims) return g_repeated * b * anp.exp(x - ans_repeated) return vjp defvjp(logsumexp, make_grad_logsumexp)
from __future__ import absolute_import import autograd.numpy as np import scipy.stats from autograd.extend import primitive, defvjp from autograd.numpy.numpy_vjps import unbroadcast_f cdf = primitive(scipy.stats.poisson.cdf) logpmf = primitive(scipy.stats.poisson.logpmf) pmf = primitive(scipy.stats.poisson.pmf) def grad_poisson_logpmf(k, mu): return np.where(k % 1 == 0, k / mu - 1, 0) defvjp(cdf, lambda ans, k, mu: unbroadcast_f(mu, lambda g: g * -pmf(np.floor(k), mu)), argnums=[1]) defvjp(logpmf, lambda ans, k, mu: unbroadcast_f(mu, lambda g: g * grad_poisson_logpmf(k, mu)), argnums=[1]) defvjp(pmf, lambda ans, k, mu: unbroadcast_f(mu, lambda g: g * ans * grad_poisson_logpmf(k, mu)), argnums=[1])
from .util import memoize, check_psd from .convolution import convolve_sum_axes, transposed_convolve_sum_axes, sum_trailing_antidiagonals, add_trailing_axis, roll_trailing_axes, unroll_trailing_axes from .einsum2 import einsum1, einsum2 def par_einsum(*args): return einsum2(*args) def convolve_trailing_axes(A, B): A = np.reshape(A, list(A.shape) + [1]) B = np.reshape(B, list(B.shape) + [1]) return convolve_sum_axes(A, B) convolve_sum_axes = primitive(convolve_sum_axes) transposed_convolve_sum_axes = primitive(transposed_convolve_sum_axes) defvjp( convolve_sum_axes, lambda ans, A, B: lambda g: transposed_convolve_sum_axes(g, B), lambda ans, A, B: lambda g: transposed_convolve_sum_axes( np.transpose(g, (0, 2, 1, 3)), A)) #convolve_sum_axes.defgrad(lambda ans, A, B: lambda g: transposed_convolve_sum_axes(g, B)) #convolve_sum_axes.defgrad(lambda ans, A, B: lambda g: transposed_convolve_sum_axes(np.transpose(g, (0, 2, 1, 3)), A), argnum=1) #convolve_sum_axes.defvjp( # lambda g, ans, vs, gvs, A, B: transposed_convolve_sum_axes(g, B)) #convolve_sum_axes.defvjp( # lambda g, ans, vs, gvs, A, B: transposed_convolve_sum_axes( # np.transpose(g, (0, 2, 1, 3)), A), # argnum=1)
return _scipy_gammainc(k, x) delta = 1e-6 defvjp( gammainc, lambda ans, a, x: unbroadcast_f( a, lambda g: g * (-gammainc(a + 2 * delta, x) + 8 * gammainc(a + delta, x) - 8 * gammainc(a - delta, x) + gammainc(a - 2 * delta, x)) / (12 * delta), ), lambda ans, a, x: unbroadcast_f( x, lambda g: g * np.exp(-x) * np.power(x, a - 1) / gamma(a)), ) gammaincc = primitive(_scipy_gammaincc) defvjp( gammaincc, lambda ans, a, x: unbroadcast_f( a, lambda g: g * (-gammaincc(a + 2 * delta, x) + 8 * gammaincc(a + delta, x) - 8 * gammaincc(a - delta, x) + gammaincc(a - 2 * delta, x)) / (12 * delta), ), lambda ans, a, x: unbroadcast_f( x, lambda g: -g * np.exp(-x) * np.power(x, a - 1) / gamma(a)), )
from __future__ import absolute_import, division import autograd.numpy as np import scipy.stats from autograd.extend import primitive, defvjp from autograd.numpy.numpy_vjps import unbroadcast_f from autograd.scipy.special import gamma cdf = primitive(scipy.stats.chi2.cdf) logpdf = primitive(scipy.stats.chi2.logpdf) pdf = primitive(scipy.stats.chi2.pdf) def grad_chi2_logpdf(x, df): return np.where(df % 1 == 0, (df - x - 2) / (2 * x), 0) defvjp(cdf, lambda ans, x, df: unbroadcast_f( x, lambda g: g * np.power(2., -df / 2) * np.exp(-x / 2) * np.power( x, df / 2 - 1) / gamma(df / 2)), argnums=[0]) defvjp( logpdf, lambda ans, x, df: unbroadcast_f(x, lambda g: g * grad_chi2_logpdf(x, df)), argnums=[0]) defvjp(pdf, lambda ans, x, df: unbroadcast_f( x, lambda g: g * ans * grad_chi2_logpdf(x, df)), argnums=[0])
from __future__ import absolute_import import scipy.special import autograd.numpy as np from autograd.extend import primitive, defvjp from autograd.numpy.numpy_vjps import unbroadcast_f ### Beta function ### beta = primitive(scipy.special.beta) betainc = primitive(scipy.special.betainc) betaln = primitive(scipy.special.betaln) defvjp(beta, lambda ans, a, b: unbroadcast_f(a, lambda g: g * ans * (psi(a) - psi(a + b))), lambda ans, a, b: unbroadcast_f(b, lambda g: g * ans * (psi(b) - psi(a + b)))) defvjp(betainc, lambda ans, a, b, x: unbroadcast_f(x, lambda g: g * np.power(x, a - 1) * np.power(1 - x, b - 1) / beta(a, b)), argnums=[2]) defvjp(betaln, lambda ans, a, b: unbroadcast_f(a, lambda g: g * (psi(a) - psi(a + b))), lambda ans, a, b: unbroadcast_f(b, lambda g: g * (psi(b) - psi(a + b)))) ### Gamma functions ### polygamma = primitive(scipy.special.polygamma) psi = primitive(scipy.special.psi) # psi(x) is just polygamma(0, x) digamma = primitive(scipy.special.digamma) # digamma is another name for psi. gamma = primitive(scipy.special.gamma) gammaln = primitive(scipy.special.gammaln) gammainc = primitive(scipy.special.gammainc) gammaincc = primitive(scipy.special.gammaincc) gammasgn = primitive(scipy.special.gammasgn) rgamma = primitive(scipy.special.rgamma)
def test_no_jvp_def(): fun = primitive(lambda x: 2. * x) with pytest.raises(NotImplementedError): deriv(fun)(1.)
from __future__ import absolute_import import scipy.misc, scipy.special from autograd.extend import primitive, defvjp import autograd.numpy as anp from autograd.numpy.numpy_vjps import repeat_to_match_shape logsumexp = primitive(scipy.special.logsumexp) def make_grad_logsumexp(ans, x, axis=None, b=1.0, keepdims=False): shape, dtype = anp.shape(x), anp.result_type(x) def vjp(g): g_repeated, _ = repeat_to_match_shape(g, shape, dtype, axis, keepdims) ans_repeated, _ = repeat_to_match_shape(ans, shape, dtype, axis, keepdims) return g_repeated * b * anp.exp(x - ans_repeated) return vjp defvjp(logsumexp, make_grad_logsumexp)
from __future__ import absolute_import import autograd.numpy as np import scipy.stats from autograd.extend import primitive, defvjp from autograd.numpy.numpy_vjps import unbroadcast_f cdf = primitive(scipy.stats.poisson.cdf) logpmf = primitive(scipy.stats.poisson.logpmf) pmf = primitive(scipy.stats.poisson.pmf) def grad_poisson_logpmf(k, mu): return np.where(k % 1 == 0, k / mu - 1, 0) defvjp( cdf, lambda ans, k, mu: unbroadcast_f(mu, lambda g: g * -pmf(np.floor(k), mu)), argnums=[1]) defvjp(logpmf, lambda ans, k, mu: unbroadcast_f( mu, lambda g: g * grad_poisson_logpmf(k, mu)), argnums=[1]) defvjp(pmf, lambda ans, k, mu: unbroadcast_f( mu, lambda g: g * ans * grad_poisson_logpmf(k, mu)), argnums=[1])
from __future__ import absolute_import import autograd.numpy as np import scipy.stats from autograd.extend import primitive, defvjp from autograd.numpy.numpy_vjps import unbroadcast_f from autograd.scipy.special import gamma, psi cdf = primitive(scipy.stats.gamma.cdf) logpdf = primitive(scipy.stats.gamma.logpdf) pdf = primitive(scipy.stats.gamma.pdf) def grad_gamma_logpdf_arg0(x, a): return (a - x - 1) / x def grad_gamma_logpdf_arg1(x, a): return np.log(x) - psi(a) defvjp(cdf, lambda ans, x, a: unbroadcast_f(x, lambda g: g * np.exp(-x) * np.power(x, a-1) / gamma(a)), argnums=[0]) defvjp(logpdf, lambda ans, x, a: unbroadcast_f(x, lambda g: g * grad_gamma_logpdf_arg0(x, a)), lambda ans, x, a: unbroadcast_f(a, lambda g: g * grad_gamma_logpdf_arg1(x, a))) defvjp(pdf, lambda ans, x, a: unbroadcast_f(x, lambda g: g * ans * grad_gamma_logpdf_arg0(x, a)), lambda ans, x, a: unbroadcast_f(a, lambda g: g * ans * grad_gamma_logpdf_arg1(x, a)))
from autograd.extend import primitive, defvjp, vspace from autograd import grad, vector_jacobian_product import autograd.numpy as npa """ Define here various primitives needed for the main code To use with both numpy and autograd backends, define the autograd primitive of a numpy function fnc as fnc_ag, and then define the vjp""" def T(x): return np.swapaxes(x, -1, -2) """=========== EXPAND ARRAY TO A GIVEN SHAPE =========== """ # extend(vals, inds, shape) makes an array of shape `shape` where indices # `inds` have values `vals` extend_ag = primitive(extend) def vjp_maker_extend(ans, vals, inds, shape): def vjp(g): return g[inds] return vjp defvjp(extend_ag, vjp_maker_extend, None, None) """=========== NP.SQRT STABLE AROUND 0 =========== """ sqrt_ag = primitive(np.sqrt) def vjp_maker_sqrt(ans, x): def vjp(g): return g * 0.5 * (x + 1e-10)**0.5/(x + 1e-10) # return np.where(np.abs(x) > 1e-10, g * 0.5 * x**-0.5, 0.)
"""Gradients of the normal distribution.""" from __future__ import absolute_import import scipy.stats import autograd.numpy as anp from autograd.extend import primitive, defvjp from autograd.numpy.numpy_vjps import unbroadcast_f pdf = primitive(scipy.stats.norm.pdf) cdf = primitive(scipy.stats.norm.cdf) sf = primitive(scipy.stats.norm.sf) logpdf = primitive(scipy.stats.norm.logpdf) logcdf = primitive(scipy.stats.norm.logcdf) logsf = primitive(scipy.stats.norm.logsf) defvjp(pdf, lambda ans, x, loc=0.0, scale=1.0: unbroadcast_f(x, lambda g: -g * ans * (x - loc) / scale**2), lambda ans, x, loc=0.0, scale=1.0: unbroadcast_f(loc, lambda g: g * ans * (x - loc) / scale**2), lambda ans, x, loc=0.0, scale=1.0: unbroadcast_f(scale, lambda g: g * ans * (((x - loc)/scale)**2 - 1.0)/scale)) defvjp(cdf, lambda ans, x, loc=0.0, scale=1.0: unbroadcast_f(x, lambda g: g * pdf(x, loc, scale)) , lambda ans, x, loc=0.0, scale=1.0: unbroadcast_f(loc, lambda g: -g * pdf(x, loc, scale)), lambda ans, x, loc=0.0, scale=1.0: unbroadcast_f(scale, lambda g: -g * pdf(x, loc, scale)*(x-loc)/scale)) defvjp(logpdf,
import use_autograd if use_autograd.use == 0: import numpy as np import numpy as npf inv = np.linalg.inv else: import autograd.numpy as np from autograd import grad from autograd.extend import primitive, defvjp from primitives import vjp_maker_inv import numpy as npf inv = primitive(npf.linalg.inv) defvjp(inv, vjp_maker_inv) def GetEpsilon_FFT(dN, eps_grid, G): '''dN = 1/Nx/Ny eps_grid is a numpy 2d array in the format of (Nx,Ny) For now, assume epsilon is isotropic if epsilon has xz,yz component, just simply add them to off-diagonal eps2 ''' eps_fft = get_conv(dN, eps_grid, G) epsinv = inv(eps_fft) # somehow block don't work with autograd # eps2 = np.block([[eps_fft,np.zeros_like(eps_fft)], # [np.zeros_like(eps_fft),eps_fft]])
""" From Appendix B in the paper Implementation of autograd """ import scipy.integrate import autograd.numpy as np from autograd.extend import primitive, defvjp_argnums from autograd import make_vjp from autograd.misc import flatten from autograd.builtins import tuple odeint = primitive(scipy.integrate.odeint) def grad_odeint_all(yt, func, y0, t, func_args, **kwargs): """ Extended from "Scalable Inference of Ordinary Differential" Equation Models of Biochemical Processes". Sec. 2.4.2 Fabian Froehlich, Carolin Loos, Jan Hasenauer, 2017 https://arxiv.org/pdf/1711.08079.pdf """ T, D = np.shape(yt) flat_args, unflatten = flatten(func_args) def flat_func(y, t, flat_args): return func(y, t, *unflatten(flat_args))
"""Gradients of the normal distribution.""" from __future__ import absolute_import import scipy.stats import autograd.numpy as anp from autograd.extend import primitive, defvjp from autograd.numpy.numpy_vjps import unbroadcast_f pdf = primitive(scipy.stats.norm.pdf) cdf = primitive(scipy.stats.norm.cdf) sf = primitive(scipy.stats.norm.sf) logpdf = primitive(scipy.stats.norm.logpdf) logcdf = primitive(scipy.stats.norm.logcdf) logsf = primitive(scipy.stats.norm.logsf) defvjp(pdf, lambda ans, x, loc=0.0, scale=1.0: unbroadcast_f( x, lambda g: -g * ans * (x - loc) / scale**2), lambda ans, x, loc=0.0, scale=1.0: unbroadcast_f( loc, lambda g: g * ans * (x - loc) / scale**2), lambda ans, x, loc=0.0, scale=1.0: unbroadcast_f( scale, lambda g: g * ans * (((x - loc) / scale)**2 - 1.0) / scale)) defvjp(cdf, lambda ans, x, loc=0.0, scale=1.0: unbroadcast_f( x, lambda g: g * pdf(x, loc, scale)), lambda ans, x, loc=0.0, scale=1.0: unbroadcast_f( loc, lambda g: -g * pdf(x, loc, scale)), lambda ans, x, loc=0.0, scale=1.0: unbroadcast_f( scale, lambda g: -g * pdf(x, loc, scale) * (x - loc) / scale)) defvjp(logpdf,
from __future__ import absolute_import import scipy.stats import autograd.numpy as np from autograd.scipy.special import digamma from autograd.extend import primitive, defvjp rvs = primitive(scipy.stats.dirichlet.rvs) pdf = primitive(scipy.stats.dirichlet.pdf) logpdf = primitive(scipy.stats.dirichlet.logpdf) defvjp(logpdf,lambda ans, x, alpha: lambda g: g * (alpha - 1) / x, lambda ans, x, alpha: lambda g: g * (digamma(np.sum(alpha)) - digamma(alpha) + np.log(x))) # Same as log pdf, but multiplied by the pdf (ans). defvjp(pdf,lambda ans, x, alpha: lambda g: g * ans * (alpha - 1) / x, lambda ans, x, alpha: lambda g: g * ans * (digamma(np.sum(alpha)) - digamma(alpha) + np.log(x)))
from __future__ import absolute_import import scipy.stats import autograd.numpy as np from autograd.numpy.numpy_vjps import unbroadcast_f from autograd.extend import primitive, defvjp pdf = primitive(scipy.stats.multivariate_normal.pdf) logpdf = primitive(scipy.stats.multivariate_normal.logpdf) entropy = primitive(scipy.stats.multivariate_normal.entropy) # With thanks to Eric Bresch. # Some formulas are from # "An extended collection of matrix derivative results # for forward and reverse mode algorithmic differentiation" # by Mike Giles # https://people.maths.ox.ac.uk/gilesm/files/NA-08-01.pdf def generalized_outer_product(x): if np.ndim(x) == 1: return np.outer(x, x) return np.matmul(x, np.swapaxes(x, -1, -2)) def covgrad(x, mean, cov, allow_singular=False): if allow_singular: raise NotImplementedError("The multivariate normal pdf is not " "differentiable w.r.t. a singular covariance matix") J = np.linalg.inv(cov) solved = np.matmul(J, np.expand_dims(x - mean, -1)) return 1./2 * (generalized_outer_product(solved) - J)