1 /* Copyright 2015 The TensorFlow Authors. All Rights Reserved. 2 3 Licensed under the Apache License, Version 2.0 (the "License"); 4 you may not use this file except in compliance with the License. 5 You may obtain a copy of the License at 6 7 http://www.apache.org/licenses/LICENSE-2.0 8 9 Unless required by applicable law or agreed to in writing, software 10 distributed under the License is distributed on an "AS IS" BASIS, 11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 See the License for the specific language governing permissions and 13 limitations under the License. 14 ==============================================================================*/ 15 16 #include "third_party/eigen3/Eigen/Core" 17 #include "tensorflow/core/framework/op.h" 18 #include "tensorflow/core/framework/op_kernel.h" 19 #include "tensorflow/core/framework/tensor_types.h" 20 #include "tensorflow/core/framework/types.h" 21 #include "tensorflow/core/kernels/linalg_ops_common.h" 22 23 namespace tensorflow { 24 25 template <typename Scalar> 26 class CholeskyGrad : public LinearAlgebraOp<Scalar> { 27 public: 28 typedef LinearAlgebraOp<Scalar> Base; 29 30 explicit CholeskyGrad(OpKernelConstruction* context) : Base(context) {} 31 32 using TensorShapes = typename Base::TensorShapes; 33 using Matrix = typename Base::Matrix; 34 using MatrixMap = typename Base::MatrixMap; 35 using MatrixMaps = typename Base::MatrixMaps; 36 using ConstMatrixMap = typename Base::ConstMatrixMap; 37 using ConstMatrixMaps = typename Base::ConstMatrixMaps; 38 using ConstRef = Eigen::Ref<const Matrix>; 39 using Ref = Eigen::Ref<Matrix>; 40 41 void ValidateInputMatrixShapes( 42 OpKernelContext* context, 43 const TensorShapes& input_matrix_shapes) const final { 44 OP_REQUIRES(context, input_matrix_shapes.size() == 2, 45 errors::InvalidArgument("Expected two input matrices, got %d.", 46 input_matrix_shapes.size())); 47 OP_REQUIRES(context, input_matrix_shapes[0] == input_matrix_shapes[1], 48 errors::InvalidArgument( 49 "Inputs (L and grad) must have the same shape.")); 50 OP_REQUIRES(context, 51 TensorShapeUtils::IsSquareMatrix(input_matrix_shapes[0]), 52 errors::InvalidArgument("Inputs must be a square matrices.")); 53 } 54 55 TensorShapes GetOutputMatrixShapes( 56 const TensorShapes& input_matrix_shapes) const final { 57 return TensorShapes({input_matrix_shapes[0]}); 58 } 59 60 void ComputeMatrix(OpKernelContext* context, const ConstMatrixMaps& inputs, 61 MatrixMaps* outputs) final { 62 const ConstMatrixMap& input_matrix_l_full = inputs[0]; 63 const ConstMatrixMap& input_matrix_grad = inputs[1]; 64 MatrixMap output_matrix = outputs->at(0); 65 66 // Algorithm only depends on lower triangular half on input_matrix_l. 67 const Matrix input_matrix_l = 68 input_matrix_l_full.template triangularView<Eigen::Lower>(); 69 // Algorithm only depends on lower triangular half on input_matrix_grad. 70 output_matrix = input_matrix_grad.template triangularView<Eigen::Lower>(); 71 72 const int64 kMatrixSize = input_matrix_l.rows(); 73 const int64 kMaxBlockSize = 32; 74 75 for (int64 block_end = kMatrixSize; block_end > 0; 76 block_end -= kMaxBlockSize) { 77 /* This shows the block structure. 78 79 / \ 80 | | 81 | R D | 82 \ B C / 83 84 Variables names representing the derivative matrix have a trailing '_bar'. 85 */ 86 87 const int64 block_begin = std::max(0ll, block_end - kMaxBlockSize); 88 const int64 block_size = block_end - block_begin; 89 const int64 trailing_size = kMatrixSize - block_end; 90 91 auto B = input_matrix_l.block(block_end, 0, trailing_size, block_begin); 92 auto B_bar = 93 output_matrix.block(block_end, 0, trailing_size, block_begin); 94 95 auto C = input_matrix_l.block(block_end, block_begin, trailing_size, 96 block_size); 97 auto C_bar = output_matrix.block(block_end, block_begin, trailing_size, 98 block_size); 99 100 auto D = input_matrix_l.block(block_begin, block_begin, block_size, 101 block_size); 102 auto D_bar = 103 output_matrix.block(block_begin, block_begin, block_size, block_size); 104 105 auto R = input_matrix_l.block(block_begin, 0, block_size, block_begin); 106 auto R_bar = output_matrix.block(block_begin, 0, block_size, block_begin); 107 108 C_bar = D.adjoint() 109 .template triangularView<Eigen::Upper>() 110 .solve(C_bar.adjoint()) 111 .adjoint(); 112 D_bar -= (C_bar.adjoint() * C).template triangularView<Eigen::Lower>(); 113 B_bar -= C_bar * R; 114 R_bar -= C_bar.adjoint() * B; 115 CholeskyGradUnblocked(D, D_bar); 116 R_bar -= (D_bar + D_bar.adjoint()) * R; 117 } 118 output_matrix = (0.5 * (output_matrix + output_matrix.transpose())).eval(); 119 } 120 121 private: 122 void CholeskyGradUnblocked(const ConstRef& l_block, Ref grad_block) { 123 const int64 kMatrixSize = l_block.rows(); 124 for (int64 k = kMatrixSize - 1; k >= 0; k--) { 125 /* This shows the block structure. 126 127 / \ 128 | | 129 | r d | 130 \ B c / 131 132 Variables names representing the derivative matrix have a trailing '_bar'. 133 */ 134 135 const int64 number_rows_B = kMatrixSize - (k + 1); 136 const int64 number_rows_r_stack_B = number_rows_B + 1; 137 138 auto r = l_block.block(k, 0, 1, k); 139 auto r_bar = grad_block.block(k, 0, 1, k); 140 auto d = l_block(k, k); // This needs to be a scalar rather than a view. 141 auto d_bar = grad_block.block(k, k, 1, 1); 142 // B is not included explicitly because it is not used on its own. 143 auto B_bar = grad_block.block(k + 1, 0, number_rows_B, k); 144 auto c = l_block.block(k + 1, k, number_rows_B, 1); 145 auto c_bar = grad_block.block(k + 1, k, number_rows_B, 1); 146 // Result of vertical stacking d_bar and c_bar. 147 auto d_stack_c_bar = grad_block.block(k, k, number_rows_r_stack_B, 1); 148 // Result of vertical stacking of r and B. 149 auto r_stack_B = l_block.block(k, 0, number_rows_r_stack_B, k); 150 d_bar -= (c.adjoint() * c_bar) / d; 151 d_stack_c_bar /= d; 152 r_bar -= d_stack_c_bar.adjoint() * r_stack_B; 153 B_bar -= c_bar * r; 154 d_bar /= 2.; 155 } 156 } 157 }; 158 159 REGISTER_LINALG_OP("CholeskyGrad", (CholeskyGrad<float>), float); 160 REGISTER_LINALG_OP("CholeskyGrad", (CholeskyGrad<double>), double); 161 REGISTER_LINALG_OP("BatchCholeskyGrad", (CholeskyGrad<float>), float); 162 REGISTER_LINALG_OP("BatchCholeskyGrad", (CholeskyGrad<double>), double); 163 164 } // namespace tensorflow 165
