ShardedArray

class brainpy.math.ShardedArray(value, dtype=None, *, keep_sharding=True)

The sharded array, which stores data across multiple devices.

A drawback of sharding is that the data may not be evenly distributed on shards.

Parameters:

• value – the array value.

• dtype (Optional[Any]) – the array type.

• keep_sharding (bool) – keep the array sharding information using jax.lax.with_sharding_constraint. Default True.

property value

The value stored in this array.

Returns:

The stored data.