easydel.kernels.gpu_ops._utils

easydel.kernels.gpu_ops._utils#

easydel.kernels.gpu_ops._utils.dtype_index(x: array) int[source]#
easydel.kernels.gpu_ops._utils.get_sharding(arr: Union[Array, ndarray, bool, number])[source]#

Gets the sharding of an array.

Parameters

arr – Array to get sharding from.

Returns

Sharding of the array.

easydel.kernels.gpu_ops._utils.get_strides(shape: tuple[int, ...]) tuple[int, ...][source]#

Calculates strides for a given shape.

Parameters

shape – Shape of the array.

Returns

Tuple of strides.

easydel.kernels.gpu_ops._utils.safe_autotune(configs, key, prune_configs_by=None, reset_to_zero=None, restore_value=None, pre_hook=None, post_hook=None, warmup=None, rep=None, use_cuda_graph=False, do_bench=None) Callable[[F], F][source]#

Applies triton.autotune safely. Falls back to the original function if autotuning fails.