[Feature]: JAX 0.8.0 is supported from rocm 7.2. Need to update rocAL

[LakshmiKumar23]

Suggestion Description

Update rocAL JAX test and Jupyter notebook to support 0.8.0.
PositionalSharding is depracated. Need to update to use NamedSharding.
https://rocm.docs.amd.com/en/latest/compatibility/compatibility-matrix.html
https://docs.jax.dev/en/latest/jax.sharding.html

Operating System

No response

GPU

No response

ROCm Component

No response

[simonCatBot]

Fix for JAX 0.8.0 Support - Replace PositionalSharding with NamedSharding

Understanding the Change

JAX 0.8.0 (supported from ROCm 7.2) deprecated PositionalSharding in favor of NamedSharding. This is part of JAX's move toward more explicit, named axes in parallelism.

Migration Guide

Before (JAX < 0.8.0):

from jax.sharding import PositionalSharding
import jax.numpy as jnp

# Create a simple sharding based on device positions
devices = jax.devices()
sharding = PositionalSharding(devices)

# Apply to array
x = jax.random.normal(jax.random.PRNGKey(0), (8, 8))
x_sharded = jax.device_put(x, sharding.reshape(4, 2))

After (JAX 0.8.0+):

from jax.sharding import NamedSharding, PartitionSpec, Mesh
from jax.lax import with_sharding_constraint
import numpy as np

# Create a mesh with named axes
devices = jax.devices()
mesh = Mesh(np.array(devices).reshape(4, 2), axis_names=('batch', 'model'))

# Create NamedSharding with explicit partition spec
sharding = NamedSharding(mesh, PartitionSpec('batch', 'model'))

# Apply to array
x = jax.random.normal(jax.random.PRNGKey(0), (8, 8))
x_sharded = jax.device_put(x, sharding)

# Or use with_sharding_constraint inside jit
@jax.jit
def f(x):
    y = with_sharding_constraint(x, NamedSharding(mesh, PartitionSpec('batch', None)))
    return y

Complete Refactored Example for rocAL

"""
rocAL JAX integration with JAX 0.8.0+ compatibility
"""

import jax
import jax.numpy as jnp
from jax.sharding import NamedSharding, PartitionSpec, Mesh
import numpy as np

class RocALJAXIterator:
    """Iterator that returns JAX arrays with proper sharding"""
    
    def __init__(self, rocal_pipeline, num_devices=None):
        self.pipeline = rocal_pipeline
        self.num_devices = num_devices or len(jax.devices())
        
        # Create mesh for data parallelism
        devices = jax.devices()
        if len(devices) > 1:
            # Reshape devices for mesh (e.g., 8 GPUs -> 4x2)
            mesh_shape = self._compute_mesh_shape(len(devices))
            self.mesh = Mesh(
                np.array(devices).reshape(mesh_shape),
                axis_names=['batch', 'model'][:len(mesh_shape)]
            )
        else:
            self.mesh = Mesh(np.array(devices).reshape(1, 1), axis_names=['batch'])
    
    def _compute_mesh_shape(self, n):
        """Compute optimal mesh shape for n devices"""
        # Simple factorization
        import math
        sqrt_n = int(math.sqrt(n))
        for i in range(sqrt_n, 0, -1):
            if n % i == 0:
                return (i, n // i)
        return (n, 1)
    
    def __iter__(self):
        return self
    
    def __next__(self):
        # Get data from rocAL pipeline
        data = self.pipeline.get_next()
        
        # Convert to JAX array
        jax_array = jnp.array(data)
        
        # Apply sharding for data parallelism
        if len(jax.devices()) > 1:
            # Shard batch dimension across devices
            sharding = NamedSharding(
                self.mesh,
                PartitionSpec('batch', None, None)  # Only shard batch dim
            )
            jax_array = jax.device_put(jax_array, sharding)
        
        return jax_array


# Usage example
if __name__ == "__main__":
    # Initialize rocAL pipeline
    fromrocal import Pipeline
    
    pipe = Pipeline(
        batch_size=32,
        num_threads=4,
        device_id=0
    )
    
    # Create JAX iterator
    iterator = RocALJAXIterator(pipe)
    
    # Iterate over sharded data
    for batch in iterator:
        print(f"Batch shape: {batch.shape}")
        print(f"Sharding: {batch.sharding}")
        # batch is automatically distributed across devices

Updated Jupyter Notebook Example

# In your rocAL JAX notebook cell:

# Cell 1: Setup with JAX 0.8.0+
import jax
from jax.sharding import Mesh, NamedSharding, PartitionSpec
import numpy as np

# Get available devices
devices = jax.devices()
print(f"Available devices: {devices}")

# Create mesh for sharding
# For 8 GPUs: reshape to (4, 2) or (8, 1) for data parallelism
mesh = Mesh(np.array(devices).reshape(-1, 1), axis_names=['data'])

# Cell 2: Load data with rocAL
from rocal import Pipeline

pipe = Pipeline(
    batch_size=64,
    num_threads=4,
    device_id=0
)

# Cell 3: Shard across devices
def shard_data(data, mesh):
    """Helper to shard rocAL output across JAX devices"""
    import jax.numpy as jnp
    
    jax_array = jnp.array(data)
    
    # Create sharding that splits batch across devices
    sharding = NamedSharding(mesh, PartitionSpec('data', None, None, None))
    
    return jax.device_put(jax_array, sharding)

# Get batch from rocAL
batch = pipe.get_next()
sharded_batch = shard_data(batch, mesh)

print(f"Sharded across: {sharded_batch.sharding}")

# Cell 4: Train step with automatic parallelism
@jax.jit
def train_step(batch, params):
    # JAX automatically handles sharded data
    # No manual device_put needed inside jit
    loss = compute_loss(batch, params)
    grads = jax.grad(loss_fn)(params, batch)
    return grads

# Training loop
for batch in pipe:
    sharded = shard_data(batch, mesh)
    grads = train_step(sharded, params)
    params = update_params(params, grads)

Backward Compatibility Layer

For supporting both JAX versions:

import jax
from packaging import version

def get_sharding(devices, shape):
    """Get sharding compatible with current JAX version"""
    
    if version.parse(jax.__version__) >= version.parse("0.8.0"):
        from jax.sharding import Mesh, NamedSharding, PartitionSpec
        
        mesh = Mesh(np.array(devices).reshape(shape), axis_names=['batch', 'model'])
        return NamedSharding(mesh, PartitionSpec('batch', 'model'))
    else:
        from jax.sharding import PositionalSharding
        
        return PositionalSharding(devices).reshape(shape)

# Usage
sharding = get_sharding(jax.devices(), (4, 2))
x_sharded = jax.device_put(x, sharding)

Testing the Migration

# test_jax_migration.py
import jax
from jax.sharding import NamedSharding, Mesh, PartitionSpec
import numpy as np

def test_sharding_migration():
    """Test that NamedSharding works correctly"""
    devices = jax.devices()
    if len(devices) < 2:
        print("Need at least 2 devices for sharding test")
        return
    
    # Create mesh
    mesh = Mesh(np.array(devices[:2]).reshape(2, 1), axis_names=['x', 'y'])
    
    # Create sharding
    sharding = NamedSharding(mesh, PartitionSpec('x', 'y'))
    
    # Create array
    x = jax.numpy.ones((4, 4))
    
    # Shard it
    x_sharded = jax.device_put(x, sharding)
    
    # Verify sharding
    assert x_sharded.sharding.is_equivalent_to(sharding, ndim=2)
    print("Sharding migration test passed!")

if __name__ == "__main__":
    test_sharding_migration()

Files to Update

1. tests/jax/test_jax_iterator.py - Update iterator
2. notebooks/rocAL_JAX_example.ipynb - Update notebook
3. docs/examples/jax_integration.py - Update documentation
4. requirements.txt - Add minimum JAX version

Verification Commands

# Check JAX version
python -c "import jax; print(jax.__version__)"

# Run tests
pytest tests/jax/ -v

# Verify no deprecation warnings
python -W error::DeprecationWarning -c "from jax.sharding import NamedSharding"