mindspore.ops.matrix_band_part

mindspore.ops.matrix_band_part(x, lower, upper)[source]

Copy a tensor setting everything outside a central band in each innermost matrix to zero.

Warning

This is an experimental API that is subject to change or deletion.

Parameters

x (Tensor) – Input tensor. \((*, m, n)\) where \(*\) means, any number of additional dimensions.
lower (Union[int, Tensor]) – Number of subdiagonals to keep. The data type must be int32 or int64. If negative, keep entire lower triangle.
upper (Union[int, Tensor]) – Number of superdiagonals to keep. The data type must be int32 or int64. If negative, keep entire upper triangle.

Returns

Tensor, has the same type and shape as x.

Raises

TypeError – If x is not a Tensor.
TypeError – If dtype of x is not valid.
TypeError – If lower is neither a number nor a Tensor.
TypeError – If upper is neither a number nor a Tensor.
TypeError – If dtype of lower is neither int32 nor int64.
TypeError – If dtype of upper is neither int32 nor int64.
ValueError – If the shape of x is not greater than or equal to 2D.
ValueError – If the shape of lower is not equal to 0D.
ValueError – If the shape of upper is not equal to 0D.

Supported Platforms:: Ascend GPU CPU

Examples

>>> import numpy as np
>>> from mindspore import Tensor, ops
>>> x = Tensor(np.ones([2, 4, 4]).astype(np.float32))
>>> output = ops.matrix_band_part(x, 2, 1)
>>> print(output)
[[[1. 1. 0. 0.]
  [1. 1. 1. 0.]
  [1. 1. 1. 1.]
  [0. 1. 1. 1.]]
 [[1. 1. 0. 0.]
  [1. 1. 1. 0.]
  [1. 1. 1. 1.]
  [0. 1. 1. 1.]]]