def test_less_than_two_dims_raises_static(self):
    x = rng.rand(3)
    y = rng.rand(1, 1)

    with self.assertRaisesRegexp(ValueError, "at least two dimensions"):
      linear_operator_util.broadcast_matrix_batch_dims([x, y])

    with self.assertRaisesRegexp(ValueError, "at least two dimensions"):
      linear_operator_util.broadcast_matrix_batch_dims([y, x])

  def test_one_batch_matrix_returned_after_tensor_conversion(self):
    arr = rng.rand(2, 3, 4)
    tensor, = linear_operator_util.broadcast_matrix_batch_dims([arr])
    self.assertTrue(isinstance(tensor, ops.Tensor))

    with self.cached_session():
      self.assertAllClose(arr, tensor.eval())

  def _to_dense(self):
    num_cols = 0
    rows = []
    broadcasted_blocks = [operator.to_dense() for operator in self.operators]
    broadcasted_blocks = linear_operator_util.broadcast_matrix_batch_dims(
        broadcasted_blocks)
    for block in broadcasted_blocks:
      batch_row_shape = array_ops.shape(block)[:-1]

      zeros_to_pad_before_shape = array_ops.concat(
          [batch_row_shape, [num_cols]], axis=-1)
      zeros_to_pad_before = array_ops.zeros(
          shape=zeros_to_pad_before_shape, dtype=block.dtype)
      num_cols += array_ops.shape(block)[-1]
      zeros_to_pad_after_shape = array_ops.concat(
          [batch_row_shape,
           [self.domain_dimension_tensor() - num_cols]], axis=-1)
      zeros_to_pad_after = array_ops.zeros(
          shape=zeros_to_pad_after_shape, dtype=block.dtype)

      rows.append(array_ops.concat(
          [zeros_to_pad_before, block, zeros_to_pad_after], axis=-1))

    mat = array_ops.concat(rows, axis=-2)
    mat.set_shape(self.shape)
    return mat

  def _operator_and_matrix(self, build_info, dtype, use_placeholder):
    shape = list(build_info.shape)
    expected_factors = build_info.__dict__["factors"]
    matrices = [
        linear_operator_test_util.random_positive_definite_matrix(
            block_shape, dtype, force_well_conditioned=True)
        for block_shape in expected_factors
    ]

    lin_op_matrices = matrices

    if use_placeholder:
      lin_op_matrices = [
          array_ops.placeholder_with_default(m, shape=None) for m in matrices]

    operator = kronecker.LinearOperatorKronecker(
        [linalg.LinearOperatorFullMatrix(
            l, is_square=True) for l in lin_op_matrices])

    matrices = linear_operator_util.broadcast_matrix_batch_dims(matrices)

    kronecker_dense = _kronecker_dense(matrices)

    if not use_placeholder:
      kronecker_dense.set_shape(shape)

    return operator, kronecker_dense

 def _diag_part(self):
   diag_list = []
   for operator in self.operators:
     # Extend the axis for broadcasting.
     diag_list += [operator.diag_part()[..., array_ops.newaxis]]
   diag_list = linear_operator_util.broadcast_matrix_batch_dims(diag_list)
   diagonal = array_ops.concat(diag_list, axis=-2)
   return array_ops.squeeze(diagonal, axis=-1)

  def _matmul(self, x, adjoint=False, adjoint_arg=False):
    split_dim = -1 if adjoint_arg else -2
    # Split input by rows normally, and otherwise columns.
    split_x = self._split_input_into_blocks(x, axis=split_dim)

    result_list = []
    for index, operator in enumerate(self.operators):
      result_list += [operator.matmul(
          split_x[index], adjoint=adjoint, adjoint_arg=adjoint_arg)]
    result_list = linear_operator_util.broadcast_matrix_batch_dims(
        result_list)
    return array_ops.concat(result_list, axis=-2)

  def _solve(self, rhs, adjoint=False, adjoint_arg=False):
    split_dim = -1 if adjoint_arg else -2
    # Split input by rows normally, and otherwise columns.
    split_rhs = self._split_input_into_blocks(rhs, axis=split_dim)

    solution_list = []
    for index, operator in enumerate(self.operators):
      solution_list += [operator.solve(
          split_rhs[index], adjoint=adjoint, adjoint_arg=adjoint_arg)]

    solution_list = linear_operator_util.broadcast_matrix_batch_dims(
        solution_list)
    return array_ops.concat(solution_list, axis=-2)

  def test_static_dims_broadcast_second_arg_higher_rank(self):
    # x.batch_shape =    [1, 2]
    # y.batch_shape = [1, 3, 1]
    # broadcast batch shape = [1, 3, 2]
    x = rng.rand(1, 2, 1, 5)
    y = rng.rand(1, 3, 2, 3, 7)
    batch_of_zeros = np.zeros((1, 3, 2, 1, 1))
    x_bc_expected = x + batch_of_zeros
    y_bc_expected = y + batch_of_zeros

    x_bc, y_bc = linear_operator_util.broadcast_matrix_batch_dims([x, y])

    with self.cached_session() as sess:
      self.assertAllEqual(x_bc_expected.shape, x_bc.get_shape())
      self.assertAllEqual(y_bc_expected.shape, y_bc.get_shape())
      x_bc_, y_bc_ = sess.run([x_bc, y_bc])
      self.assertAllClose(x_bc_expected, x_bc_)
      self.assertAllClose(y_bc_expected, y_bc_)

  def _broadcast_batch_dims(self, x, spectrum):
    """Broadcast batch dims of batch matrix `x` and spectrum."""
    # spectrum.shape = batch_shape + block_shape
    # First make spectrum a batch matrix with
    #   spectrum.shape = batch_shape + [prod(block_shape), 1]
    spec_mat = array_ops.reshape(
        spectrum, array_ops.concat(
            (self.batch_shape_tensor(), [-1, 1]), axis=0))
    # Second, broadcast, possibly requiring an addition of array of zeros.
    x, spec_mat = linear_operator_util.broadcast_matrix_batch_dims((x,
                                                                    spec_mat))
    # Third, put the block shape back into spectrum.
    batch_shape = array_ops.shape(x)[:-2]
    spectrum = array_ops.reshape(
        spec_mat,
        array_ops.concat((batch_shape, self.block_shape_tensor()), axis=0))

    return x, spectrum

  def _operator_and_mat_and_feed_dict(self, build_info, dtype, use_placeholder):
    shape = list(build_info.shape)
    expected_blocks = (
        build_info.__dict__["blocks"] if "blocks" in build_info.__dict__
        else [shape])
    matrices = [
        linear_operator_test_util.random_positive_definite_matrix(
            block_shape, dtype, force_well_conditioned=True)
        for block_shape in expected_blocks
    ]

    if use_placeholder:
      matrices_ph = [
          array_ops.placeholder(dtype=dtype) for _ in expected_blocks
      ]
      # Evaluate here because (i) you cannot feed a tensor, and (ii)
      # values are random and we want the same value used for both mat and
      # feed_dict.
      matrices = self.evaluate(matrices)
      operator = block_diag.LinearOperatorBlockDiag(
          [linalg.LinearOperatorFullMatrix(
              m_ph, is_square=True) for m_ph in matrices_ph],
          is_square=True)
      feed_dict = {m_ph: m for (m_ph, m) in zip(matrices_ph, matrices)}
    else:
      operator = block_diag.LinearOperatorBlockDiag(
          [linalg.LinearOperatorFullMatrix(
              m, is_square=True) for m in matrices])
      feed_dict = None
      # Should be auto-set.
      self.assertTrue(operator.is_square)

    # Broadcast the shapes.
    expected_shape = list(build_info.shape)

    matrices = linear_operator_util.broadcast_matrix_batch_dims(matrices)

    block_diag_dense = _block_diag_dense(expected_shape, matrices)

    if not use_placeholder:
      block_diag_dense.set_shape(
          expected_shape[:-2] + [expected_shape[-1], expected_shape[-1]])

    return operator, block_diag_dense, feed_dict

  def operator_and_matrix(
      self, shape_info, dtype, use_placeholder,
      ensure_self_adjoint_and_pd=False):
    shape = list(shape_info.shape)
    expected_blocks = (
        shape_info.__dict__["blocks"] if "blocks" in shape_info.__dict__
        else [shape])
    matrices = [
        linear_operator_test_util.random_positive_definite_matrix(
            block_shape, dtype, force_well_conditioned=True)
        for block_shape in expected_blocks
    ]

    lin_op_matrices = matrices

    if use_placeholder:
      lin_op_matrices = [
          array_ops.placeholder_with_default(
              matrix, shape=None) for matrix in matrices]

    operator = block_diag.LinearOperatorBlockDiag(
        [linalg.LinearOperatorFullMatrix(
            l,
            is_square=True,
            is_self_adjoint=True if ensure_self_adjoint_and_pd else None,
            is_positive_definite=True if ensure_self_adjoint_and_pd else None)
         for l in lin_op_matrices])

    # Should be auto-set.
    self.assertTrue(operator.is_square)

    # Broadcast the shapes.
    expected_shape = list(shape_info.shape)

    matrices = linear_operator_util.broadcast_matrix_batch_dims(matrices)

    block_diag_dense = _block_diag_dense(expected_shape, matrices)

    if not use_placeholder:
      block_diag_dense.set_shape(
          expected_shape[:-2] + [expected_shape[-1], expected_shape[-1]])

    return operator, block_diag_dense

  def test_dynamic_dims_broadcast_32bit_second_arg_higher_rank(self):
    # x.batch_shape =    [1, 2]
    # y.batch_shape = [3, 4, 1]
    # broadcast batch shape = [3, 4, 2]
    x = rng.rand(1, 2, 1, 5).astype(np.float32)
    y = rng.rand(3, 4, 1, 3, 7).astype(np.float32)
    batch_of_zeros = np.zeros((3, 4, 2, 1, 1)).astype(np.float32)
    x_bc_expected = x + batch_of_zeros
    y_bc_expected = y + batch_of_zeros

    x_ph = array_ops.placeholder(dtypes.float32)
    y_ph = array_ops.placeholder(dtypes.float32)

    x_bc, y_bc = linear_operator_util.broadcast_matrix_batch_dims([x_ph, y_ph])

    with self.cached_session() as sess:
      x_bc_, y_bc_ = sess.run([x_bc, y_bc], feed_dict={x_ph: x, y_ph: y})
      self.assertAllClose(x_bc_expected, x_bc_)
      self.assertAllClose(y_bc_expected, y_bc_)

  def _operator_and_mat_and_feed_dict(self, build_info, dtype, use_placeholder):
    shape = list(build_info.shape)
    expected_blocks = (
        build_info.__dict__["blocks"] if "blocks" in build_info.__dict__
        else [shape])
    diag_matrices = [
        linear_operator_test_util.random_uniform(
            shape=block_shape[:-1], minval=1., maxval=20., dtype=dtype)
        for block_shape in expected_blocks
    ]

    if use_placeholder:
      diag_matrices_ph = [
          array_ops.placeholder(dtype=dtype) for _ in expected_blocks
      ]
      diag_matrices = self.evaluate(diag_matrices)
      # Evaluate here because (i) you cannot feed a tensor, and (ii)
      # values are random and we want the same value used for both mat and
      # feed_dict.
      operator = block_diag.LinearOperatorBlockDiag(
          [linalg.LinearOperatorDiag(m_ph) for m_ph in diag_matrices_ph])
      feed_dict = {m_ph: m for (m_ph, m) in zip(
          diag_matrices_ph, diag_matrices)}
    else:
      operator = block_diag.LinearOperatorBlockDiag(
          [linalg.LinearOperatorDiag(m) for m in diag_matrices])
      feed_dict = None
      # Should be auto-set.
      self.assertTrue(operator.is_square)

    # Broadcast the shapes.
    expected_shape = list(build_info.shape)

    matrices = linear_operator_util.broadcast_matrix_batch_dims(
        [array_ops.matrix_diag(diag_block) for diag_block in diag_matrices])

    block_diag_dense = _block_diag_dense(expected_shape, matrices)
    if not use_placeholder:
      block_diag_dense.set_shape(
          expected_shape[:-2] + [expected_shape[-1], expected_shape[-1]])

    return operator, block_diag_dense, feed_dict

  def _operator_and_matrix(
      self, build_info, dtype, use_placeholder,
      ensure_self_adjoint_and_pd=False):
    # Kronecker products constructed below will be from symmetric
    # positive-definite matrices.
    del ensure_self_adjoint_and_pd
    shape = list(build_info.shape)
    expected_factors = build_info.__dict__["factors"]
    matrices = [
        linear_operator_test_util.random_positive_definite_matrix(
            block_shape, dtype, force_well_conditioned=True)
        for block_shape in expected_factors
    ]

    lin_op_matrices = matrices

    if use_placeholder:
      lin_op_matrices = [
          array_ops.placeholder_with_default(m, shape=None) for m in matrices]

    operator = kronecker.LinearOperatorKronecker(
        [linalg.LinearOperatorFullMatrix(
            l,
            is_square=True,
            is_self_adjoint=True,
            is_positive_definite=True)
         for l in lin_op_matrices])

    matrices = linear_operator_util.broadcast_matrix_batch_dims(matrices)

    kronecker_dense = _kronecker_dense(matrices)

    if not use_placeholder:
      kronecker_dense.set_shape(shape)

    return operator, kronecker_dense

 def test_zero_batch_matrices_returned_as_empty_list(self):
   self.assertAllEqual([],
                       linear_operator_util.broadcast_matrix_batch_dims([]))