custom_kernels

This is summary documentation for hippynn.custom_kernels module. Click here for full documentation.

API Summary for sub-modules:

• autograd_wrapper
• env_pytorch
• fast_convert
• tensor_wrapper
• utils

Links to Full Documentation:

Custom Kernels for hip-nn interaction sum.

This module provides implementations in pytorch, numba, cupy, and triton.

Pytorch implementations take extra memory, but launch faster than numba kernels. Numba kernels use far less memory, but do come with some launching overhead on GPUs. Cupy kernels only work on the GPU, but are faster than numba. Cupy kernels require numba for CPU operations. Triton custom kernels only work on the GPU, and are generaly faster than CUPY. Triton kernels revert to numba or pytorch as available on module import.

On import, this module attempts to set the custom kernels as specified by the user in hippynn.settings.

Module Attributes

CUSTOM_KERNELS_AVAILABLE	List of available kernel implementations based on currently installed packages..
CUSTOM_KERNELS_ACTIVE	Which custom kernel implementation is currently active.
envsum(sensitivities, features, pair_first, ...)	Computes outer product of sensitivities of pairs and atom features from pair_second, whilst accumulating them onto indices pair_first.
sensesum(env, features, pair_first, pair_second)	Computes product of environment at pair_first with features from pair_second, whilst summing over feature indices.
featsum(env, sense, pair_first, pair_second)	Compute inner product of sensitivities with environment tensor over atoms from pair_first, while accumulating them on to pair_second.

Functions

populate_custom_kernel_availability()	Check available imports and populate the list of available custom kernels.
set_custom_kernels([active])	Activate or deactivate custom kernels for interaction.

Exceptions

CustomKernelError

Modules

autograd_wrapper	Wraps non-pytorch implementations for use with pytorch autograd.
env_pytorch	Pure pytorch implementation of envsum operations.
fast_convert	This module implements a version of converting
tensor_wrapper	Tools for converting between torch and numba-compatible arrays
utils