MCPcopy Create free account
hub / github.com/apache/tvm / batch_to_space_nd_python

Function batch_to_space_nd_python

python/tvm/topi/testing/batch_to_space_nd.py:25–99  ·  view source on GitHub ↗

Batch to Space operator in python for NHWC layout. Parameters ---------- data : np.ndarray N-D with shape [batch, spatial_shape, remaining_shapes], where spatial_shape has M dimensions. block_shape : list of ints 1-D array of size [M] where M is number of sp

(data, block_shape, crop_begin_list, crop_end_list)

Source from the content-addressed store, hash-verified

23
24
25def batch_to_space_nd_python(data, block_shape, crop_begin_list, crop_end_list):
26 """Batch to Space operator in python for NHWC layout.
27
28 Parameters
29 ----------
30 data : np.ndarray
31 N-D with shape [batch, spatial_shape, remaining_shapes],
32 where spatial_shape has M dimensions.
33
34 block_shape : list of ints
35 1-D array of size [M] where M is number of spatial dims, specifies block
36 size for each spatial dimension.
37
38 crop_begin_list : list of ints
39 list of shape [M] where M is number of spatial dims, specifies
40 begin crop size for each spatial dimension.
41
42 crop_end_list : list of ints
43 list of shape [M] where M is number of spatial dims, specifies
44 end crop size for each spatial dimension.
45
46 Returns
47 -------
48 b2s_out : np.ndarray
49 N-D with shape
50 [batch / prod(block_shape),
51 in_shape[1] * block_shape[0] - crop_begin_list[0] - crop_end_list[0], ...,
52 in_shape[M] * block_shape[M-1] - crop_begin_list[M-1] - crop_end_list[M-1],
53 remaining_shape]
54 """
55 in_shape = data.shape
56 N = len(in_shape)
57 M = len(block_shape)
58 block_shape_prod = np.prod(block_shape)
59 in_batch = data.shape[0]
60 axis = []
61 r_p_shape = []
62
63 r_shape = [block_shape[i] for i in range(0, M)]
64 axis.append(len(r_shape))
65 r_shape.append(in_batch // block_shape_prod)
66
67 for i in range(1, N):
68 axis.append(len(r_shape))
69 if len(axis) < (M + N):
70 axis.append(len(r_shape) - (M + 1))
71 r_shape.append(in_shape[i])
72
73 r_p_shape.append(int(in_batch / block_shape_prod))
74 for i in range(1, M + 1):
75 r_p_shape.append(in_shape[i] * block_shape[i - 1])
76 for i in range(M + 1, N):
77 r_p_shape.append(in_shape[i])
78
79 b2s_out = np.reshape(data, newshape=r_shape)
80 b2s_out = np.transpose(b2s_out, axes=axis)
81 b2s_out = np.reshape(b2s_out, newshape=r_p_shape)
82

Callers

nothing calls this directly

Calls 5

strided_slice_pythonFunction · 0.90
prodMethod · 0.80
appendMethod · 0.45
reshapeMethod · 0.45
transposeMethod · 0.45

Tested by

no test coverage detected

Used in the wild real call sites across dependent graphs

searching dependent graphs…