Python get_tensor_ptrの例

コード例 #1

ファイル: gloo_collective_group.py プロジェクト: wxj0916/ray

 def collective_fn(input_tensor, output_tensor, context):
     pygloo.allreduce(
         context, gloo_util.get_tensor_ptr(input_tensor),
         gloo_util.get_tensor_ptr(output_tensor),
         gloo_util.get_tensor_n_elements(input_tensor),
         gloo_util.get_gloo_tensor_dtype(input_tensor),
         gloo_util.get_gloo_reduce_op(allreduce_options.reduceOp))

コード例 #2

ファイル: gloo_collective_group.py プロジェクト: krfricke/ray

 def collective_fn(input_tensor, output_tensor, context):
     pygloo.allgather(
         context,
         gloo_util.get_tensor_ptr(input_tensor),
         gloo_util.get_tensor_ptr(output_tensor),
         gloo_util.get_tensor_n_elements(input_tensor),
         gloo_util.get_gloo_tensor_dtype(input_tensor),
     )

コード例 #3

ファイル: gloo_collective_group.py プロジェクト: wxj0916/ray

 def collective_fn(input_tensor, output_tensor, context):
     size = gloo_util.get_tensor_n_elements(input_tensor)
     world_size = self._gloo_context.size
     pygloo.reduce_scatter(
         context, gloo_util.get_tensor_ptr(input_tensor),
         gloo_util.get_tensor_ptr(output_tensor), size,
         [size // world_size for _ in range(world_size)],
         gloo_util.get_gloo_tensor_dtype(output_tensor),
         gloo_util.get_gloo_reduce_op(reducescatter_options.reduceOp))

コード例 #4

ファイル: gloo_collective_group.py プロジェクト: krfricke/ray

 def collective_fn(input_tensor, output_tensor, context):
     pygloo.broadcast(
         context,
         gloo_util.get_tensor_ptr(input_tensor),
         gloo_util.get_tensor_ptr(output_tensor),
         gloo_util.get_tensor_n_elements(input_tensor),
         gloo_util.get_gloo_tensor_dtype(input_tensor),
         root_rank,
     )

コード例 #5

ファイル: gloo_collective_group.py プロジェクト: krfricke/ray

 def p2p_fn(tensor, context, peer):
     pygloo.recv(
         context,
         gloo_util.get_tensor_ptr(tensor),
         gloo_util.get_tensor_n_elements(tensor),
         gloo_util.get_gloo_tensor_dtype(tensor),
         peer,
     )