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) – The input tensor.
lower (Union[int, Tensor]) – Number of subdiagonals to keep. If negative, keep entire lower triangle.
upper (Union[int, Tensor]) – Number of superdiagonals to keep. If negative, keep entire upper triangle.

Returns

Tensor

Supported Platforms:: Ascend GPU CPU

Examples

>>> import mindspore
>>> x = mindspore.ops.ones([2, 4, 4])
>>> output = mindspore.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.]]]