Searched refs:ShardedDeviceArray (Results 1 – 2 of 2) sorted by relevance
155 class ShardedDeviceArray : xla::DeviceArrayBase {157 ShardedDeviceArray( in ShardedDeviceArray() function
409 py::class_<ShardedDeviceArray> sda(pmap_lib, "ShardedDeviceArray"); in BuildPmapSubmodule()411 .def_property_readonly("aval", &ShardedDeviceArray::GetAval) in BuildPmapSubmodule()413 &ShardedDeviceArray::GetShardingSpec) in BuildPmapSubmodule()415 &ShardedDeviceArray::GetDeviceBuffers); in BuildPmapSubmodule()