jax.lax.axis_size

jax.lax.axis_size(axis_name)

Return the size of the mapped axis axis_name.

Parameters:

axis_name (AxisName) – hashable Python object used to name the mapped axis.

Returns:

An integer representing the size.

Return type:

int

For example, with 8 XLA devices available:

>>> mesh = jax.make_mesh((8,), 'i', axis_types=(jax.sharding.AxisType.Explicit,))
>>> @jax.shard_map(mesh=mesh, in_specs=jax.P('i'), out_specs=jax.P())
... def f(_):
...   return lax.axis_size('i')
...
>>> f(jnp.zeros(16))
Array(8, dtype=int32, weak_type=True)

>>> mesh = jax.make_mesh((4, 2), ('i', 'j'),
...                       axis_types=(jax.sharding.AxisType.Explicit,) * 2)
>>> @jax.shard_map(mesh=mesh, in_specs=jax.P('i', 'j'), out_specs=jax.P())
... def f(_):
...   return lax.axis_size(('i', 'j'))
...
>>> f(jnp.zeros((16, 8)))
Array(8, dtype=int32, weak_type=True)