nvalchemi.distributed.to_local

nvalchemi.distributed.to_local(x)

Return the plain local tensor backing a ShardTensor, else x unchanged.

Call this before handing a tensor to a kernel that must not see a ShardTensor. On the halo policy the result is the rank’s owned+ghost block.

Parameters:

x (Any) – A tensor, ShardTensor, or any non-tensor value.

Returns:

x.to_local() for a ShardTensor; x unchanged otherwise (a plain tensor or non-tensor passes straight through, so the call is always safe).

Return type:

Any