MCPcopy Create free account
hub / github.com/MegEngine/MegEngine / arange

Function arange

imperative/python/megengine/xla/rules/tensor.py:296–317  ·  view source on GitHub ↗
(
    start: Union[int, float, HLOTensor] = 0,
    stop: Union[int, float, HLOTensor] = None,
    step: Union[int, float, HLOTensor] = 1,
    num: int = None,
    dtype="float32",
)

Source from the content-addressed store, hash-verified

294
295
296def arange(
297 start: Union[int, float, HLOTensor] = 0,
298 stop: Union[int, float, HLOTensor] = None,
299 step: Union[int, float, HLOTensor] = 1,
300 num: int = None,
301 dtype="float32",
302):
303 if stop is None:
304 start, stop = 0, start
305
306 if isinstance(step, int):
307 assert step != 0, "step should not be zero"
308
309 if all([isinstance(x, (int, float)) for x in [start, stop, step]]):
310 assert num is None, "cannot specify num when all args are int or float"
311 num = int(np.ceil((stop - start) / step))
312 stop = start + step * (num - 1)
313 else:
314 assert num is not None, "must specify num when hlotensor exists in args"
315 stop = start + step * (num - 1)
316
317 return linspace(start, stop, num, dtype=dtype)
318
319
320def pad(inp, pad_value, padding):

Callers 2

_get_tableMethod · 0.50

Calls 2

linspaceFunction · 0.70
ceilMethod · 0.45

Tested by

no test coverage detected