pylibcugraphops.dimenet.agg_edge_to_edge_bwd2_grad

pylibcugraphops.dimenet.agg_edge_to_edge_bwd2_grad = <nanobind.nb_func object>

Computes the second-order backward pass for Dimenet++ interaction block

aggregation layer with respect to the gradient of output embeddings.

agg_edge_to_edge_bwd2_grad(
    output_grad_grad_embedding: device array,
    input_vector: device array, input_rbf: device array,
    input_grad_grad_rbf: device array, input_embedding: device array,
    input_grad_grad_embedding: device array, input_weights: device array,
    coo_idx: device array, dst_offsets: device array, dst_edge_index: device array,
    mma_operation: pylibcugraphops.MMAOp, cuda_stream: int = 0
)

We define the following dimensions:

• n_spherical: number of spherical basis functions. Must be 7.

• n_radial: number of radial basis functions. Must be 6.

• n_vec: number of vector/position dimensions. Must be 3.

• n_emb: input/output embedding dimension. Assumed to be at most 64.

• n_mid: we project the spherical basis features twice: once to feature

  dimension n_mid, and then to feature dimension n_emb. Assumed to be at most 8.

Parameters:

output_grad_grad_embeddingdevice array type

Device array containing the output gradients for input gradients (from backward) of output embeddings (from forward). Dimension is assumed to be [#edges, n_emb].

input_vectordevice array type

Device array containing the vector values from forward. Dimension is assumed to be [#edges, n_vec].

input_rbfdevice array type

Device array containing the radial basis features from forward. Dimension is assumed to be [#edges, n_spherical * n_radial].

input_grad_grad_rbfdevice array type

Device array containing the gradients for output gradients (from backward) of radial basis features (from forward). Dimension is assumed to be [#edges, n_spherical * n_radial].

input_embeddingdevice array type

Device array containing the input embeddings values from forward. Dimension is assumed to be [#edges, n_emb].

input_grad_grad_embeddingdevice array type

Device array containing the gradients for output gradients (from backward) of input embeddings (from forward). Dimension is assumed to be [#edges, n_emb].

input_weightsdevice array type

Device array containing the weights used in forward. Dimension is assumed to be [(n_spherical * n_radial + n_emb), n_mid].

coo_idxdevice array type

Device array containing the COO index of the graph. Dimension is assumed to be [2, #edges].

dst_offsetsdevice array type

Device array containing the CSR-like offsets of the destination nodes. Dimension is assumed to be [#nodes + 1].

dst_edge_indexdevice array type

Device array containing the CSR-like indices of mapping the neighbors of destination nodes to edge IDs. Dimension is assumed to be [#edges].

mma_operationpylibcugraphops.MMAOp

MMA precision: pylibcugraphops.MMAOp.HighPrecision performs 3x TF32 operations while pylibcugraphops.MMAOp.LowPrecision performs 1x TF32 MMA operation

cuda_streamint, default=0

CUDA stream as an integer representing the raw pointer