enerzyme.models.functional

Functions

gather_nd(params, indices)	The same as tf.gather_nd but batched gather is not supported yet.
segment_sum_coo(src, idx[, dim_size])
softplus_inverse(x)	Inverse of the softplus function.