38 batch_num *= input_shape.
Dims(dim);
41 const T row_num = input_shape.
Dims(last_dim - 1);
42 const T col_num = input_shape.
Dims(last_dim);
44 if (!(num_lower_diags <= row_num))
45 throw std::runtime_error(
46 "MatrixBandPart : num_lower must be negative or less or equal to number of rows");
48 if (!(num_upper_diags <= col_num))
49 throw std::runtime_error(
50 "MatrixBandPart : num_upper must be negative or less or equal to number of columns");
52 std::fill(output_data, output_data +
output_shape.FlatSize(), 0);
55 for (T batch = 0; batch < batch_num; ++batch)
57 for (T row = 0; row < row_num; ++row)
59 auto output = output_data + (batch * row_num * col_num + row * col_num);
60 auto input = input_data + (batch * row_num * col_num + row * col_num);
63 num_lower_diags < 0 ? 0 : std::min(col_num, std::max(T{0}, row - num_lower_diags));
64 const T band_end = num_upper_diags < 0
66 : std::min(
static_cast<T
>(col_num), row + num_upper_diags + 1);
68 for (T band_idx = band_start; band_idx < band_end; band_idx++)
70 output[band_idx] = input[band_idx];
void MatrixBandPart(const T num_lower_diags, const T num_upper_diags, const Shape &input_shape, const float *input_data, const Shape &output_shape, float *output_data)