MCPcopy Create free account
hub / github.com/zai-org/CodeGeeX / set_tensor_model_parallel_attributes

Function set_tensor_model_parallel_attributes

codegeex/megatron/mpu/layers.py:55–62  ·  view source on GitHub ↗
(tensor, is_parallel, dim, stride)

Source from the content-addressed store, hash-verified

53
54
55def set_tensor_model_parallel_attributes(tensor, is_parallel, dim, stride):
56 # Make sure the attributes are not set.
57 for attribute in _MODEL_PARALLEL_ATTRIBUTE_DEFAULTS:
58 assert not hasattr(tensor, attribute)
59 # Set the attributes.
60 setattr(tensor, "tensor_model_parallel", is_parallel)
61 setattr(tensor, "partition_dim", dim)
62 setattr(tensor, "partition_stride", stride)
63
64
65def set_defaults_if_not_set_tensor_model_parallel_attributes(tensor):

Callers 3

__init__Method · 0.85

Calls

no outgoing calls

Tested by

no test coverage detected