pylibcugraphops.dimenet.agg_edge_to_edge_bwd

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

Computes the backward pass for Dimenet++ interaction block aggregation layer.

agg_edge_to_edge_bwd(
    output_grad_rbf: device array,
    output_grad_embedding: device array, output_grad_weights: device array,
    input_vector: device array, input_rbf: device array,
    input_embedding: device array, input_grad_embedding: device array,
    input_weights: device array, coo_idx: device array, src_offsets: 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_rbfdevice array type

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

output_grad_embeddingdevice array type

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

output_grad_weightsdevice array type

Device array containing the output gradients for weights. Dimension is assumed to be [(n_spherical * n_radial + n_emb), n_mid].

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_embeddingdevice array type

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

input_grad_embeddingdevice array type

Device array containing the input gradients for output embeddings (of 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].

src_offsetsdevice array type

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

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