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/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 CholeskyGrad(OpKernelConstruction * context)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 ValidateInputMatrixShapes(OpKernelContext * context,const TensorShapes & input_matrix_shapes) const41 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 GetOutputMatrixShapes(const TensorShapes & input_matrix_shapes) const55 TensorShapes GetOutputMatrixShapes( 56 const TensorShapes& input_matrix_shapes) const final { 57 return TensorShapes({input_matrix_shapes[0]}); 58 } 59 ComputeMatrix(OpKernelContext * context,const ConstMatrixMaps & inputs,MatrixMaps * outputs)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_t kMatrixSize = input_matrix_l.rows(); 73 const int64_t kMaxBlockSize = 32; 74 75 for (int64_t 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_t block_begin = 88 std::max(int64_t{0}, block_end - kMaxBlockSize); 89 const int64_t block_size = block_end - block_begin; 90 const int64_t trailing_size = kMatrixSize - block_end; 91 92 auto B = input_matrix_l.block(block_end, 0, trailing_size, block_begin); 93 auto B_bar = 94 output_matrix.block(block_end, 0, trailing_size, block_begin); 95 96 auto C = input_matrix_l.block(block_end, block_begin, trailing_size, 97 block_size); 98 auto C_bar = output_matrix.block(block_end, block_begin, trailing_size, 99 block_size); 100 101 auto D = input_matrix_l.block(block_begin, block_begin, block_size, 102 block_size); 103 auto D_bar = 104 output_matrix.block(block_begin, block_begin, block_size, block_size); 105 106 auto R = input_matrix_l.block(block_begin, 0, block_size, block_begin); 107 auto R_bar = output_matrix.block(block_begin, 0, block_size, block_begin); 108 109 C_bar = D.adjoint() 110 .template triangularView<Eigen::Upper>() 111 .solve(C_bar.adjoint()) 112 .adjoint(); 113 D_bar -= (C_bar.adjoint() * C).template triangularView<Eigen::Lower>(); 114 B_bar -= C_bar * R; 115 R_bar -= C_bar.adjoint() * B; 116 CholeskyGradUnblocked(D, D_bar); 117 R_bar -= (D_bar + D_bar.adjoint()) * R; 118 } 119 output_matrix = (0.5 * (output_matrix + output_matrix.transpose())).eval(); 120 } 121 122 private: CholeskyGradUnblocked(const ConstRef & l_block,Ref grad_block)123 void CholeskyGradUnblocked(const ConstRef& l_block, Ref grad_block) { 124 const int64_t kMatrixSize = l_block.rows(); 125 for (int64_t k = kMatrixSize - 1; k >= 0; k--) { 126 /* This shows the block structure. 127 128 / \ 129 | | 130 | r d | 131 \ B c / 132 133 Variables names representing the derivative matrix have a trailing '_bar'. 134 */ 135 136 const int64_t number_rows_B = kMatrixSize - (k + 1); 137 const int64_t number_rows_r_stack_B = number_rows_B + 1; 138 139 auto r = l_block.block(k, 0, 1, k); 140 auto r_bar = grad_block.block(k, 0, 1, k); 141 auto d = l_block(k, k); // This needs to be a scalar rather than a view. 142 auto d_bar = grad_block.block(k, k, 1, 1); 143 // B is not included explicitly because it is not used on its own. 144 auto B_bar = grad_block.block(k + 1, 0, number_rows_B, k); 145 auto c = l_block.block(k + 1, k, number_rows_B, 1); 146 auto c_bar = grad_block.block(k + 1, k, number_rows_B, 1); 147 // Result of vertical stacking d_bar and c_bar. 148 auto d_stack_c_bar = grad_block.block(k, k, number_rows_r_stack_B, 1); 149 // Result of vertical stacking of r and B. 150 auto r_stack_B = l_block.block(k, 0, number_rows_r_stack_B, k); 151 d_bar -= (c.adjoint() * c_bar) / d; 152 d_stack_c_bar /= d; 153 r_bar -= d_stack_c_bar.adjoint() * r_stack_B; 154 B_bar -= c_bar * r; 155 d_bar /= 2.; 156 } 157 } 158 }; 159 160 REGISTER_LINALG_OP("CholeskyGrad", (CholeskyGrad<float>), float); 161 REGISTER_LINALG_OP("CholeskyGrad", (CholeskyGrad<double>), double); 162 REGISTER_LINALG_OP("BatchCholeskyGrad", (CholeskyGrad<float>), float); 163 REGISTER_LINALG_OP("BatchCholeskyGrad", (CholeskyGrad<double>), double); 164 165 } // namespace tensorflow 166