Home
last modified time | relevance | path

Searched refs:NeighborExchange (Results 1 – 7 of 7) sorted by relevance

/third_party/mindspore/mindspore/core/ops/
Dneighborexchange.h30 class MS_CORE_API NeighborExchange : public PrimitiveC {
33 NeighborExchange() : PrimitiveC(kNameNeighborExchange) {} in NeighborExchange() function
35 ~NeighborExchange() = default;
36 MS_DECLARE_PARENT(NeighborExchange, PrimitiveC);
42 using PrimNeighborExchangePtr = std::shared_ptr<NeighborExchange>;
Dneighborexchange.cc190 REGISTER_PRIMITIVE_EVAL_IMPL(NeighborExchange, prim::kPrimNeighborExchange, NeighborExchangeInfer, …
/third_party/mindspore/tests/ut/python/parallel/
Dtest_neighborexchange.py24 from mindspore.ops.operations.comm_ops import NeighborExchange
52 self.alltoallv = NeighborExchange(send_rank_ids=[0, 1], recv_rank_ids=[1, 2],
80 …self.alltoallv = NeighborExchange(send_rank_ids=[0], recv_rank_ids=[1, 2], recv_shapes=([32, 32], …
105 … self.alltoallv = NeighborExchange(send_rank_ids=[], recv_rank_ids=[1], recv_shapes=([1],),
127 self.alltoallv = NeighborExchange(send_rank_ids=[0], recv_rank_ids=[], recv_shapes=(),
149 self.alltoallv = NeighborExchange(send_rank_ids=[], recv_rank_ids=[], recv_shapes=(),
173 …self.alltoallv = NeighborExchange(send_rank_ids=[0], recv_rank_ids=[1, 2], recv_shapes=([32, 32],),
201 self.alltoallv = NeighborExchange(send_rank_ids=[0, 1], recv_rank_ids=[1, 2],
230 self.alltoallv = NeighborExchange(send_rank_ids=[0, 1], recv_rank_ids=[1, 2],
259 …self.alltoallv = NeighborExchange(send_rank_ids=[0], recv_rank_ids=[1, 2], recv_shapes=([32, 32], …
[all …]
/third_party/mindspore/tests/ut/cpp/python_input/gtest_input/pre_activate/
Dall_to_all_unify_mindir_test.py17 from mindspore.ops.operations.comm_ops import AlltoAll, NeighborExchange
37 …neighbor = NeighborExchange(send_rank_ids=[0], recv_rank_ids=[1], recv_shapes=([2, 2],), send_shap…
/third_party/mindspore/mindspore/ops/_grad/
Dgrad_comm_ops.py25 …m_ops import (AllGather, _MiniStepAllGather, _HostAllGather, AllReduce, NeighborExchange, AlltoAll,
360 @bprop_getters.register(NeighborExchange)
369 …neighborexchange_grad = NeighborExchange(send_rank_ids, recv_rank_ids, recv_shapes, send_shapes, r…
/third_party/mindspore/mindspore/ops/operations/
D__init__.py38 from .comm_ops import (AllGather, AllReduce, NeighborExchange, AlltoAll, AllSwap, ReduceScatter, Br…
Dcomm_ops.py626 class NeighborExchange(Primitive): class