def _nest_translation_rule(*args, backend, name, call_jaxpr, scope, **_):
return xla._xla_call_translation_rule( # pylint: disable=protected-access
*args,
name=jax_util.wrap_name(name, f'nest[{scope}]'),
backend=backend,
call_jaxpr=call_jaxpr,
donated_invars=(False, ) * len(args))
def _nest_lowering(ctx, *args, name, call_jaxpr, scope, **_):
return mlir._xla_call_lower( # pylint: disable=protected-access
ctx,
*args,
name=jax_util.wrap_name(name, f'nest[{scope}]'),
call_jaxpr=call_jaxpr,
donated_invars=(False, ) * len(args))
def process_higher_order_primitive(self, primitive, f, tracers, params,
is_map):
name = params.pop('name', f.__name__)
tracers = safe_map(self.instantiate_const, tracers)
vals = [t.val for t in tracers]
context = trace_util.get_dynamic_context(self)
active_tag = context.settings.tag
plants = context.plants
if primitive is nest_p:
plants = plants.get(params['scope'], {})
if is_map:
# TODO(sharadmv): figure out if invars are mapped or unmapped
params = params.copy()
out_axes_thunk = params['out_axes_thunk']
@jax_util.as_hashable_function(key=('harvest', out_axes_thunk))
def new_out_axes_thunk():
out_axes = out_axes_thunk()
assert all(out_axis == 0 for out_axis in out_axes)
return (0, ) * out_tree().num_leaves
new_params = dict(
params,
in_axes=(0, ) * len(tree_util.tree_leaves(plants)) +
params['in_axes'],
out_axes_thunk=new_out_axes_thunk)
else:
new_params = dict(params)
all_args, all_tree = tree_util.tree_flatten((plants, vals))
num_plants = len(all_args) - len(vals)
if 'donated_invars' in params:
new_params['donated_invars'] = ((False, ) * num_plants +
params['donated_invars'])
f, out_tree = harvest_eval(f, self, context.settings, all_tree)
out_flat = primitive.bind(f,
*all_args,
**new_params,
name=jax_util.wrap_name(name, 'harvest'))
out, reaps = tree_util.tree_unflatten(out_tree(), out_flat)
out_tracers = safe_map(self.pure, out)
reap_tracers = tree_util.tree_map(self.pure, reaps)
if primitive is nest_p and reap_tracers:
flat_tracers, tree = tree_util.tree_flatten(reap_tracers)
self.handle_sow(*flat_tracers,
name=params['scope'],
tag=active_tag,
mode='strict',
tree=tree)
else:
for name, reap_tracer in reap_tracers.items():
flat_tracers, tree = tree_util.tree_flatten(reap_tracer)
self.handle_sow(*flat_tracers,
name=name,
tag=active_tag,
mode='strict',
tree=tree)
return out_tracers
def process_call(self, call_primitive, f, tracers, params):
assert call_primitive.multiple_results
heads, tails = unzip2((t.head, t.tail) for t in tracers)
nonzero_tails, in_tree_def = tree_flatten(tails)
f_double, out_tree_def = screen_nones(doubling_subtrace(f, self.main),
len(heads), in_tree_def)
name = params.get('name', f.__name__)
new_params = dict(params, name=wrap_name(name, 'doubledouble'),
donated_invars=(False,) * (len(heads) + len(nonzero_tails)))
result = call_primitive.bind(f_double, *heads, *nonzero_tails, **new_params)
heads_out, tails_out = tree_unflatten(out_tree_def(), result)
return [DoublingTracer(self, h, t) for h, t in zip(heads_out, tails_out)]
def process_higher_order_primitive(self, trace, call_primitive, f, tracers,
params, is_map):
del is_map
name = jax_util.wrap_name(params.pop('name', f.__name__), 'reap')
context = trace_util.get_dynamic_context(trace)
vals = [t.val for t in tracers]
plants = context.plants
if 'in_axes' in params:
# TODO(b/199459308): figure out if invars are mapped or unmapped
params = dict(params,
in_axes=(0, ) * len(tree_util.tree_leaves(plants)) +
params['in_axes'])
if 'donated_invars' in params:
params = dict(params)
params['donated_invars'] = (
(False, ) * len(tree_util.tree_leaves(plants)) +
params['donated_invars'])
elif call_primitive is nest_p:
plants = plants.get(params['scope'], {})
all_vals, all_tree = tree_util.tree_flatten((plants, vals))
f = plant_eval(f, trace, self.settings, all_tree)
out_vals = call_primitive.bind(f, *all_vals, name=name, **params)
return jax_util.safe_map(trace.pure, out_vals)
def reap_higher_order_primitive(self, trace, call_primitive, f, tracers,
params, is_map):
"""Wraps the inner function with a reap trace."""
name = jax_util.wrap_name(params.pop('name', f.__name__), 'reap')
vals = [t.val for t in tracers]
f, aux = reap_eval(f, trace, self.settings)
if is_map:
out_axes_thunk = params['out_axes_thunk']
@jax_util.as_hashable_function(closure=('harvest', out_axes_thunk))
def new_out_axes_thunk():
out_axes = out_axes_thunk()
assert all(out_axis == 0 for out_axis in out_axes)
out_tree, _ = aux()
return (0, ) * out_tree.num_leaves
params = dict(params, out_axes_thunk=new_out_axes_thunk)
out_flat = call_primitive.bind(f, *vals, name=name, **params)
out_tree, metadata = aux()
out_vals, reaps = tree_util.tree_unflatten(out_tree, out_flat)
out_tracers = jax_util.safe_map(trace.pure, out_vals)
reap_tracers = tree_util.tree_map(trace.pure, reaps)
return out_tracers, reap_tracers, metadata
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_pvals = [t.pval for t in tracers]
if is_map:
unknown = pe.PartialVal.unknown
in_pvals = [
pval if pval.is_known() or in_axis is None else unknown(
mapped_aval(params['axis_size'], in_axis, pval[0]))
for pval, in_axis in zip(in_pvals, params['in_axes'])
]
out_axes_thunk = params['out_axes_thunk']
@jax_util.as_hashable_function(closure=('unzip', out_axes_thunk))
def new_out_axes_thunk():
out_axes = out_axes_thunk()
assert all(out_axis == 0 for out_axis in out_axes)
_, num_outputs, _ = aux()
return (0, ) * num_outputs
new_params = dict(params, out_axes_thunk=new_out_axes_thunk)
else:
new_params = params
pvs, in_consts = jax_util.unzip2(t.pval for t in tracers)
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, **new_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,
new_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, new_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=new_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, new_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
