tensorflow/compiler/tf2xla/kernels/diag_op.cc
// Create a diagonal / batch diagonal matrix with 'input' on the diagonal.
Create a diagonal / batch diagonal matrix with 'input' on the diagonal.
Create a diagonal / batch diagonal matrix with ' input ' on the diagonal .
xla::XlaOp CreateDiagonal(xla::XlaOp input, int64_t last_dim_size,
                          absl::Span<const int64_t> other_dims) {
  xla::XlaBuilder* builder = input.builder();
   [[0, 0, 0, 0]            [[0, 1, 2, 3]
    [1, 1, 1, 1]             [0, 1, 2, 3]
    [2, 2, 2, 2]             [0, 1, 2, 3]
    [3, 3, 3, 3]]            [0, 1, 2, 3]]
   This produces a predicate matrix of the right size, with "true" on the
   diagonal.
  xla::XlaOp iota = xla::Iota(builder, xla::S32, last_dim_size);
  xla::XlaOp iota_broadcast = xla::Broadcast(iota, {last_dim_size});
  xla::XlaOp mask = xla::Eq(iota_broadcast, iota, {0});
   If this is a batched diagonal, broadcast the mask across the other
   dimensions.
  if (!other_dims.empty()) {
    mask = xla::Broadcast(mask, other_dims);
  }
   Broadcast the input, and then use the mask computed above to select the
   diagonal:
   e.g, in 2D:
           [[t, f, f]    [[1, 1, 1]    [[0, 0, 0]      [[1, 0, 0]
   select(  [f, t, f]  ,  [4, 4, 4]  ,  [0, 0, 0]  ) =  [0, 4, 0]
            [f, f, t]]    [9, 9, 9]]    [0, 0, 0]]      [0, 0, 9]]
  std::vector<int64_t> out_dim_sizes(other_dims.begin(), other_dims.end());
  out_dim_sizes.push_back(last_dim_size);
  out_dim_sizes.push_back(last_dim_size);
   Broadcast into the second to last dimension.
  std::vector<int64_t> broadcast_dimensions(other_dims.size() + 1);
  absl::c_iota(broadcast_dimensions, 0);
  ++broadcast_dimensions.back();
  xla::XlaOp input_broadcast =
      xla::BroadcastInDim(input, out_dim_sizes, broadcast_dimensions);
  return xla::Select(mask, input_broadcast, xla::ZerosLike(input_broadcast));
}
xla :: XlaOp CreateDiagonal ( xla :: XlaOp input , int64_t last_dim_size , absl :: Span < const int64_t > other_dims ) { xla :: XlaBuilder * builder = input . builder ( ) ; xla :: XlaOp iota = xla :: Iota ( builder , xla :: S32 , last_dim_size ) ; xla :: XlaOp iota_broadcast = xla :: Broadcast ( iota , { last_dim_size } ) ; xla :: XlaOp mask = xla :: Eq ( iota_broadcast , iota , { 0 } ) ; if ( ! other_dims . empty ( ) ) { mask = xla :: Broadcast ( mask , other_dims ) ; } std :: vector < int64_t > out_dim_sizes ( other_dims . begin ( ) , other_dims . end ( ) ) ; out_dim_sizes . push_back ( last_dim_size ) ; out_dim_sizes . push_back ( last_dim_size ) ; std :: vector < int64_t > broadcast_dimensions ( other_dims . size ( ) + 1 ) ; absl :: c_iota ( broadcast_dimensions , 0 ) ; ++ broadcast_dimensions . back ( ) ; xla :: XlaOp input_broadcast = xla :: BroadcastInDim ( input , out_dim_sizes , broadcast_dimensions ) ; return xla :: Select ( mask , input_broadcast , xla :: ZerosLike ( input_broadcast ) ) ; }
Create a diagonal / batch diagonal matrix with 'input' on the diagonal.
Create a diagonal / batch diagonal matrix with ' input ' on the diagonal .
// Create two matrices that have the following forms, and compare them: // // [[0, 0, 0, 0]            [[0, 1, 2, 3] //  [1, 1, 1, 1]             [0, 1, 2, 3] //  [2, 2, 2, 2]             [0, 1, 2, 3] //  [3, 3, 3, 3]]            [0, 1, 2, 3]] // // This produces a predicate matrix of the right size, with "true" on the // diagonal. // If this is a batched diagonal, broadcast the mask across the other // dimensions. // Broadcast the input, and then use the mask computed above to select the // diagonal: // e.g, in 2D: //         [[t, f, f]    [[1, 1, 1]    [[0, 0, 0]      [[1, 0, 0] // select(  [f, t, f]  ,  [4, 4, 4]  ,  [0, 0, 0]  ) =  [0, 4, 0] //          [f, f, t]]    [9, 9, 9]]    [0, 0, 0]]      [0, 0, 9]] // // Broadcast into the second to last dimension.