def check_vjp(fun, arg):
vs_in = vspace(arg)
vs_out = vspace(fun(arg))
autograd_jac = linear_fun_to_matrix(
flatten_fun(make_vjp(fun)(arg)[0], vs_out), vs_out).T
numerical_jac = linear_fun_to_matrix(
numerical_deriv(flatten_fun(fun, vs_in), vspace_flatten(arg)), vs_in)
assert np.allclose(autograd_jac, numerical_jac)
def numerical_jacobian(fun, argnum, args, kwargs):
def vector_fun(x):
args_tmp = list(args)
args_tmp[argnum] = vs_in.unflatten(vs_in.flatten(args[argnum]) + x)
return vs_out.flatten(fun(*args_tmp, **kwargs))
vs_in = vspace(args[argnum])
vs_out = vspace(fun(*args, **kwargs))
return np.stack([(vector_fun(dx) - vector_fun(-dx)) / EPS
for dx in np.eye(vs_in.size) * EPS / 2]).T
def numerical_jacobian(fun, argnum, args, kwargs):
def vector_fun(x):
args_tmp = list(args)
args_tmp[argnum] = vs_in.unflatten(vs_in.flatten(args[argnum]) + x)
return vs_out.flatten(fun(*args_tmp, **kwargs))
vs_in = vspace(args[argnum])
vs_out = vspace(fun(*args, **kwargs))
return np.stack([(vector_fun(dx) - vector_fun(-dx)) / EPS
for dx in np.eye(vs_in.size) * EPS / 2]).T
def check_equivalent(A, B, rtol=RTOL, atol=ATOL):
A_vspace = vspace(A)
B_vspace = vspace(B)
A_flat = vspace_flatten(A)
B_flat = vspace_flatten(B)
assert A_vspace == B_vspace, \
"VSpace mismatch:\nanalytic: {}\nnumeric: {}".format(A_vspace, B_vspace)
assert np.allclose(vspace_flatten(A), vspace_flatten(B), rtol=rtol, atol=atol), \
"Diffs are:\n{}.\nanalytic is:\n{}.\nnumeric is:\n{}.".format(
A_flat - B_flat, A_flat, B_flat)
def check_equivalent(A, B, rtol=RTOL, atol=ATOL):
A_vspace = vspace(A)
B_vspace = vspace(B)
A_flat = vspace_flatten(A)
B_flat = vspace_flatten(B)
assert A_vspace == B_vspace, \
"VSpace mismatch:\nanalytic: {}\nnumeric: {}".format(A_vspace, B_vspace)
assert np.allclose(vspace_flatten(A), vspace_flatten(B), rtol=rtol, atol=atol), \
"Diffs are:\n{}.\nanalytic is:\n{}.\nnumeric is:\n{}.".format(
A_flat - B_flat, A_flat, B_flat)
def dict_untake(x, idx, template):
def mut_add(A):
A[idx] = vs.shape[idx].mut_add(A[idx], x)
return A
vs = vspace(template)
return SparseObject(vs, mut_add)
def flatten(value):
"""Flattens any nesting of tuples, arrays, or dicts.
Returns 1D numpy array and an unflatten function.
Doesn't preserve mixed numeric types (e.g. floats and ints).
Assumes dict keys are sortable."""
vs = vspace(value)
return vs.flatten(value), vs.unflatten
def check_args(fun, argnum, args, kwargs):
ans = fun(*args)
in_vspace = vspace(args[argnum])
ans_vspace = vspace(ans)
jac = numerical_jacobian(fun, argnum, args, kwargs)
for outgrad in ans_vspace.examples():
result = fun.vjps[argnum](
outgrad, ans, in_vspace, ans_vspace, *args, **kwargs)
result_vspace = vspace(result)
result_reals = vspace_flatten(result, True)
nd_result_reals = np.dot(vspace_flatten(outgrad, True), jac)
assert result_vspace == in_vspace, \
report_mismatch(fun, argnum, args, kwargs, outgrad,
in_vspace, result_vspace)
assert np.allclose(result_reals, nd_result_reals),\
report_nd_failure(fun, argnum, args, kwargs, outgrad,
nd_result_reals, result_reals)
def match_complex(vs, x):
x_iscomplex = vspace(getval(x)).iscomplex
if x_iscomplex and not vs.iscomplex:
return anp.real(x)
elif not x_iscomplex and vs.iscomplex:
return x + 0j
else:
return x
def check_args(fun, argnum, args, kwargs):
ans = fun(*args)
in_vspace = vspace(args[argnum])
ans_vspace = vspace(ans)
jac = numerical_jacobian(fun, argnum, args, kwargs)
for outgrad in ans_vspace.examples():
result = fun.vjps[argnum](
outgrad, ans, in_vspace, ans_vspace, *args, **kwargs)
result_vspace = vspace(result)
result_reals = vspace_flatten(result, True)
nd_result_reals = np.dot(vspace_flatten(outgrad, True), jac)
assert result_vspace == in_vspace, \
report_mismatch(fun, argnum, args, kwargs, outgrad,
in_vspace, result_vspace)
assert np.allclose(result_reals, nd_result_reals),\
report_nd_failure(fun, argnum, args, kwargs, outgrad,
nd_result_reals, result_reals)
def match_complex(vs, x):
x_iscomplex = vspace(getval(x)).iscomplex
if x_iscomplex and not vs.iscomplex:
return anp.real(x)
elif not x_iscomplex and vs.iscomplex:
return x + 0j
else:
return x
def ggnvp_maker(*args, **kwargs):
f_vjp, f_x = make_vjp(f, f_argnum)(*args, **kwargs)
g_hvp, grad_g_x = make_vjp(grad(g))(f_x)
f_vjp_vjp, _ = make_vjp(f_vjp)(vspace(getval(grad_g_x)).zeros())
def ggnvp(v):
return f_vjp(g_hvp(f_vjp_vjp(v)))
return ggnvp
def array_from_args_fwd_gradmaker(argnum, g, ans, gvs, vs, args, kwargs):
result = list()
for i, arg in enumerate(args):
if i == argnum:
result.append(g)
else:
result.append(
arg.vspace.zeros() if isnode(arg) else vspace(arg).zeros())
return nw.array_from_args(*result)
def unary_nd(f, x, eps=EPS):
vs = vspace(x)
nd_grad = np.zeros(vs.size)
x_flat = vs.flatten(x)
for d in range(vs.size):
dx = np.zeros(vs.size)
dx[d] = eps/2
nd_grad[d] = ( f(vs.unflatten(x_flat + dx))
- f(vs.unflatten(x_flat - dx)) ) / eps
return vs.unflatten(nd_grad, True)
def unary_nd(f, x, eps=EPS):
vs = vspace(x)
nd_grad = np.zeros(vs.size)
x_flat = vs.flatten(x)
for d in range(vs.size):
dx = np.zeros(vs.size)
dx[d] = eps / 2
nd_grad[d] = (f(vs.unflatten(x_flat + dx)) -
f(vs.unflatten(x_flat - dx))) / eps
return vs.unflatten(nd_grad, True)
def forward_mode_pass(v):
ac.assert_vspace_match(v, start_node.vspace, None)
start_node.progenitors[start_node] = v
active_forward_progenitors[start_node] = True
end_node = fun(*args, **kwargs)
active_forward_progenitors.pop(start_node)
if not ac.isnode(end_node) or start_node not in end_node.progenitors:
warnings.warn("Output seems independent of input.")
return (end_node, end_node.vspace.zeros() if ac.isnode(end_node) else
ac.vspace(end_node).zeros())
return end_node, end_node.progenitors[start_node]
def flatten(value):
"""Flattens any nesting of tuples, arrays, or dicts.
Returns 1D numpy array and an unflatten function.
Doesn't preserve mixed numeric types (e.g. floats and ints).
Assumes dict keys are sortable."""
try:
vs = vspace(getval(value))
except TypeError:
raise Exception("Don't know how to flatten type {}".format(
type(value)))
return vs.flatten(value), vs.unflatten
def as_scalar(x):
vs = vspace(getval(x))
if vs.iscomplex:
x = np.real(x)
if vs.shape == ():
return x
elif vs.size == 1:
return x.reshape(())
else:
raise TypeError("Output {} can't be cast to float. "
"Function grad requires a scalar-valued function. "
"Try jacobian or elementwise_grad.".format(getval(x)))
def as_scalar(x):
vs = vspace(getval(x))
if vs.iscomplex:
x = np.real(x)
if vs.shape == ():
return x
elif vs.size == 1:
return x.reshape(())
else:
raise TypeError(
"Output {} can't be cast to float. "
"Function grad requires a scalar-valued function. "
"Try jacobian or elementwise_grad.".format(getval(x)))
def make_node(fun, args, kwargs):
# infer shape data by running fun on dummy arguments
argvals = [
arg.vs.ones() if isinstance(arg, ExprNode) else arg for arg in args
]
vs = vspace(fun(*argvals, **kwargs))
new_node = ExprNode.__new__(ExprNode)
new_node.fun = fun
new_node.args = list(args)
new_node.kwargs = kwargs
new_node.vs = vs
return new_node
def calc_jacobian(start, end):
# if the end_box is not a box - autograd can not track back
if not isbox(end):
return vspace(start.shape).zeros()
# the final jacobian matrices
jac = []
# the backward pass is done for each objective function once
for j in range(end.shape[1]):
b = anp.zeros(end.shape)
b[:, j] = 1
n = new_box(b, 0, VJPNode.new_root())
_jac = backward_pass(n, end._node)
jac.append(_jac)
jac = anp.stack(jac, axis=1)
return jac
def __init__(self, value):
self.shape = [vspace(x) for x in value]
self.size = sum(s.size for s in self.shape)
self.sequence_type = type(value)
assert self.sequence_type in (tuple, list)
def time_vspace_float():
core.vspace(1.)
def jvp_maker(*args, **kwargs):
vjp, y = make_vjp(fun, argnum)(*args, **kwargs)
vjp_vjp, _ = make_vjp(vjp)(vspace(y).zeros())
return vjp_vjp # vjp_vjp is just jvp by linearity
def gradfun(*args,**kwargs):
args = list(args)
args[argnum] = safe_type(args[argnum])
vjp, ans = make_vjp(fun, argnum)(*args, **kwargs)
return vjp(vspace(ans).ones())
def __init__(self, value): self.shape = [ag_core.vspace(x) for x in value] self.size = sum(s.size for s in self.shape) self.sequence_type = type(value)
def __init__(self, value):
self.shape = {k: vspace(v) for k, v in iteritems(value)}
self.size = sum(s.size for s in self.shape.values())
def time_vspace_float():
core.vspace(1.)
def gradfun(*args,**kwargs):
args = list(args)
args[argnum] = safe_type(args[argnum])
vjp, ans = make_vjp(fun, argnum)(*args, **kwargs)
return vjp(vspace(getval(ans)).ones())
def add_preserves_vspace(x, y):
return vs == vspace(add(x, y))
def time_vspace_array():
vspace(A)
nd_result_reals = np.dot(vspace_flatten(outgrad, True), jac)
assert result_vspace == in_vspace, \
report_mismatch(fun, argnum, args, kwargs, outgrad,
in_vspace, result_vspace)
assert np.allclose(result_reals, nd_result_reals),\
report_nd_failure(fun, argnum, args, kwargs, outgrad,
nd_result_reals, result_reals)
def check_primitive(fun, vspace_instances, kwargs={}, argnums=[0]):
arg_sets = [concat([vs.examples() for vs in vsi])
for vsi in vspace_instances]
for argnum, args in it.product(argnums, it.product(*arg_sets)):
check_args(fun, argnum, args, kwargs)
array_shapes = [(), (1,), (1,1), (2,), (3,1), (1,2), (3,2), (2,3,1)]
real_arrays = [vspace(np.zeros(s) ) for s in array_shapes]
complex_arrays = [vspace(np.zeros(s, dtype=complex)) for s in array_shapes]
composite_values = [[0.0], [0.0, np.zeros((2,1))],
[0.0, np.zeros((2,1)), [0.0]]]
lists = list(map(vspace, composite_values))
tuples = list(map(vspace, map(tuple, composite_values)))
dicts = list(map(vspace, [dict(zip(it.count(), x)) for x in composite_values]))
all_arrays = real_arrays + complex_arrays
everything = all_arrays + lists + tuples + dicts
report_mismatch = \
'''
Vspace mismatch
(function {} argnum {} args {}, kwargs {}, outgrad {}):
expected {}
def scalar_mul_preserves_vspace(x, a):
return vs == vspace(scalar_mul(x, a))
def randn_correct_vspace():
return vs == vspace(randn())
def zeros_correct_vspace():
return vs == vspace(zeros())
def fwd_grad_make_sequence(argnum, g, ans, gvs, vs, args, kwargs):
typ, elts = args[0], args[1:]
zeros = list(elt.vspace.zeros() if isnode(elt) else vspace(elt).zeros()
for elt in elts)
zeros[argnum - 1] = g
return ct.make_sequence(typ, *zeros)
def jvp_maker(*args, **kwargs):
vjp, y = make_vjp(fun, argnum)(*args, **kwargs)
vjp_vjp, _ = make_vjp(vjp)(vspace(getval(y)).zeros())
return vjp_vjp # vjp_vjp is just jvp by linearity
def fwd_grad_sequence_extend_right(argnum, g, ans, gvs, vs, args, kwargs):
zeros = list(arg.vspace.zeros() if isnode(arg) else vspace(arg).zeros()
for arg in args)
zeros[argnum] = g
return ct.sequence_extend_right(*zeros)
def jacfun(*args, **kwargs):
vjp, ans = make_vjp(fun, argnum)(*args, **kwargs)
ans_vspace = vspace(getval(ans))
jacobian_shape = ans_vspace.shape + vspace(getval(args[argnum])).shape
grads = map(vjp, ans_vspace.standard_basis())
return np.reshape(np.stack(grads), jacobian_shape)
def ones_correct_vspace():
return vs == vspace(ones())
def __init__(self, value):
self.shape = [vspace(x) for x in value]
self.size = sum(s.size for s in self.shape)
self.sequence_type = type(value)
assert self.sequence_type in (tuple, list)
def time_vspace_array():
core.vspace(A)
def __init__(self, value): self.shape = [ag_core.vspace(x) for x in value] self.size = sum(s.size for s in self.shape) self.sequence_type = type(value)
def time_vspace_float():
vspace(1.)
def ggnvp_maker(*args, **kwargs):
f_vjp, f_x = make_vjp(f, f_argnum)(*args, **kwargs)
g_hvp, grad_g_x = make_vjp(grad(g))(f_x)
f_jvp, _ = make_vjp(f_vjp)(vspace(grad_g_x).zeros())
def ggnvp(v): return f_vjp(g_hvp(f_jvp(v)))
return ggnvp
def basis_correct_vspace():
return (vs == vspace(x) for x in standard_basis())
def jacfun(*args, **kwargs):
vjp, ans = make_vjp(fun, argnum)(*args, **kwargs)
ans_vspace = vspace(ans)
jacobian_shape = ans_vspace.shape + vspace(args[argnum]).shape
grads = map(vjp, ans_vspace.standard_basis())
return np.reshape(np.stack(grads), jacobian_shape)
def assert_vspace_match(x, expected_vspace, fun, fwd=False):
grad_string = "Forward grad" if fwd else "Grad"
assert expected_vspace == ac.vspace(ac.getval(x)), \
"\n{} of {} returned unexpected vector space" \
"\nVector space is {}" \
"\nExpected {}".format(grad_string, fun, ac.vspace(ac.getval(x)), expected_vspace)
def untake(x, idx, template):
def mut_add(A):
np.add.at(A, idx, x)
return A
return SparseObject(vspace(template), mut_add)
def time_vspace_array():
core.vspace(A)
def __init__(self, value):
self.shape = {k : vspace(v) for k, v in iteritems(value)}
self.size = sum(s.size for s in self.shape.values())
def time_vspace_float():
vspace(1.)
nd_result_reals = np.dot(vspace_flatten(outgrad, True), jac)
assert result_vspace == in_vspace, \
report_mismatch(fun, argnum, args, kwargs, outgrad,
in_vspace, result_vspace)
assert np.allclose(result_reals, nd_result_reals),\
report_nd_failure(fun, argnum, args, kwargs, outgrad,
nd_result_reals, result_reals)
def check_primitive(fun, vspace_instances, kwargs={}, argnums=[0]):
arg_sets = [concat([vs.examples() for vs in vsi])
for vsi in vspace_instances]
for argnum, args in it.product(argnums, it.product(*arg_sets)):
check_args(fun, argnum, args, kwargs)
array_shapes = [(), (1,), (1,1), (2,), (3,1), (1,2), (3,2), (2,3,1)]
real_arrays = [vspace(np.zeros(s) ) for s in array_shapes]
complex_arrays = [vspace(np.zeros(s, dtype=complex)) for s in array_shapes]
composite_values = [[0.0], [0.0, np.zeros((2,1))],
[0.0, np.zeros((2,1)), [0.0]]]
lists = list(map(vspace, composite_values))
tuples = list(map(vspace, map(tuple, composite_values)))
dicts = list(map(vspace, [dict(zip(it.count(), x)) for x in composite_values]))
all_arrays = real_arrays + complex_arrays
everything = all_arrays + lists + tuples + dicts
report_mismatch = \
'''
Vspace mismatch
(function {} argnum {} args {}, kwargs {}, outgrad {}):
expected {}
def grad_sequence_take(g, ans, vs, gvs, A, idx):
return sequence_untake(g, idx, vspace(getval(A)))
def time_vspace_array():
vspace(A)
def check_vspace(value):
vs = vspace(value)
# --- required attributes ---
size = vs.size
add = vs.add
scalar_mul = vs.scalar_mul
inner_prod = vs.inner_prod
randn = vs.randn
zeros = vs.zeros
ones = vs.ones
standard_basis = vs.standard_basis
# --- util ---
def randns(N=2):
return [randn() for i in range(N)]
def rand_scalar():
return float(np.random.randn())
def rand_scalars(N=2):
return [rand_scalar() for i in range(N)]
def vector_close(x, y):
z = randn()
return scalar_close(inner_prod(z, x), inner_prod(z, y))
# --- vector space axioms ---
def associativity_of_add(x, y, z):
return vector_close(add(x, add(y, z)),
add(add(x, y), z))
def commutativity_of_add(x, y):
return vector_close(add(x, y), add(y, x))
def identity_element_of_add(x):
return vector_close(add(zeros(), x), x)
def inverse_elements_of_add(x):
return vector_close(zeros(), add(x, scalar_mul(x, -1.0)))
def compatibility_of_scalar_mul_with_field_mul(x, a, b):
return vector_close(scalar_mul(x, a * b),
scalar_mul(scalar_mul(x, a), b))
def identity_element_of_scalar_mul(x):
return vector_close(scalar_mul(x, 1.0), x)
def distributivity_of_scalar_mul_wrt_vector_add(x, y, a):
return vector_close(scalar_mul(add(x, y), a),
add(scalar_mul(x, a),
scalar_mul(y, a)))
def distributivity_of_scalar_mul_wrt_scalar_add(x, a, b):
return vector_close(scalar_mul(x, a + b),
add(scalar_mul(x, a),
scalar_mul(x, b)))
# --- closure ---
def add_preserves_vspace(x, y):
return vs == vspace(add(x, y))
def scalar_mul_preserves_vspace(x, a):
return vs == vspace(scalar_mul(x, a))
# --- inner product axioms ---
def symmetry(x, y): return scalar_close(inner_prod(x, y), inner_prod(y, x))
def linearity(x, y, a): return scalar_close(inner_prod(scalar_mul(x, a), y),
a * inner_prod(x, y))
def positive_definitive(x): return 0 < inner_prod(x, x)
def inner_zeros(): return scalar_close(0, inner_prod(zeros(), zeros()))
# --- basis vectors and special vectors---
def basis_orthonormality():
return all(
[scalar_close(inner_prod(x, y), 1.0 * (ix == iy))
for (ix, x), (iy, y) in it.product(enumerate(standard_basis()),
enumerate(standard_basis()))])
def ones_sum_of_basis_vects():
return vector_close(reduce(add, standard_basis()), ones())
def basis_correct_size():
return len(list(standard_basis())) == size
def basis_correct_vspace():
return (vs == vspace(x) for x in standard_basis())
def zeros_correct_vspace():
return vs == vspace(zeros())
def ones_correct_vspace():
return vs == vspace(ones())
def randn_correct_vspace():
return vs == vspace(randn())
assert associativity_of_add(*randns(3))
assert commutativity_of_add(*randns())
assert identity_element_of_add(randn())
assert inverse_elements_of_add(randn())
assert compatibility_of_scalar_mul_with_field_mul(randn(), *rand_scalars())
assert identity_element_of_scalar_mul(randn())
assert distributivity_of_scalar_mul_wrt_vector_add(randn(), randn(), rand_scalar())
assert distributivity_of_scalar_mul_wrt_scalar_add(randn(), *rand_scalars())
assert add_preserves_vspace(*randns())
assert scalar_mul_preserves_vspace(randn(), rand_scalar())
assert symmetry(*randns())
assert linearity(randn(), randn(), rand_scalar())
assert positive_definitive(randn())
assert inner_zeros()
assert basis_orthonormality()
assert ones_sum_of_basis_vects()
assert basis_correct_size()
assert basis_correct_vspace()
assert zeros_correct_vspace()
assert ones_correct_vspace()
assert randn_correct_vspace()
# --- grads of basic operations ---
check_grads(add)(*randns())
check_grads(scalar_mul)(randn(), rand_scalar())
check_grads(inner_prod)(*randns())
