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

Function _broadcast_nz

codegeex/mindspore/scripts/layer_norm.py:312–327  ·  view source on GitHub ↗

broadcast_nz

(tensor, shape)

Source from the content-addressed store, hash-verified

310
311
312def _broadcast_nz(tensor, shape):
313 """
314 broadcast_nz
315 """
316 broadcast_axes = []
317 src_shape = shape_util.shape_to_list(tensor.shape)
318 for i, _ in enumerate(shape):
319 if shape[i] != src_shape[i]:
320 broadcast_axes.append(i)
321 if len(broadcast_axes) == 2 and \
322 broadcast_axes[1] - broadcast_axes[0] != 1 and \
323 broadcast_axes[1] + 1 == len(shape):
324 temp_shape = src_shape[:-1] + [shape[-1]]
325 tensor = tbe.broadcast(tensor, temp_shape)
326 tensor = tbe.broadcast(tensor, shape)
327 return tensor
328
329
330def _check_vector_to_cube(dtype, ori_shape_x, shape_x, begin_norm_axis, impl_mode):

Callers 2

nz_non_alignedFunction · 0.70
layer_norm_compute_nzFunction · 0.70

Calls

no outgoing calls

Tested by

no test coverage detected