Python array_ops.broadcast_to函数代码示例

  def benchmarkBatchMatMulBroadcast(self):
    for (a_shape, b_shape) in self.shape_pairs:
      with compat.forward_compatibility_horizon(2019, 4, 26):
        with ops.Graph().as_default(), \
            session.Session(config=benchmark.benchmark_config()) as sess, \
            ops.device("/cpu:0"):
          matrix_a = variables.Variable(
              GetRandomNormalInput(a_shape, np.float32))
          matrix_b = variables.Variable(
              GetRandomNormalInput(b_shape, np.float32))
          variables.global_variables_initializer().run()

          # Use batch matmul op's internal broadcasting.
          self.run_op_benchmark(
              sess,
              math_ops.matmul(matrix_a, matrix_b),
              min_iters=50,
              name="batch_matmul_cpu_{}_{}".format(a_shape, b_shape))

          # Manually broadcast the input matrices using the broadcast_to op.
          broadcasted_batch_shape = array_ops.broadcast_static_shape(
              matrix_a.shape[:-2], matrix_b.shape[:-2])
          broadcasted_a_shape = broadcasted_batch_shape.concatenate(
              matrix_a.shape[-2:])
          broadcasted_b_shape = broadcasted_batch_shape.concatenate(
              matrix_b.shape[-2:])
          self.run_op_benchmark(
              sess,
              math_ops.matmul(
                  array_ops.broadcast_to(matrix_a, broadcasted_a_shape),
                  array_ops.broadcast_to(matrix_b, broadcasted_b_shape)),
              min_iters=50,
              name="batch_matmul_manual_broadcast_cpu_{}_{}".format(
                  a_shape, b_shape))

 def testBroadcastScalarToNonScalar(self):
   with self.session(use_gpu=True):
     x = np.array(1.0, dtype=np.float)
     v_tf = array_ops.broadcast_to(constant_op.constant(1.0), [2, 3, 4,
                                                               1, 1, 1])
     v_np = np.broadcast_to(x, [2, 3, 4, 1, 1, 1])
     self.assertAllEqual(v_tf.eval(), v_np)

 def testBroadcastToBasic(self):
   for dtype in [np.uint8, np.uint16, np.int8, np.int16, np.int32, np.int64]:
     with self.test_session(use_gpu=True):
       x = np.array([1, 2, 3], dtype=dtype)
       v_tf = array_ops.broadcast_to(constant_op.constant(x), [3, 3])
       v_np = np.broadcast_to(x, [3, 3])
       self.assertAllEqual(v_tf.eval(), v_np)

 def testBroadcastToBadOutputShape(self):
   with context.eager_mode():
     with self.assertRaisesRegexp(errors.InvalidArgumentError,
                                  "Unable to broadcast tensor of shape"):
       self.evaluate(
           array_ops.broadcast_to(
               constant_op.constant([0, 1]), constant_op.constant([2, 1])))

 def testBroadcastToShapeLargerDim2(self):
   input_shape = [2, 1, 3, 2, 2, 2, 1, 1, 1]
   output_shape = [1, 1, 1, 2, 5, 3, 2, 2, 2, 3, 3, 3]
   with self.cached_session(use_gpu=True):
     x = np.array(np.random.randint(5, size=input_shape), dtype=np.int32)
     v_tf = array_ops.broadcast_to(constant_op.constant(x), output_shape)
     v_np = np.broadcast_to(x, output_shape)
     self.assertAllEqual(v_tf.eval(), v_np)

def _broadcast_to_uniform_shape(rt_input, shape, broadcast_inner_dimensions):
  """Broadcasts rt_input to the uniform shape `shape`."""
  if isinstance(rt_input, ragged_tensor.RaggedTensor):
    raise ValueError('Incompatible with shape: ragged rank mismatch')
  if broadcast_inner_dimensions:
    return array_ops.broadcast_to(rt_input, shape.inner_dim_sizes)
  else:
    return rt_input

 def testGradientForScalar(self):
   x = constant_op.constant(1, dtype=dtypes.float32)
   v = array_ops.broadcast_to(x, [2, 4, 3])
   out = 2 * v
   with self.cached_session():
     err = gradient_checker.compute_gradient_error(x, x.get_shape(), out,
                                                   out.get_shape())
   self.assertLess(err, 1e-4)

 def testGradientWithBroadcastAllDimensions(self):
   x = constant_op.constant([[1, 2, 3], [4, 5, 6]], dtype=dtypes.float32)
   v = array_ops.broadcast_to(x, [5, 4, 6])
   out = 2 * v
   with self.test_session():
     err = gradient_checker.compute_gradient_error(x, x.get_shape(),
                                                   out, out.get_shape())
   self.assertLess(err, 1e-4)

  def _verifyLu(self, x, output_idx_type=dtypes.int64):
    # Verify that Px = LU.
    lu, perm = linalg_ops.lu(x, output_idx_type=output_idx_type)

    # Prepare the lower factor of shape num_rows x num_rows
    lu_shape = np.array(lu.shape.as_list())
    batch_shape = lu_shape[:-2]
    num_rows = lu_shape[-2]
    num_cols = lu_shape[-1]

    lower = array_ops.matrix_band_part(lu, -1, 0)

    if num_rows > num_cols:
      eye = linalg_ops.eye(
          num_rows, batch_shape=batch_shape, dtype=lower.dtype)
      lower = array_ops.concat([lower, eye[..., num_cols:]], axis=-1)
    elif num_rows < num_cols:
      lower = lower[..., :num_rows]

    # Fill the diagonal with ones.
    ones_diag = array_ops.ones(
        np.append(batch_shape, num_rows), dtype=lower.dtype)
    lower = array_ops.matrix_set_diag(lower, ones_diag)

    # Prepare the upper factor.
    upper = array_ops.matrix_band_part(lu, 0, -1)

    verification = math_ops.matmul(lower, upper)

    # Permute the rows of product of the Cholesky factors.
    if num_rows > 0:
      # Reshape the product of the triangular factors and permutation indices
      # to a single batch dimension. This makes it easy to apply
      # invert_permutation and gather_nd ops.
      perm_reshaped = array_ops.reshape(perm, [-1, num_rows])
      verification_reshaped = array_ops.reshape(verification,
                                                [-1, num_rows, num_cols])
      # Invert the permutation in each batch.
      inv_perm_reshaped = map_fn.map_fn(array_ops.invert_permutation,
                                        perm_reshaped)
      batch_size = perm_reshaped.shape.as_list()[0]
      # Prepare the batch indices with the same shape as the permutation.
      # The corresponding batch index is paired with each of the `num_rows`
      # permutation indices.
      batch_indices = math_ops.cast(
          array_ops.broadcast_to(
              math_ops.range(batch_size)[:, None], perm_reshaped.shape),
          dtype=output_idx_type)
      permuted_verification_reshaped = array_ops.gather_nd(
          verification_reshaped,
          array_ops.stack([batch_indices, inv_perm_reshaped], axis=-1))

      # Reshape the verification matrix back to the original shape.
      verification = array_ops.reshape(permuted_verification_reshaped,
                                       lu_shape)

    self._verifyLuBase(x, lower, upper, perm, verification,
                       output_idx_type)

 def testGradientWithIncreasingRank(self):
   x = constant_op.constant([[1], [2]],
                            dtype=dtypes.float32)
   v = array_ops.broadcast_to(x, [5, 2, 3])
   out = 2 * v
   with self.test_session():
     err = gradient_checker.compute_gradient_error(x, x.get_shape(),
                                                   out, out.get_shape())
   self.assertLess(err, 1e-4)

 def testGradientWithSameRank(self):
   x = constant_op.constant(np.reshape(np.arange(6), (2, 1, 3)),
                            dtype=dtypes.float32)
   v = array_ops.broadcast_to(x, [2, 5, 3])
   out = 2 * v
   with self.test_session():
     err = gradient_checker.compute_gradient_error(x, x.get_shape(),
                                                   out, out.get_shape())
   self.assertLess(err, 1e-4)

 def testBroadcastToShape(self):
   for input_dim in range(1, 6):
     for output_dim in range(input_dim, 6):
       with self.test_session(use_gpu=True):
         input_shape = [2] * input_dim
         output_shape = [2] * output_dim
         x = np.array(np.random.randint(5, size=input_shape), dtype=np.int32)
         v_tf = array_ops.broadcast_to(constant_op.constant(x), output_shape)
         v_np = np.broadcast_to(x, output_shape)
         self.assertAllEqual(v_tf.eval(), v_np)

 def testGradientForScalar(self):
   # TODO(alextp): There is a bug with broadcast_to on GPU from scalars,
   # hence we make this test cpu-only.
   with ops.device("cpu:0"):
     x = constant_op.constant(1, dtype=dtypes.float32)
     v = array_ops.broadcast_to(x, [2, 4, 3])
     out = 2 * v
     with self.test_session():
       err = gradient_checker.compute_gradient_error(x, x.get_shape(),
                                                     out, out.get_shape())
   self.assertLess(err, 1e-4)

 def testGradientWithLargeDim(self):
   input_shape = [2, 1, 3, 2, 2, 2, 1, 1, 1]
   output_shape = [1, 1, 1, 2, 5, 3, 2, 2, 2, 3, 3, 3]
   x = constant_op.constant(np.array(np.random.randn(*input_shape),
                                     dtype=np.float32))
   v = array_ops.broadcast_to(x, output_shape)
   out = 2 * v
   with self.cached_session():
     err = gradient_checker.compute_gradient_error(x, x.get_shape(),
                                                   out, out.get_shape())
   self.assertLess(err, 1e-4)

 def testBroadcastToShapeTypeAndInference(self):
   for dtype in [dtypes.int32, dtypes.int64]:
     with self.test_session(use_gpu=True):
       x = np.array([1, 2, 3])
       v_tf = array_ops.broadcast_to(
           constant_op.constant(x),
           constant_op.constant([3, 3], dtype=dtype))
       shape = v_tf.get_shape().as_list()
       v_np = np.broadcast_to(x, [3, 3])
       self.assertAllEqual(v_tf.eval(), v_np)
       # check shape inference when shape input is constant
       self.assertAllEqual(shape, v_np.shape)