通过部分求值选择保存中间值还是在求导时重算

1322remat3 = bool_state(1323    name='jax_remat3',1324    default=False,1325    upgrade=True,1326    help='If True, embrace the future of remat.',1327    include_in_jit_key=True,1328    include_in_trace_context=True,1329)

378    jaxpr, consts, out_tree = _trace_to_jaxpr(fun_, in_tree, tuple(in_avals), debug)379    if config.mutable_array_checks.value:380      api_util._check_no_aliased_closed_over_refs(debug, consts, args_flat)381    if isinstance(prevent_cse, tuple):382      cse_args = (tuple(args), kwargs) if kwargs else tuple(args)383      cse = (False,) * len(consts) + tuple(broadcast_prefix(prevent_cse, cse_args))384    else:385      cse = prevent_cse386    out_flat = remat_p.bind(387        *consts, *args_flat, jaxpr=jaxpr, prevent_cse=cse, differentiated=False,388        policy=policy)389    return tree_unflatten(out_tree, out_flat)

83def dots_with_no_batch_dims_saveable(prim, *args, **params) -> bool:84  """This is a useful heuristic for transformers."""85  if prim is lax_internal.dot_general_p:86    (_, _), (lhs_b, rhs_b) = params['dimension_numbers']87    if not lhs_b and not rhs_b:88      return True8990  if prim.name == "scaled_matmul_wrapper":  # avoid cudnn import91    return args[0].shape[0] == 19293  return False

586def remat_partial_eval(trace: pe.JaxprTrace, *tracers: core.Tracer,587                       jaxpr: core.Jaxpr, prevent_cse, **params):588  assert not jaxpr.constvars589  disallowed_effects = effects.remat_allowed_effects.filter_not_in(jaxpr.effects)590  if disallowed_effects:591    raise NotImplementedError(592        f'Effects not supported in AD of `checkpoint`/`remat`: {disallowed_effects}')593  policy = params['policy'] or nothing_saveable594  in_unknowns = [not t.is_known() for t in tracers]595  jaxpr_known, jaxpr_staged, out_unknowns, out_inst, num_res = \596      pe.partial_eval_jaxpr_custom(597          jaxpr, in_unknowns, [True] * len(in_unknowns), False, False, policy)598599  # DCE jaxpr_staged, keeping only instantiated outputs which are unknown600  _, out_inst_unknown = partition_list(out_inst, out_unknowns)601  jaxpr_unknown, in_used_staged = pe.dce_jaxpr(jaxpr_staged, out_inst_unknown)602  used_res, in_used_staged = split_list(in_used_staged, [num_res])603604  # DCE jaxpr_known, keeping all known outputs but discarding dce'd res605  out_used_known = [True] * (len(out_unknowns) - sum(out_unknowns)) + used_res606  jaxpr_known, in_used_known = pe.dce_jaxpr(jaxpr_known, out_used_known)607  num_res = sum(used_res)

609  # To avoid precision mismatches in fwd and bwd passes due to XLA excess610  # precision, insert explicit x = reduce_precision(x, **finfo(x.dtype)) calls611  # on producers of any residuals. See https://github.com/jax-ml/jax/pull/22244.612  jaxpr_known_ = _insert_reduce_precision(jaxpr_known, num_res)613614  # Compute known outputs and residuals (hoisted out of remat primitive)615  _, in_consts_ = unzip2(t.pval for t in tracers if t.pval.is_known())616  _, in_consts = partition_list(in_used_known, in_consts_)617  out_consts = core.eval_jaxpr(jaxpr_known_, (), *in_consts)618  out_knowns, residuals = split_list(out_consts, [len(out_consts)-num_res])

620  # set up unknown outputs with a recipe to call remat621  res_tracers = map(trace.new_instantiated_const, residuals)622  _, tracers_staged = partition_list(in_used_staged, tracers)623  in_jaxpr_tracers = res_tracers + map(trace.instantiate_const, tracers_staged)  # pyrefly: ignore[bad-argument-type]624  out_jaxpr_tracers = [pe.JaxprTracer(trace, pe.PartialVal.unknown(x.aval), None)625                       for x in jaxpr_unknown.outvars]626  if isinstance(prevent_cse, tuple):627    _, prevent_cse_ = partition_list(in_used_staged, prevent_cse)628    prevent_cse = (True,) * len(res_tracers) + tuple(prevent_cse_)629  new_params = dict(params, jaxpr=jaxpr_unknown, differentiated=True,630                    prevent_cse=prevent_cse)631  recipe = pe.new_eqn_recipe(trace, in_jaxpr_tracers, out_jaxpr_tracers, remat_p,632                             new_params, core.positional_effects(jaxpr_unknown),633                             source_info_util.current())