Python get_tensor_ptr Exemples

Exemple #1

Fichier : gloo_collective_group.py Projet : 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))

Exemple #2

Fichier : gloo_collective_group.py Projet : 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),
     )

Exemple #3

Fichier : gloo_collective_group.py Projet : 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))

Exemple #4

Fichier : gloo_collective_group.py Projet : 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,
     )

Exemple #5

Fichier : gloo_collective_group.py Projet : 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,
     )