[JAX] Add attention_pair_bias frontend

[Supernova-45]

Adds a public cuequivariance_jax.attention_pair_bias frontend for attention pair bias, mirroring the cuet.attention_pair_bias API.

What it does

• Public API with raw pair input (B, N, N, D_z) and cached pair-projection mode (return_z_proj=True -> pass back with is_cached_z_proj=True), multiplicity over single_repr, optional generalized Q/K projection parameters, and keyword-only precision.
• Portable pure-JAX implementation: CPU and GPU, jax.jit, VJP, and nested vmap; this is also the fallback path.
• When the optional cuequivariance_ops_jax package is installed, eligible CUDA cells (FP16/BF16, head dim divisible by 8 and <= 128, sequence length divisible by 8) are delegated to its triangle-attention FFI backend. Installing the CUDA package is an optimization, not an API requirement.
• Masking is a finite additive penalty (default inf=1e6) so floating masks stay differentiable; fully masked rows remain unsupported per the APB contract.

Tests

• New attention_pair_bias_test.py: parity of raw vs cached forms, gradient checks, JIT/vmap.
• Full cuequivariance_jax test suite passes locally against current main (83 passed; GPU-only skipped).

Related

• Backend (cuequivariance_ops_jax) FFI acceleration is staged separately in kernelcatcher MR !634.

CC @NVIDIA/cuEquivariance-maintainers

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.