Function TridiagonalMatMulShapeFn

tensorflow/core/ops/linalg_ops.cc:213–257

Inputs: [...,1,M], [...,1,M], [...,1,M],[...,M,N]. Output is [...,M,N].

// Inputs: [...,1,M], [...,1,M], [...,1,M],[...,M,N].
// Output is [...,M,N].
Status TridiagonalMatMulShapeFn(InferenceContext* c) {
  ShapeHandle superdiag;
  ShapeHandle maindiag;
  ShapeHandle subdiag;
  ShapeHandle rhs;

  // Check that rank is at least 2.
  TF_RETURN_IF_ERROR(c->WithRankAtLeast(c->input(0), 2, &superdiag));
  TF_RETURN_IF_ERROR(c->WithRankAtLeast(c->input(1), 2, &maindiag));
  TF_RETURN_IF_ERROR(c->WithRankAtLeast(c->input(2), 2, &subdiag));
  TF_RETURN_IF_ERROR(c->WithRankAtLeast(c->input(3), 2, &rhs));

  // Extract batch dimensions and check they are the same.
  ShapeHandle superdiag_batch_shape;
  ShapeHandle maindiag_batch_shape;
  ShapeHandle subdiag_batch_shape;
  ShapeHandle rhs_batch_shape;
  TF_RETURN_IF_ERROR(c->Subshape(superdiag, 0, -2, &superdiag_batch_shape));
  TF_RETURN_IF_ERROR(c->Subshape(maindiag, 0, -2, &maindiag_batch_shape));
  TF_RETURN_IF_ERROR(c->Subshape(subdiag, 0, -2, &subdiag_batch_shape));
  TF_RETURN_IF_ERROR(c->Subshape(rhs, 0, -2, &rhs_batch_shape));
  TF_RETURN_IF_ERROR(c->Merge(superdiag, maindiag, &superdiag));
  TF_RETURN_IF_ERROR(
      c->Merge(maindiag_batch_shape, rhs_batch_shape, &rhs_batch_shape));
  TF_RETURN_IF_ERROR(
      c->Merge(subdiag_batch_shape, rhs_batch_shape, &rhs_batch_shape));

  // Check that diagonals have the same shape.
  TF_RETURN_IF_ERROR(c->Merge(superdiag, maindiag, &maindiag));
  TF_RETURN_IF_ERROR(c->Merge(subdiag, maindiag, &maindiag));

  // Check that size of tri-diagonal matrix is the same as height of matrix on
  // the right.
  DimensionHandle m_lhs = c->Dim(maindiag, -1);
  DimensionHandle m_rhs = c->Dim(rhs, -2);
  TF_RETURN_IF_ERROR(c->Merge(m_lhs, m_rhs, &m_lhs));

  // Check that next-to-last dimension of diagonals is 1.
  DimensionHandle unused;
  TF_RETURN_IF_ERROR(c->WithValue(c->Dim(maindiag, -2), 1, &unused));

  // The output shape is the same as rhs shape.
  c->set_output(0, rhs);
  return Status::OK();
}

// The first input is [...,3,M] and second input is [...,M,K].
// Output is [...,M,K].