OGBG workload variants

[priyakasimbeg]

Add OGBG workload variants:

• SiLU : replace RELU activation functions with SiLU activation functions.
• GELU: replace RELU activation functions with GELU activation functions. Note we explicitly set PyTorch GELU approximate arg to tanh to match the JAX GELU default implementation.
• Model Size:
  • Change num_ message_passing_steps from 5 to 3
  • Add dense layer in update functions by changing hidden_dims for update functions from (256) to (256, 256)
  • Reduce latent_dim from 256 to 128

Add modeldiff test for OGBG. The original test was broken w a KeyError. I rewrote the key_transform for translating PyTorch layer names to JAX layer names so that it generalizes across OGBG variants.

Fixed formulas for calculating input dimensions in update functions in PyTorch OGBG model.

[github-actions]

MLCommons CLA bot All contributors have signed the MLCommons CLA ✍️ ✅

[runame]

LGTM!