Searched refs:psum_p (Results 1 – 3 of 3) sorted by relevance
593 psum_p = core.Primitive('psum') variable594 psum_p.multiple_results = True596 pxla.soft_pmap_rules[psum_p] = \599 ad.deflinear2(psum_p, _psum_transpose_rule)601 batching.primitive_batchers[psum_p] = partial(_collective_batcher, psum_p)602 batching.collective_rules[psum_p] = \604 psum_p,610 @psum_p.def_custom_bind1035 out_tup = xla.parallel_translations[psum_p](1063 psum_p.bind = partial(core.Primitive.bind, psum_p) # type: ignore[all …]
343 psum_p,
832 lax_parallel.pmin_p, lax_parallel.ppermute_p, lax_parallel.psum_p,