Function LuShapeFn

tensorflow/core/ops/linalg_ops.cc:115–134

Input is [...,N,N]. First and second outputs are: [...,N,N]; [...,N].

// First and second outputs are:
//   [...,N,N]; [...,N].
Status LuShapeFn(InferenceContext* c) {
  ShapeHandle input;
  TF_RETURN_IF_ERROR(c->WithRankAtLeast(c->input(0), 2, &input));

  DimensionHandle n;
  TF_RETURN_IF_ERROR(c->Merge(c->Dim(input, -2), c->Dim(input, -1), &n));

  ShapeHandle batch_shape;
  TF_RETURN_IF_ERROR(c->Subshape(input, 0, -2, &batch_shape));

  ShapeHandle lu_shape;
  ShapeHandle p_shape;

  TF_RETURN_IF_ERROR(c->Concatenate(batch_shape, c->Matrix(n, n), &lu_shape));
  TF_RETURN_IF_ERROR(c->Concatenate(batch_shape, c->Vector(n), &p_shape));

  c->set_output(0, lu_shape);
  c->set_output(1, p_shape);
  return Status::OK();
}

// Input is [...,M,N].
// First and second outputs are: