!12530 add meshgrid support for start <= stop & complex step with real component
From: @jachua Reviewed-by: @guoqi1024,@liangchenghui Signed-off-by: @liangchenghui
This commit is contained in:
commit
842ca43df3
|
@ -1291,6 +1291,12 @@ def meshgrid(*xi, sparse=False, indexing='xy'):
|
|||
if indexing not in ('xy', 'ij'):
|
||||
_raise_type_error("Valid values for `indexing` are 'xy' and 'ij'.")
|
||||
|
||||
shape_out = ()
|
||||
for x in xi:
|
||||
shape_out += (x.size,)
|
||||
if _is_shape_empty(shape_out):
|
||||
return ones(shape_out)
|
||||
|
||||
grids = []
|
||||
for x in xi:
|
||||
if F.rank(x) == 1:
|
||||
|
@ -1351,7 +1357,7 @@ class nd_grid:
|
|||
else:
|
||||
step = 1
|
||||
if isinstance(step, complex):
|
||||
v = linspace(k.start, k.stop, int(abs(step.imag)))
|
||||
v = linspace(k.start, k.stop, int(abs(step)))
|
||||
else:
|
||||
v = arange(k.start, k.stop, step)
|
||||
xi.append(v)
|
||||
|
@ -1362,6 +1368,8 @@ class nd_grid:
|
|||
if self.sparse:
|
||||
return grids
|
||||
|
||||
if isinstance(grids, Tensor_):
|
||||
return grids
|
||||
expanded = []
|
||||
for grid in grids:
|
||||
expanded.append(F.expand_dims(grid, 0))
|
||||
|
@ -1380,6 +1388,11 @@ class mGridClass(nd_grid):
|
|||
as specifying the number of points to create between the start and
|
||||
stop values, where the stop value is inclusive.
|
||||
|
||||
Note:
|
||||
Unlike Numpy, if the step length is a complex number with a real
|
||||
component, the step length is handled as equivalent to
|
||||
``int(abs(step))``.
|
||||
|
||||
Returns:
|
||||
Tensor or tuple of tensor, a meshgrid.
|
||||
|
||||
|
@ -1422,6 +1435,11 @@ class oGridClass(nd_grid):
|
|||
as specifying the number of points to create between the start and
|
||||
stop values, where the stop value is inclusive.
|
||||
|
||||
Note:
|
||||
Unlike Numpy, if the step length is a complex number with a real
|
||||
component, the step length is handled as equivalent to
|
||||
``int(abs(step))``.
|
||||
|
||||
Raises:
|
||||
TypeError: if slicing indices are not integers.
|
||||
|
||||
|
|
Loading…
Reference in New Issue