Searched refs:new_eqn_recipe (Results 1 – 3 of 3) sorted by relevance
| /dports/math/py-jax/jax-0.2.9/jax/interpreters/ |
| H A D | partial_eval.py | 155 eqn = new_eqn_recipe(tracers, out_tracers, primitive, params, source)
160 out_tracer.recipe = new_eqn_recipe(tracers, [out_tracer], primitive,
237 eqn = new_eqn_recipe(in_tracers, unknown_tracers_out, primitive, new_params,
278 eqn = new_eqn_recipe(in_tracers, out_tracers, primitive, new_params,
331 eqn = new_eqn_recipe(in_tracers, out_tracers, prim.initial_style,
370 eqn = new_eqn_recipe(in_tracers, out_tracers, prim.initial_style,
546 def new_eqn_recipe(invars: Sequence[JaxprTracer], function
827 eqn = new_eqn_recipe(in_tracers, unknown_output_tracers, remat_call_p, new_params,
|
| /dports/math/py-jax/jax-0.2.9/jax/_src/lax/ |
| H A D | parallel.py | 1076 eqn = pe.new_eqn_recipe([], [out_tracer], axis_index_p,
|
| H A D | control_flow.py | 910 eqn = pe.new_eqn_recipe(
1643 eqn = pe.new_eqn_recipe(int_res_tracers + new_tracers + ext_res_tracers,
|