Searched defs:pos_batch (Results 1 – 1 of 1) sorted by relevance
| /dports/math/py-jax/jax-0.2.9/jax/_src/lax/ |
| H A D | parallel.py | 356 def pdot(x, y, axis_name, pos_contract=((), ()), pos_batch=((), ())): argument
990 def _pdot_impl(x, y, *, axis_name, pos_contract, pos_batch): argument
995 def _pdot_abstract_eval(x, y, *, axis_name, pos_contract, pos_batch): argument
1003 pos_contract, pos_batch): argument
1020 pos_batch): argument
1029 def _pdot_translation_rule(c, x, y, *, axis_name, pos_contract, pos_batch, argument
1044 def _pdot_transpose_lhs(g, y, *, axis_name, pos_contract, pos_batch): argument
1049 def _pdot_transpose_rhs(g, x, *, axis_name, pos_contract, pos_batch): argument
|