Fix MMV axis name error and improve single-device support - Fixes #603

[nsrawat0333]

Problem

Issue #603 reported NameError: unbound axis name: i when running MMV (Self-Supervised MultiModal Versatile Networks). This error occurs because the model uses cross-replica batch normalization expecting a pmap axis named 'i', but this axis only exists in multi-device contexts, not single-device runs.

Root Cause Analysis

The error happens in normalization.py where:

1. Cross-replica batch norm defaults to axis name 'i'
2. JAX pmap creates this axis in multi-device setups
3. Single-device runs don't have this axis defined
4. Results in runtime error during model initialization

Solution

Implemented comprehensive fix for both single and multi-device scenarios:

🔧 Smart Normalization Layer (normalization.py)

# Before: Always assumes axis 'i' exists
kwargs['cross_replica_axis'] = 'i'

# After: Intelligent device detection
if len(jax.devices()) > 1:
    kwargs['cross_replica_axis'] = 'i'  # Multi-device
else:
    kwargs['cross_replica_axis'] = None  # Single-device

[polarbe]

Dooray! Failure Notice Failure Notice Your message sent to ***@***.***) has failed to be delivered. Please refer to the below for details. * Recipient :  ***@***.***) * Sent time :  2025-08-11T04:16:03 * Subject :  [google-deepmind/deepmind-research] Fix MMV axis name error and improve single-device support - Fixes #603 (PR #640) * Remote host said :  Your mail was denied from the receiver. This message was sent from a notification-only address that cannot accept incoming email. For more information, please contact ***@***.*** © Dooray!.