Searched refs:_partitionmap (Results 1 – 1 of 1) sorted by relevance
| /dports/math/py-jax/jax-0.2.9/jax/interpreters/ |
| H A D | xla.py | 402 def _partitionmap(func: Callable, vars: Sequence, nodes: Sequence): function
428 _partitionmap(write, [core.unitvar], [_make_unit_constant(c)])
429 _partitionmap(write, jaxpr.constvars, consts)
430 _partitionmap(write, jaxpr.invars, args)
473 _partitionmap(write, eqn.outvars, out_nodes)
|