Searched refs:DynamicJaxprTracer (Results 1 – 2 of 2) sorted by relevance
| /dports/math/py-jax/jax-0.2.9/jax/interpreters/ |
| H A D | partial_eval.py | 899 class DynamicJaxprTracer(core.Tracer): class
1014 tracer = DynamicJaxprTracer(self, aval, source_info_util.current())
1022 tracer = DynamicJaxprTracer(self, aval, source_info_util.current())
1061 out_tracers = [DynamicJaxprTracer(self, a, source_info) for a in out_avals]
1074 out_tracers = [DynamicJaxprTracer(self, a, source_info) for a in out_avals]
1104 out_tracers = [DynamicJaxprTracer(self, a, source_info) for a in out_avals]
1129 out_tracers = [DynamicJaxprTracer(self, a) for a in out_avals]
1150 out_tracers = [DynamicJaxprTracer(self, a) for a in out_avals]
|
| /dports/math/py-jax/jax-0.2.9/jax/experimental/ |
| H A D | maps.py | 376 trace_to_subjaxpr_dynamic, DynamicJaxprTracer, source_info_util,
390 out_tracers = [DynamicJaxprTracer(self, a, source_info) for a in out_avals]
|