MCPcopy Create free account
hub / github.com/NVIDIA/DALI / get_dali_tensor_gpu

Function get_dali_tensor_gpu

dali/test/python/jax_plugin/utils.py:26–48  ·  view source on GitHub ↗

Helper function to create DALI TensorGPU. Args: value : Value to fill the tensor with. shape : Shape for the tensor. dtype : Data type for the tensor. Returns: TensorGPU: DALI TensorGPU with provided shape and dtype filled with provided value.

(value, shape, dtype, device_id=0)

Source from the content-addressed store, hash-verified

24
25
26def get_dali_tensor_gpu(value, shape, dtype, device_id=0) -> TensorGPU:
27 """Helper function to create DALI TensorGPU.
28
29 Args:
30 value : Value to fill the tensor with.
31 shape : Shape for the tensor.
32 dtype : Data type for the tensor.
33
34 Returns:
35 TensorGPU: DALI TensorGPU with provided shape and dtype filled
36 with provided value.
37 """
38
39 @pipeline_def(num_threads=1, batch_size=1, exec_dynamic=True)
40 def dali_pipeline():
41 values = types.Constant(value=np.full(shape, value, dtype), device="gpu")
42
43 return values
44
45 pipe = dali_pipeline(device_id=device_id)
46 dali_output = pipe.run()
47
48 return dali_output[0][0]
49
50
51def sequential_pipeline(batch_size, shape):

Callers 4

run_sharding_testFunction · 0.90
test_lax_workflowFunction · 0.70

Calls 2

dali_pipelineFunction · 0.70
runMethod · 0.45

Tested by 3

run_sharding_testFunction · 0.72
test_lax_workflowFunction · 0.56