shard_map fails with jax.random.key but not jax.random.PRNGKey

[jwtkeeble]

Hi, I'm trying to update my code from jax.random.PRNGKey to jax.random.key, however, I'm getting a weird issue with shard_map when trying to use the new jax.random.key datatype.

If you un-comment the lines marked with # uncoment and use the old jax.random.PRNGKey layout, the script works, but with the new jax.random.key layout it gives an error (see below).

Versions:

jax 0.5.0
jaxlib 0.5.0 

Here's a minimal reproducible example,

NCPUS = 16

import os
os.environ["XLA_FLAGS"] = (
    f'--xla_force_host_platform_device_count={NCPUS} '
)

import jax
jax.config.update('jax_enable_x64',True)
jax.config.update('jax_traceback_filtering','off') 

import jax.numpy as jnp
from jax import pmap, vmap

from jax.sharding import PartitionSpec as P
from jax.experimental.shard_map import shard_map
from jax.sharding import Mesh
from jax.experimental import mesh_utils

n_devices = jax.device_count()
devices = mesh_utils.create_device_mesh((n_devices,))
mesh = Mesh(devices.reshape(-1,1), axis_names=('batch','model'))

def func(key, prob):
    print('{}',key)
    return key

def vmap_func(keys, probs):
    probs = vmap(fun=func, in_axes=(0,0), out_axes=(0))(keys, probs)
    return probs

def pmap_func(keys, probs):
    probs = pmap(fun=vmap_func, axis_name=mesh.shape['batch'], in_axes=(0,0), out_axes=(0))(keys, probs)
    return probs

shard_map_func = shard_map(pmap_func, mesh=mesh,
                           in_specs=(P('batch'), P('batch',None)), # chain_keys, probs,
                           out_specs=(P('batch',None)), # probs
                           check_rep=False)
shard_map_func = jax.jit(shard_map_func)

seed=42
key = jax.random.key(seed=seed, impl='threefry2x32')
#key = jax.random.PRNGKey(seed=seed, impl='threefry2x32') # uncomment

key, subkey = jax.random.split(key, 2)
device_keys = jax.random.split(subkey, num=128)

device_keys = device_keys.reshape(NCPUS, 8) 
#device_keys = device_keys.reshape(NCPUS, 8, 2) # uncomment
probs = jnp.zeros(shape=(NCPUS,8), dtype=jnp.float64)

ids = jnp.arange(NCPUS)

out = shard_map_func(device_keys, probs)
print(out.shape)

Here's the stacktrace,

{} Traced<ShapedArray(key<fry>[])>with<BatchTrace> with
  val = Traced<ShapedArray(key<fry>[8])>with<DynamicJaxprTrace>
  batch_dim = 0
loc("jit(pmap_func)/jit(main)/jit(shmap_body)/xla_pmap"(callsite("pmap_func"("/home/user/main/key_shard_map.py":37:0) at "<module>"("/home/user/main/key_shard_map.py":59:0)))): error: has mismatched number of start indices (2) and the rank of operand (3)
Traceback (most recent call last):
  File "/home/user/main/key_shard_map.py", line 59, in <module>
    out = shard_map_func(device_keys, probs)
  File "/home/user/main/key_shard_map.py", line 37, in pmap_func
    probs = pmap(fun=vmap_func, axis_name=mesh.shape['batch'], in_axes=(0,0), out_axes=(0))(keys, probs)
jax._src.source_info_util.JaxStackTraceBeforeTransformation: ValueError: Operation creation failed

The preceding stack trace is the source of the JAX operation that, once transformed by JAX, triggered the following exception.

--------------------

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/home/user/main/key_shard_map.py", line 59, in <module>
    out = shard_map_func(device_keys, probs)
  File "/home/user/pip_venvs/jax_0_4_50/lib/python3.10/site-packages/jax/_src/traceback_util.py", line 180, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "/home/user/pip_venvs/jax_0_4_50/lib/python3.10/site-packages/jax/_src/pjit.py", line 340, in cache_miss
    pgle_profiler) = _python_pjit_helper(fun, jit_info, *args, **kwargs)
  File "/home/user/pip_venvs/jax_0_4_50/lib/python3.10/site-packages/jax/_src/pjit.py", line 198, in _python_pjit_helper
    out_flat, compiled, profiler = _pjit_call_impl_python(*args_flat, **p.params)
  File "/home/user/pip_venvs/jax_0_4_50/lib/python3.10/site-packages/jax/_src/pjit.py", line 1660, in _pjit_call_impl_python
    compiled = _resolve_and_lower(
  File "/home/user/pip_venvs/jax_0_4_50/lib/python3.10/site-packages/jax/_src/pjit.py", line 1627, in _resolve_and_lower
    return _pjit_lower(
  File "/home/user/pip_venvs/jax_0_4_50/lib/python3.10/site-packages/jax/_src/pjit.py", line 1792, in _pjit_lower
    return pxla.lower_sharding_computation(
  File "/home/user/pip_venvs/jax_0_4_50/lib/python3.10/site-packages/jax/_src/profiler.py", line 333, in wrapper
    return func(*args, **kwargs)
  File "/home/user/pip_venvs/jax_0_4_50/lib/python3.10/site-packages/jax/_src/interpreters/pxla.py", line 2330, in lower_sharding_computation
    nreps, tuple_args, shape_poly_state) = _cached_lowering_to_hlo(
  File "/home/user/pip_venvs/jax_0_4_50/lib/python3.10/site-packages/jax/_src/interpreters/pxla.py", line 1960, in _cached_lowering_to_hlo
    lowering_result = mlir.lower_jaxpr_to_module(
  File "/home/user/pip_venvs/jax_0_4_50/lib/python3.10/site-packages/jax/_src/interpreters/mlir.py", line 1194, in lower_jaxpr_to_module
    lower_jaxpr_to_fun(
  File "/home/user/pip_venvs/jax_0_4_50/lib/python3.10/site-packages/jax/_src/interpreters/mlir.py", line 1678, in lower_jaxpr_to_fun
    out_vals, tokens_out = jaxpr_subcomp(
  File "/home/user/pip_venvs/jax_0_4_50/lib/python3.10/site-packages/jax/_src/interpreters/mlir.py", line 1950, in jaxpr_subcomp
    ans = lower_per_platform(rule_ctx, str(eqn.primitive),
  File "/home/user/pip_venvs/jax_0_4_50/lib/python3.10/site-packages/jax/_src/interpreters/mlir.py", line 2068, in lower_per_platform
    output = kept_rules[0](ctx, *rule_args, **rule_kwargs)
  File "/home/user/pip_venvs/jax_0_4_50/lib/python3.10/site-packages/jax/experimental/shard_map.py", line 706, in _shard_map_lowering
    out_nodes_, tokens_out = mlir.call_lowering(
  File "/home/user/pip_venvs/jax_0_4_50/lib/python3.10/site-packages/jax/_src/interpreters/mlir.py", line 2264, in call_lowering
    func_op, output_types, effects = lower_called_computation(
  File "/home/user/pip_venvs/jax_0_4_50/lib/python3.10/site-packages/jax/_src/interpreters/mlir.py", line 2246, in lower_called_computation
    func_op = _lower_jaxpr_to_fun_cached(
  File "/home/user/pip_venvs/jax_0_4_50/lib/python3.10/site-packages/jax/_src/interpreters/mlir.py", line 2206, in _lower_jaxpr_to_fun_cached
    func_op = lower_jaxpr_to_fun(
  File "/home/user/pip_venvs/jax_0_4_50/lib/python3.10/site-packages/jax/_src/interpreters/mlir.py", line 1678, in lower_jaxpr_to_fun
    out_vals, tokens_out = jaxpr_subcomp(
  File "/home/user/pip_venvs/jax_0_4_50/lib/python3.10/site-packages/jax/_src/interpreters/mlir.py", line 1950, in jaxpr_subcomp
    ans = lower_per_platform(rule_ctx, str(eqn.primitive),
  File "/home/user/pip_venvs/jax_0_4_50/lib/python3.10/site-packages/jax/_src/interpreters/mlir.py", line 2068, in lower_per_platform
    output = kept_rules[0](ctx, *rule_args, **rule_kwargs)
  File "/home/user/pip_venvs/jax_0_4_50/lib/python3.10/site-packages/jax/_src/interpreters/pxla.py", line 1536, in _pmap_lowering
    sharded_outs, _ = mlir.jaxpr_subcomp(
  File "/home/user/pip_venvs/jax_0_4_50/lib/python3.10/site-packages/jax/_src/interpreters/pxla.py", line 1529, in <genexpr>
    _hlo_shard(aval, new_env, in_node, in_axis)
  File "/home/user/pip_venvs/jax_0_4_50/lib/python3.10/site-packages/jax/_src/interpreters/pxla.py", line 1443, in _hlo_shard
    dynamic_slice_result = hlo.dynamic_slice(
  File "/home/user/pip_venvs/jax_0_4_50/lib/python3.10/site-packages/jaxlib/mlir/dialects/_stablehlo_ops_gen.py", line 2518, in dynamic_slice
    return DynamicSliceOp(operand=operand, start_indices=start_indices, slice_sizes=slice_sizes, loc=loc, ip=ip).result
  File "/home/user/pip_venvs/jax_0_4_50/lib/python3.10/site-packages/jaxlib/mlir/dialects/_stablehlo_ops_gen.py", line 2492, in __init__
    super().__init__(self.build_generic(attributes=attributes, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip))
ValueError: Operation creation failed

[jakevdp]

The difference between jax.random.key and jax.random.PRNGKey is that in the first, the pair of integers backing the key is implicit, and in the second, the pair of integers representing the key is explicit. This is reflected in the shapes:

In [1]: import jax

In [2]: jax.random.split(jax.random.key(0), 10).shape
Out[2]: (10,)

In [3]: jax.random.split(jax.random.PRNGKey(0), 10).shape
Out[3]: (10, 2)

The sharding specification you used is P('batch'), which is appropriate for one-dimensional arrays, and so it is valid in the first case, but not in the second case.

If you changed your sharding spec to in_specs=(P('batch', None), P('batch',None)), then it would be valid in the second case, but not the first.

Does that make sense?

[jwtkeeble]

Hi @jakevdp, thanks for the quick response! That makes more sense, given the shapes of the keys. But what confuses me is the following.

Let's say I have some arbitrary function which takes a key and a scalar value, called prob, and I have 128 values (for both key and prob) which I want to compute over a set of devices and then within each device I want to vectorize over a device-specific subset.

So, what I had above is,

NCPUS = 8 # except 8 devices rather than 16

import os
os.environ["XLA_FLAGS"] = (
    f'--xla_force_host_platform_device_count={NCPUS} '
)

import jax
jax.config.update('jax_enable_x64',True)
jax.config.update('jax_traceback_filtering','off') 

import jax.numpy as jnp
from jax import pmap, vmap

from jax.sharding import PartitionSpec as P
from jax.experimental.shard_map import shard_map
from jax.sharding import Mesh
from jax.experimental import mesh_utils

n_devices = jax.device_count()
devices = mesh_utils.create_device_mesh((n_devices,))
mesh = Mesh(devices.reshape(-1,1), axis_names=('batch','model'))

def func(key, prob):
    jax.debug.print('{}',key) # my single input function
    return key

def vmap_func(keys, probs):
    probs = vmap(fun=func, in_axes=(0,0), out_axes=(0))(keys, probs) # vectorize over a batch (on each GPU)
    return probs

def pmap_func(keys, probs):
    probs = pmap(fun=vmap_func, axis_name=mesh.shape['batch'], in_axes=(0,0), out_axes=(0))(keys, probs) # handle parrallelisation of each GPU?
    return probs

# Tell JAX how to distribute over each GPU (for inputs and outputs)
shard_map_func = shard_map(pmap_func, mesh=mesh,
                           in_specs=(P('batch'), P('batch',None)), # chain_keys, probs,
                           out_specs=(P('batch',None)), # probs
                           check_rep=False)
shard_map_func = jax.jit(shard_map_func)

seed=42
key = jax.random.key(seed=seed, impl='threefry2x32')

key, subkey = jax.random.split(key, 2)
device_keys = jax.random.split(subkey, num=128)
device_keys = device_keys.reshape(NCPUS, 16)

jax.debug.print('device_keys: {} \n',device_keys.shape)

probs = jnp.zeros(shape=(NCPUS,16), dtype=jnp.float64)

jax.debug.print('probs: {} \n',probs.shape)

out = shard_map_func(device_keys, probs) # fails here
print(out.shape)

The error message I get is the same as above, that being,

device_keys: (Array(8, dtype=int64, weak_type=True), Array(16, dtype=int64, weak_type=True)) 

probs: (Array(8, dtype=int64, weak_type=True), Array(16, dtype=int64, weak_type=True)) 

loc("jit(pmap_func)/jit(main)/jit(shmap_body)/xla_pmap"(callsite("pmap_func"("main.py":33:0) at "<module>"("main.py":55:0)))): error: has mismatched number of start indices (2) and the rank of operand (3)
Traceback (most recent call last):
  File "main.py", line 55, in <module>
    out = shard_map_func(device_keys, probs) # fails here
  File "main.py", line 33, in pmap_func
    probs = pmap(fun=vmap_func, axis_name=mesh.shape['batch'], in_axes=(0,0), out_axes=(0))(keys, probs)
jax._src.source_info_util.JaxStackTraceBeforeTransformation: ValueError: Operation creation failed

The preceding stack trace is the source of the JAX operation that, once transformed by JAX, triggered the following exception.

What confuses me is, why does it state that I have a mismatched number of start indices? I assume it relates to the implicit/expilict shape difference between jax.random.key and jax.random.PRNGKey, but why? Both my input arrays are two-dimension objects, so I can't see (at least within my mind), why JAX states an operand has a rank of 3?

Also, a side-comment about my code example here. Is the wrapping of shard_map over pmap acceptable use here? I used pmap here to remove the trailing device dimension in my JAX array (so vmap only see the 16 probs per device), but I just wanted to check if what I'm doing is correct here (or if it's redundant with shard_map composed over pmap)

Thanks for the help, it's appreciated!

[jakevdp]

I suspect this issue comes from attempting to combine pmap and shard_map. pmap is deprecated, and is incompatible with shard_map. You should prefer to do the whole computation using shard_map.

[jwtkeeble]

Ok, so I should use shard_map to distribute between GPUs, then vmap to vectorise over the subset on each GPU? This is allowed within JAX, right? Thanks for the help!