Tensorflow LU ops

linalg_ops.cc

/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
    http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

#include "tensorflow/core/framework/common_shape_fns.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/shape_inference.h"

namespace tensorflow {

using shape_inference::DimensionHandle;
using shape_inference::InferenceContext;
using shape_inference::ShapeHandle;

namespace {

// Input is [...,M,M].
// First and second outputs are:
//   [...,M,M]; [...,M,M]; [...,M,M]; [...,M,M]
Status ShapeFn(InferenceContext* c) {
  ShapeHandle l_shape;
  ShapeHandle u_shape;
  ShapeHandle p_shape;
  ShapeHandle q_shape;
  c->set_output(0, l_shape);
  c->set_output(1, u_shape);
  c->set_output(2, p_shape);
  c->set_output(3, q_shape);
  return Status::OK();
}


REGISTER_OP("Lu")
    .Input("input: T")
    .Output("l: T")
    .Output("u: T")
    .Output("p: T")
    .Output("q: T")
    .Attr("T: {double, float, complex64, complex128}")
    .SetShapeFn(LuShapeFn)  
    .Doc(R"doc(
Computes the LU decomposition of one or more square matrices.
The input is a tensor of shape `[..., M, M]` whose inner-most 2 dimensions
form square matrices.
```python
# l is a tensor of lower triangular matrices.
# u is a tensor of upper triangular matrices.
# p is a tensor.
# q is a tensor.
# a = p.inverse()*l*u*q.inverse()
l, u, p, q = lu(a) 
```
The input has to be square matrix. 
The output is a tensor of the same shape as the input
containing the LU decompositions for all input submatrices `[..., :, :]`.
input: Shape is `[..., M, M]`.
l: Shape is `[..., M, M]`.
u: Shape is `[..., M, M]`.
p: Shape is `[..., M, M]`.
q: Shape is `[..., M, M]`.
)doc");

}  // namespace tensorflow