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

Function _check_vector_to_cube

codegeex/mindspore/scripts/layer_norm.py:330–350  ·  view source on GitHub ↗

judge case using cube to handle reducesum only supported follow case in Ascend910 and Ascend710: ori_shape: ((batch), m, 1024(768)), "shape": ((batch), 64(48), m//16, 16, 16), "dtype": fp16

(dtype, ori_shape_x, shape_x, begin_norm_axis, impl_mode)

Source from the content-addressed store, hash-verified

328
329
330def _check_vector_to_cube(dtype, ori_shape_x, shape_x, begin_norm_axis, impl_mode):
331 """
332 judge case using cube to handle reducesum
333 only supported follow case in Ascend910 and Ascend710:
334 ori_shape: ((batch), m, 1024(768)), "shape": ((batch), 64(48), m//16, 16, 16), "dtype": fp16
335 """
336
337 def _check_shape_and_dtype():
338 if dtype != "float16":
339 return False
340 if len(ori_shape_x) not in (2, 3) or ori_shape_x[-1] not in (1024, 768, 96, 384, 192, 128, 512, 256):
341 return False
342 if len(shape_x) not in (4, 5) or shape_x[-4] not in (64, 48, 6, 12, 24, 16, 32):
343 return False
344 if "Ascend910" not in get_soc_spec(SOC_VERSION) and "Ascend710" not in get_soc_spec(SOC_VERSION):
345 return False
346 if begin_norm_axis != (len(ori_shape_x) - 1):
347 return False
348 return True
349
350 return impl_mode == "high_performance" and _check_shape_and_dtype()
351
352
353# 'pylint: disable=too-many-locals,too-many-statements,too-many-branches

Callers 1

layer_normFunction · 0.85

Calls 1

_check_shape_and_dtypeFunction · 0.85

Tested by

no test coverage detected