def global_norm(updates: Updates) -> Updates: return jnp.sqrt(sum([jnp.sum(jnp.square(x)) for x in tree_leaves(updates)]))
def _vdot_tree(x, y): return sum(tree_leaves(tree_map(partial( jnp.vdot, precision=lax.Precision.HIGHEST), x, y)))
def _norm(x): xs = tree_leaves(x) return jnp.sqrt(sum(map(_vdot_real_part, xs, xs)))
def testAllLeavesWithTrees(self, tree): leaves = tree_util.tree_leaves(tree) self.assertTrue(tree_util.all_leaves(leaves)) self.assertFalse(tree_util.all_leaves([tree]))
def _iterative_classical_gram_schmidt(Q, x, xnorm, max_iterations=2): """ Orthogonalize x against the columns of Q. The process is repeated up to `max_iterations` times, or fewer if the condition ||r|| < (1/sqrt(2)) ||x|| is met earlier (see below for the meaning of r and x). Parameters ---------- Q : array or tree of arrays A matrix of orthonormal columns. x : array or tree of arrays A vector. It will be replaced with a new vector q which is orthonormal to the columns of Q, such that x in span(col(Q), q). xnorm : float Norm of x. Returns ------- q : array or tree of arrays A unit vector, orthonormal to each column of Q, such that x in span(col(Q), q). r : array Stores the overlaps of x with each vector in Q. """ # "twice is enough" # http://slepc.upv.es/documentation/reports/str1.pdf # TODO(shoyer): consider switching to only one iteration, like SciPy? # This assumes that Q's leaves all have the same dimension in the last # axis. r = jnp.zeros((tree_leaves(Q)[0].shape[-1])) q = x xnorm_scaled = xnorm / jnp.sqrt(2) def body_function(carry): k, q, r, qnorm_scaled = carry h = _project_on_columns(Q, q) Qh = tree_map(lambda X: _dot(X, h), Q) q = _sub(q, Qh) r = _add(r, h) def qnorm_cond(carry): k, not_done, _, _ = carry return jnp.logical_and(not_done, k < (max_iterations - 1)) def qnorm(carry): k, _, q, qnorm_scaled = carry _, qnorm = _safe_normalize(q) qnorm_scaled = qnorm / jnp.sqrt(2) return (k, False, q, qnorm_scaled) init = (k, True, q, qnorm_scaled) _, _, q, qnorm_scaled = lax.while_loop(qnorm_cond, qnorm, init) return (k + 1, q, r, qnorm_scaled) def cond_function(carry): k, _, r, qnorm_scaled = carry _, rnorm = _safe_normalize(r) return jnp.logical_and(k < (max_iterations - 1), rnorm < qnorm_scaled) k, q, r, qnorm_scaled = body_function((0, q, r, xnorm_scaled)) k, q, r, _ = lax.while_loop(cond_function, body_function, (k, q, r, qnorm_scaled)) return q, r
def _infer_shape_jax(f, *vals, **params): avals = map(abstractify, vals) return pe.abstract_eval_fun( lambda *a, **k: tree_util.tree_leaves(f(*a, **k)), *avals, **params)
def flatten(structure, expand_composites=False): """Add expand_composites support for JAX.""" if expand_composites and JAX_MODE: from jax import tree_util # pylint: disable=g-import-not-at-top return tree_util.tree_leaves(structure) return dm_flatten(structure)
def tree_size(tree: PyTree) -> int: """ Returns the sum of the size of all leaves in the tree. It's equivalent to the number of scalars in the pytree. """ return sum(tree_leaves(tree_map(lambda x: x.size, tree)))
def donate_argnums(self): """Flat tuple of donated argument indices.""" return tuple( i for i, x in enumerate(tree_util.tree_leaves(self.args_info)) if x.donated)
def _ApplyGraphNet(graph): """Applies a configured GraphNetwork to a graph. This implementation follows Algorithm 1 in https://arxiv.org/abs/1806.01261 There is one difference. For the nodes update the class aggregates over the sender edges and receiver edges separately. This is a bit more general the algorithm described in the paper. The original behaviour can be recovered by using only the receiver edge aggregations for the update. In addition this implementation supports softmax attention over incoming edge features. Many popular Graph Neural Networks can be implemented as special cases of GraphNets, for more information please see the paper. Args: graph: a `GraphsTuple` containing the graph. Returns: Updated `GraphsTuple`. """ # pylint: disable=g-long-lambda nodes, edges, receivers, senders, globals_, n_node, n_edge = graph # Equivalent to jnp.sum(n_node), but jittable sum_n_node = tree.tree_leaves(nodes)[0].shape[0] sum_n_edge = senders.shape[0] if not tree.tree_all( tree.tree_map(lambda n: n.shape[0] == sum_n_node, nodes)): raise ValueError( 'All node arrays in nest must contain the same number of nodes.' ) sent_attributes = tree.tree_map(lambda n: n[senders], nodes) received_attributes = tree.tree_map(lambda n: n[receivers], nodes) # Here we scatter the global features to the corresponding edges, # giving us tensors of shape [num_edges, global_feat]. global_edge_attributes = tree.tree_map( lambda g: jnp.repeat( g, n_edge, axis=0, total_repeat_length=sum_n_edge), globals_) if update_edge_fn: edges = update_edge_fn(edges, sent_attributes, received_attributes, global_edge_attributes) if attention_logit_fn: logits = attention_logit_fn(edges, sent_attributes, received_attributes, global_edge_attributes) tree_calculate_weights = functools.partial(attention_normalize_fn, segment_ids=receivers, num_segments=sum_n_node) weights = tree.tree_map(tree_calculate_weights, logits) edges = attention_reduce_fn(edges, weights) if update_node_fn: sent_attributes = tree.tree_map( lambda e: aggregate_edges_for_nodes_fn(e, senders, sum_n_node), edges) received_attributes = tree.tree_map( lambda e: aggregate_edges_for_nodes_fn(e, receivers, sum_n_node ), edges) # Here we scatter the global features to the corresponding nodes, # giving us tensors of shape [num_nodes, global_feat]. global_attributes = tree.tree_map( lambda g: jnp.repeat( g, n_node, axis=0, total_repeat_length=sum_n_node), globals_) nodes = update_node_fn(nodes, sent_attributes, received_attributes, global_attributes) if update_global_fn: n_graph = n_node.shape[0] graph_idx = jnp.arange(n_graph) # To aggregate nodes and edges from each graph to global features, # we first construct tensors that map the node to the corresponding graph. # For example, if you have `n_node=[1,2]`, we construct the tensor # [0, 1, 1]. We then do the same for edges. node_gr_idx = jnp.repeat(graph_idx, n_node, axis=0, total_repeat_length=sum_n_node) edge_gr_idx = jnp.repeat(graph_idx, n_edge, axis=0, total_repeat_length=sum_n_edge) # We use the aggregation function to pool the nodes/edges per graph. node_attributes = tree.tree_map( lambda n: aggregate_nodes_for_globals_fn( n, node_gr_idx, n_graph), nodes) edge_attribtutes = tree.tree_map( lambda e: aggregate_edges_for_globals_fn( e, edge_gr_idx, n_graph), edges) # These pooled nodes are the inputs to the global update fn. globals_ = update_global_fn(node_attributes, edge_attribtutes, globals_) # pylint: enable=g-long-lambda return gn_graph.GraphsTuple(nodes=nodes, edges=edges, receivers=receivers, senders=senders, globals=globals_, n_node=n_node, n_edge=n_edge)
def _vdot(x, y): f = partial(jnp.vdot, precision=lax.Precision.HIGHEST) return sum(tree_leaves(tree_multimap(f, x, y)))
def new_body(carry, x): flat_args = tree_leaves((carry, x)) out = body_fun(*(const_vals + flat_args)) out_carry, y = split_list(out, [num_carry]) return out_carry, y
def test_var_tree_flatten(self): newsym = core.gensym() aval = core.ShapedArray((), np.dtype('int32')) a, b, c, d = (newsym(aval), newsym(aval), newsym(aval), newsym(aval)) syms = {c: d, a: b} assert 'bd' == ''.join(map(str, tree_leaves(syms)))
def debug_callback_impl(*flat_args, callback: Callable[..., Any], effect: DebugEffect, in_tree: tree_util.PyTreeDef): del effect args, kwargs = tree_util.tree_unflatten(in_tree, flat_args) out = callback(*args, **kwargs) return tree_util.tree_leaves(out)
def cg(A, b, x0=None, *, tol=1e-5, atol=0.0, maxiter=None, M=None): """Use Conjugate Gradient iteration to solve ``Ax = b``. The numerics of JAX's ``cg`` should exact match SciPy's ``cg`` (up to numerical precision), but note that the interface is slightly different: you need to supply the linear operator ``A`` as a function instead of a sparse matrix or ``LinearOperator``. Derivatives of ``cg`` are implemented via implicit differentiation with another ``cg`` solve, rather than by differentiating *through* the solver. They will be accurate only if both solves converge. Parameters ---------- A : function Function that calculates the matrix-vector product ``Ax`` when called like ``A(x)``. ``A`` must represent a hermitian, positive definite matrix, and must return array(s) with the same structure and shape as its argument. b : array or tree of arrays Right hand side of the linear system representing a single vector. Can be stored as an array or Python container of array(s) with any shape. Returns ------- x : array or tree of arrays The converged solution. Has the same structure as ``b``. info : None Placeholder for convergence information. In the future, JAX will report the number of iterations when convergence is not achieved, like SciPy. Other Parameters ---------------- x0 : array Starting guess for the solution. Must have the same structure as ``b``. tol, atol : float, optional Tolerances for convergence, ``norm(residual) <= max(tol*norm(b), atol)``. We do not implement SciPy's "legacy" behavior, so JAX's tolerance will differ from SciPy unless you explicitly pass ``atol`` to SciPy's ``cg``. maxiter : integer Maximum number of iterations. Iteration will stop after maxiter steps even if the specified tolerance has not been achieved. M : function Preconditioner for A. The preconditioner should approximate the inverse of A. Effective preconditioning dramatically improves the rate of convergence, which implies that fewer iterations are needed to reach a given error tolerance. See also -------- scipy.sparse.linalg.cg jax.lax.custom_linear_solve """ if x0 is None: x0 = tree_map(jnp.zeros_like, b) b, x0 = device_put((b, x0)) if maxiter is None: size = sum(bi.size for bi in tree_leaves(b)) maxiter = 10 * size # copied from scipy if M is None: M = _identity if tree_structure(x0) != tree_structure(b): raise ValueError('x0 and b must have matching tree structure: ' f'{tree_structure(x0)} vs {tree_structure(b)}') if _shapes(x0) != _shapes(b): raise ValueError('arrays in x0 and b must have matching shapes: ' f'{_shapes(x0)} vs {_shapes(b)}') cg_solve = partial(_cg_solve, x0=x0, tol=tol, atol=atol, maxiter=maxiter, M=M) # real-valued positive-definite linear operators are symmetric def real_valued(x): return not issubclass(x.dtype.type, np.complexfloating) symmetric = all(map(real_valued, tree_leaves(b))) x = lax.custom_linear_solve(A, b, solve=cg_solve, transpose_solve=cg_solve, symmetric=symmetric) info = None # TODO(shoyer): return the real iteration count here return x, info
def wait_until_computed(x): for leaf in tree_leaves(x): leaf.block_until_ready()
def test_var_tree_flatten(self): newsym = core.gensym() a, b, c, d = (newsym(core.abstract_unit), newsym(core.abstract_unit), newsym(core.abstract_unit), newsym(core.abstract_unit)) syms = {c: d, a: b} assert 'bd' == ''.join(map(str, tree_leaves(syms)))
def custom_layer_cau_batch(vals, dims, *, num_consts, in_tree, out_tree, kwargs, **params): """Batching rule for layer_cau primitive to handle custom layers.""" if all(dim is batching.not_mapped for dim in dims): return layer_cau_p.bind(*vals, num_consts=num_consts, in_tree=in_tree, out_tree=out_tree, kwargs=kwargs, **params) orig_vals, orig_dims = vals, dims vals, dims = vals[num_consts:], dims[num_consts:] args = tree_util.tree_unflatten(in_tree, vals) dims_ = [not_mapped if dim is None else dim for dim in dims] layer, args = args[0], args[1:] if hasattr(layer, '_call_and_update_batched'): num_params = len(tree_util.tree_leaves(layer)) layer_dims, arg_dims = dims_[:num_params], dims_[num_params:] if kwargs['has_rng']: rng, args = args[0], args[1:] rng_dim, arg_dims = arg_dims[0], arg_dims[1:] mapping_over_layer = all(layer_dim is not not_mapped for layer_dim in layer_dims) mapping_over_args = all(arg_dim is not not_mapped for arg_dim in arg_dims) assert mapping_over_layer or mapping_over_args, (layer_dims, arg_dims) if not mapping_over_layer and mapping_over_args: if kwargs['has_rng']: if rng_dim is not not_mapped: arg_dims = tuple(None if dim is not_mapped else dim for dim in arg_dims) map_fun = jax.vmap( lambda layer, rng, *args: _layer_cau_batched( layer, rng, *args, # pylint: disable=unnecessary-lambda, g-long-lambda **kwargs), in_axes=(None, rng_dim) + (None, ) * len(arg_dims)) else: map_fun = lambda layer, *args: _layer_cau_batched( layer, *args, # pylint: disable=unnecessary-lambda, g-long-lambda **kwargs) vals_out, update_out = map_fun(layer, rng, *args) else: vals_out, update_out = _layer_cau_batched( layer, *args, **kwargs) vals_out = tree_util.tree_leaves(vals_out) update_out = tree_util.tree_leaves(update_out) assert all(dim == 0 for dim in arg_dims) # Assume dimensions out are consistent dims_out = (0, ) * len(vals_out) dims_update = (None, ) * len(update_out) assert len(vals_out) == len(dims_out) assert len(update_out) == len(dims_update) return vals_out + update_out, dims_out + dims_update batched, out_dims = primitive.batch_fun( lu.wrap_init( layer_cau_p.impl, dict(params, num_consts=num_consts, in_tree=in_tree, out_tree=out_tree, kwargs=kwargs)), orig_dims) return batched.call_wrapped(*orig_vals), out_dims()
def wrapped(*args): mapped_args = mapping_fn(*args) ildjs = inverse.ildj(mapping_fn, *args)(mapped_args) return target_log_prob(mapped_args) - np.sum( np.array(tree_util.tree_leaves(ildjs)))
def global_norm(items): return jnp.sqrt(jnp.sum([jnp.sum(x**2) for x in tree_leaves(items)]))
def transpose(res_arg, ct_out): args_flat = tree_leaves((res_arg, ct_out)) ct_ins = core.jaxpr_as_fun(transpose_jaxpr)(*transpose_consts, *args_flat) return tree_unflatten(lin_tree, ct_ins)
def _gmres(A, b, x0=None, *, tol=1e-5, atol=0.0, restart=20, maxiter=None, M=None, qr_mode=False): """ GMRES solves the linear system A x = b for x, given A and b. A is specified as a function performing A(vi) -> vf = A @ vi, and in principle need not have any particular special properties, such as symmetry. However, convergence is often slow for nearly symmetric operators. Parameters ---------- A: function Function that calculates the linear map (matrix-vector product) ``Ax`` when called like ``A(x)``. ``A`` must return array(s) with the same structure and shape as its argument. b : array or tree of arrays Right hand side of the linear system representing a single vector. Can be stored as an array or Python container of array(s) with any shape. Returns ------- x : array or tree of arrays The converged solution. Has the same structure as ``b``. info : None Placeholder for convergence information. In the future, JAX will report the number of iterations when convergence is not achieved, like SciPy. Other Parameters ---------------- x0 : array, optional Starting guess for the solution. Must have the same structure as ``b``. If this is unspecified, zeroes are used. tol, atol : float, optional Tolerances for convergence, ``norm(residual) <= max(tol*norm(b), atol)``. We do not implement SciPy's "legacy" behavior, so JAX's tolerance will differ from SciPy unless you explicitly pass ``atol`` to SciPy's ``gmres``. restart : integer, optional Size of the Krylov subspace ("number of iterations") built between restarts. GMRES works by approximating the true solution x as its projection into a Krylov space of this dimension - this parameter therefore bounds the maximum accuracy achievable from any guess solution. Larger values increase both number of iterations and iteration cost, but may be necessary for convergence. The algorithm terminates early if convergence is achieved before the full subspace is built. Default is 20. maxiter : integer Maximum number of times to rebuild the size-``restart`` Krylov space starting from the solution found at the last iteration. If GMRES halts or is very slow, decreasing this parameter may help. Default is infinite. M : function Preconditioner for A. The preconditioner should approximate the inverse of A. Effective preconditioning dramatically improves the rate of convergence, which implies that fewer iterations are needed to reach a given error tolerance. qr_mode : bool If True, the algorithm builds an internal Krylov subspace using a QR based algorithm, which reduces overhead and improved stability. However, it may degrade performance significantly on GPUs or TPUs, in which case this flag should be set False. See also -------- scipy.sparse.linalg.gmres jax.lax.custom_linear_solve """ if x0 is None: x0 = tree_map(jnp.zeros_like, b) if M is None: M = _identity b, x0 = device_put((b, x0)) size = sum(bi.size for bi in tree_leaves(b)) if maxiter is None: maxiter = 10 * size # copied from scipy restart = min(restart, size) if tree_structure(x0) != tree_structure(b): raise ValueError('x0 and b must have matching tree structure: ' f'{tree_structure(x0)} vs {tree_structure(b)}') b_norm = _norm_tree(b) if b_norm == 0: return b, 0 outer_tol = jnp.maximum(tol * b_norm, atol) Mb = M(b) Mb_norm = _norm_tree(Mb) inner_tol = Mb_norm * min(1.0, outer_tol / b_norm) if qr_mode: def _solve(A, b): return _gmres_solve(A, b, x0, outer_tol, inner_tol, restart, maxiter, M, _gmres_plain) else: def _solve(A, b): return _gmres_solve(A, b, x0, outer_tol, inner_tol, restart, maxiter, M, _gmres_qr) x = lax.custom_linear_solve(A, b, solve=_solve, transpose_solve=_solve) failed = jnp.isnan(_norm_tree(x)) info = jnp.where(failed, x=-1, y=0) return x, info
def _shapes(pytree): return map(jnp.shape, tree_leaves(pytree))
def _vdot_tree(x, y): return sum(tree_leaves(tree_multimap(_vdot, x, y)))
def _vdot_real_tree(x, y): return sum(tree_leaves(tree_map(_vdot_real_part, x, y)))
def _cau_jaxpr(self, *args, **kwargs): flat_args = tree_util.tree_leaves(args) out_flat = eval_jaxpr_with_kwargs(self._jaxpr.jaxpr, self._jaxpr.literals, *flat_args, **kwargs) return tree_util.tree_unflatten(self._out_tree, out_flat)
def gmres(A, b, x0=None, *, tol=1e-5, atol=0.0, restart=20, maxiter=None, M=None, solve_method='batched'): """ GMRES solves the linear system A x = b for x, given A and b. A is specified as a function performing A(vi) -> vf = A @ vi, and in principle need not have any particular special properties, such as symmetry. However, convergence is often slow for nearly symmetric operators. Parameters ---------- A: ndarray or function 2D array or function that calculates the linear map (matrix-vector product) ``Ax`` when called like ``A(x)``. ``A`` must return array(s) with the same structure and shape as its argument. b : array or tree of arrays Right hand side of the linear system representing a single vector. Can be stored as an array or Python container of array(s) with any shape. Returns ------- x : array or tree of arrays The converged solution. Has the same structure as ``b``. info : None Placeholder for convergence information. In the future, JAX will report the number of iterations when convergence is not achieved, like SciPy. Other Parameters ---------------- x0 : array or tree of arrays, optional Starting guess for the solution. Must have the same structure as ``b``. If this is unspecified, zeroes are used. tol, atol : float, optional Tolerances for convergence, ``norm(residual) <= max(tol*norm(b), atol)``. We do not implement SciPy's "legacy" behavior, so JAX's tolerance will differ from SciPy unless you explicitly pass ``atol`` to SciPy's ``gmres``. restart : integer, optional Size of the Krylov subspace ("number of iterations") built between restarts. GMRES works by approximating the true solution x as its projection into a Krylov space of this dimension - this parameter therefore bounds the maximum accuracy achievable from any guess solution. Larger values increase both number of iterations and iteration cost, but may be necessary for convergence. The algorithm terminates early if convergence is achieved before the full subspace is built. Default is 20. maxiter : integer Maximum number of times to rebuild the size-``restart`` Krylov space starting from the solution found at the last iteration. If GMRES halts or is very slow, decreasing this parameter may help. Default is infinite. M : ndarray or function Preconditioner for A. The preconditioner should approximate the inverse of A. Effective preconditioning dramatically improves the rate of convergence, which implies that fewer iterations are needed to reach a given error tolerance. solve_method : 'incremental' or 'batched' The 'incremental' solve method builds a QR decomposition for the Krylov subspace incrementally during the GMRES process using Givens rotations. This improves numerical stability and gives a free estimate of the residual norm that allows for early termination within a single "restart". In contrast, the 'batched' solve method solves the least squares problem from scratch at the end of each GMRES iteration. It does not allow for early termination, but has much less overhead on GPUs. See also -------- scipy.sparse.linalg.gmres jax.lax.custom_linear_solve """ if x0 is None: x0 = tree_map(jnp.zeros_like, b) if M is None: M = _identity A = _normalize_matvec(A) M = _normalize_matvec(M) b, x0 = device_put((b, x0)) size = sum(bi.size for bi in tree_leaves(b)) if maxiter is None: maxiter = 10 * size # copied from scipy restart = min(restart, size) if tree_structure(x0) != tree_structure(b): raise ValueError( 'x0 and b must have matching tree structure: ' f'{tree_structure(x0)} vs {tree_structure(b)}') b_norm = _norm(b) atol = jnp.maximum(tol * b_norm, atol) Mb = M(b) Mb_norm = _norm(Mb) ptol = Mb_norm * jnp.minimum(1.0, atol / b_norm) if solve_method == 'incremental': gmres_func = _gmres_incremental elif solve_method == 'batched': gmres_func = _gmres_batched else: raise ValueError(f"invalid solve_method {solve_method}, must be either " "'incremental' or 'batched'") def _solve(A, b): return _gmres_solve(A, b, x0, atol, ptol, restart, maxiter, M, gmres_func) x = lax.custom_linear_solve(A, b, solve=_solve, transpose_solve=_solve) failed = jnp.isnan(_norm(x)) info = jnp.where(failed, x=-1, y=0) return x, info
def handle_call_primitive(self, call_primitive, f, tracers, params, is_map): """Handler for call_primitives, like jit or layer_call. When an UnzipTracer hits a call primitive, there is either a variable inside of the call primitive, in which case the input function needs to be unzipped into two, or there are no variables in the function, so the call_primitive is recorded in the trace as-is. We use `unzip_eval_wrapper`, which returns whether or not an unzip was successful or not. If it was successful, we record two new Jaxprs into the trace (one for init, one for apply). Otherwise, we just record the Jaxpr corresponding to the function call. Args: call_primitive: a call primitive like xla_call f: a jax.linear_util wrapped function to be called tracers: inputs to the function params: parameters of the primitives is_map: whether or not the primitive is a map primitive (e.g. xla_pmap) Returns: A list of output tracers """ name = params.get('name', f.__name__) settings = trace_util.get_dynamic_context(self).settings tracers = safe_map(self.instantiate_const_abstracted, tracers) if call_primitive in current_custom_rules(): return current_custom_rules()[call_primitive](self, f, *tracers, **params) if call_primitive in pe.call_partial_eval_rules: raise NotImplementedError in_pvs, in_consts = jax_util.unzip2(t.pval for t in tracers) if is_map: pvs = [ None if pv is None else mapped_aval(params['axis_size'], pv) for pv in in_pvs ] else: pvs = in_pvs keys = tuple(t.is_key() for t in tracers) new_settings = UnzipSettings(settings.tag, call_primitive in block_registry) fun, aux = unzip_eval(f, self, keys, tuple(pvs), new_settings) out_flat = call_primitive.bind(fun, *in_consts, **params) success, results = aux() if not success: out_pvs, out_keys, jaxpr, env = results out_pv_consts, consts = jax_util.split_list( out_flat, [len(out_pvs)]) out_tracers = self._bound_output_tracers(call_primitive, params, jaxpr, consts, env, tracers, out_pvs, out_pv_consts, out_keys, name, is_map) return out_tracers init_name = jax_util.wrap_name(name, 'init') apply_name = jax_util.wrap_name(name, 'apply') init_pvs, num_init_consts, apply_pvs = results[0] init_jaxpr, apply_jaxpr = results[1] init_env, apply_env = results[2] variable_names, variable_tree, apply_keys = results[3] key_tracers = [t for t in tracers if t.is_key()] abstract_tracers = [t for t in tracers if not t.is_key()] all_init_consts, all_apply_consts = jax_util.split_list( out_flat, [len(init_pvs) + num_init_consts]) init_pv_consts, init_consts = jax_util.split_list( all_init_consts, [len(init_pvs)]) apply_pv_consts, apply_consts = jax_util.split_list( all_apply_consts, [len(apply_pvs)]) variable_tracers = self._bound_output_tracers(call_primitive, params, init_jaxpr, init_consts, init_env, key_tracers, init_pvs, init_pv_consts, [True] * len(init_pvs), init_name, is_map) unflat_variables = tree_util.tree_unflatten(variable_tree, variable_tracers) if call_primitive is harvest.nest_p: variable_dict = harvest.sow(dict( safe_zip(variable_names, unflat_variables)), tag=settings.tag, name=params['scope'], mode='strict') unflat_variables = tuple(variable_dict[name] for name in variable_names) else: unflat_variables = [ harvest.sow( # pylint: disable=g-complex-comprehension unflat_variable, tag=settings.tag, name=name, mode='strict') for unflat_variable, name in safe_zip( unflat_variables, variable_names) ] variable_tracers = tree_util.tree_leaves(unflat_variables) out_tracers = self._bound_output_tracers( call_primitive, params, apply_jaxpr, apply_consts, apply_env, variable_tracers + abstract_tracers, apply_pvs, apply_pv_consts, apply_keys, apply_name, is_map) return out_tracers
def _scan_harvest_rule(trace: HarvestTrace, *tracers, length, reverse, jaxpr, num_consts, num_carry, linear): """Collects and injects values into/from the scan body.""" context = trace_util.get_dynamic_context(trace) settings = context.settings values = [t.val for t in tracers] consts, init, xs = jax_util.split_list(values, [num_consts, num_carry]) active_sows = _find_sows(jaxpr, settings.tag) active_modes = [params['mode'] for params in active_sows] if any(mode == 'strict' for mode in active_modes): raise ValueError('Cannot use strict mode in a scan.') active_names = [params['name'] for params in active_sows] sow_modes = {name: mode for name, mode in zip(active_names, active_modes)} carry_plants = { name: context.plants[name] for name in active_names if name in context.plants and sow_modes[name] == 'clobber' } xs_plants = { name: context.plants[name] for name in active_names if name in context.plants and sow_modes[name] == 'append' } def jaxpr_fun(carry, x): body_out = jax_core.eval_jaxpr(jaxpr.jaxpr, [], *(consts + carry + x)) carry, y = jax_util.split_list(body_out, [num_carry]) return carry, y harvest_body = harvest(jaxpr_fun, tag=settings.tag, allowlist=settings.allowlist, blocklist=settings.blocklist) def body(carry, x): x_plants, x_vals = x (carry, y), reaps = harvest_body({ **carry_plants, **x_plants }, carry, x_vals) return carry, (y, reaps) xs_flat = tree_util.tree_leaves((xs_plants, xs)) x_avals = [] for x in xs_flat: x_aval = jax_core.get_aval(x) if x_aval is jax_core.abstract_unit: x_avals.append(x_aval) else: x_shape, x_dtype = masking.padded_shape_as_value( x.shape[1:]), x.dtype x_avals.append(abstract_arrays.ShapedArray(x_shape, x_dtype)) x_avals = tuple(x_avals) init_avals = tuple( abstract_arrays.raise_to_shaped(jax_core.get_aval(a)) for a in init) in_flat, in_tree = tree_util.tree_flatten((init, (xs_plants, xs))) body_jaxpr, new_consts, out_tree = ( jax.lax.lax_control_flow._initial_style_jaxpr( # pylint: disable=protected-access body, in_tree, init_avals + x_avals)) new_values = list(new_consts) + in_flat num_xs_plants = len(new_values) - len(init) - len(xs) - len(new_consts) remaining_linear = linear[num_consts:] new_linear = ((False, ) * len(new_consts) + remaining_linear[:len(init)] + (False, ) * num_xs_plants + remaining_linear[len(init):]) assert len(new_linear) == len(new_values) outs = lax.scan_p.bind(*new_values, length=length, reverse=reverse, jaxpr=body_jaxpr, num_consts=len(new_consts), num_carry=num_carry, linear=new_linear) outs = safe_map(trace.pure, outs) carry, (ys, reaps) = tree_util.tree_unflatten(out_tree, outs) out_reaps = {} for k, val in reaps.items(): mode = sow_modes.get(k, 'strict') if mode == 'append': val = tree_util.tree_map(np.concatenate, val) elif mode == 'clobber': val = tree_util.tree_map(lambda x: x[-1], val) out_reaps[k] = sow(val, tag=settings.tag, name=k, mode='strict') (carry, ys) = prim.tie_in(out_reaps, (carry, ys)) return carry + ys
def _plant_scan_rule(trace: HarvestTrace, *tracers, length, reverse, jaxpr, num_consts, num_carry, linear, unroll): """Injects values into a scan according to their sow mode.""" const_tracers, carry_tracers, xs_tracers = jax_util.split_list( tracers, [num_consts, num_carry]) carry_avals, xs_avals = tree_util.tree_map(lambda x: x.aval, (carry_tracers, xs_tracers)) const_vals, carry_vals, xs_vals = tree_util.tree_map( lambda x: x.val, (const_tracers, carry_tracers, xs_tracers)) context = trace_util.get_dynamic_context(trace) settings = context.settings x_tracers = [t[0] if hasattr(t, '_getitem') else t for t in xs_tracers] x_avals = [t.aval for t in x_tracers] metadata = _get_harvest_metadata( jaxpr, settings, *(const_tracers + carry_tracers + x_tracers)) plants = context.plants plant_modes = collections.defaultdict(set) plant_xs_avals = {} for name, meta in metadata.items(): mode = meta['mode'] aval = meta['aval'] if mode == 'strict': raise ValueError( f'Cannot use strict mode for \'{name}\' inside `scan`.') plant_modes[mode].add(name) if mode == 'append' and name in plants: plant_xs_avals[name] = aval body_fun = jax_core.jaxpr_as_fun(jaxpr) clobber_plants = { name: value for name, value in plants.items() if name in plant_modes['clobber'] } append_plants = { name: value for name, value in plants.items() if name in plant_modes['append'] } plant_xs_flat_avals, _ = tree_util.tree_flatten(plant_xs_avals) plant_xs_in_tree = tree_util.tree_structure( (carry_avals, (xs_avals, plant_xs_avals))) def new_body(carry, x): x, plants = x all_plants = {**plants, **clobber_plants} all_values = const_vals + tree_util.tree_leaves((carry, x)) out = plant(body_fun, tag=settings.tag, allowlist=settings.allowlist, blocklist=settings.blocklist, exclusive=settings.exclusive)(all_plants, *all_values) carry_out, y = jax_util.split_list(out, [num_carry]) return carry_out, y new_body_jaxpr, consts, _ = lcf._initial_style_jaxpr( # pylint: disable=protected-access new_body, plant_xs_in_tree, tuple(carry_avals + x_avals + plant_xs_flat_avals)) plant_vals = tree_util.tree_leaves(append_plants) out = lcf.scan_p.bind(*(consts + carry_vals + xs_vals + plant_vals), reverse=reverse, length=length, jaxpr=new_body_jaxpr, num_consts=len(consts), num_carry=num_carry, linear=linear + (False, ) * len(plant_vals), unroll=unroll) return out