def benchmarkTridiagonalMulOp(self): devices = [('/cpu:0', 'cpu')] for device_id, device_name in devices: for batch_size, matrix_size in self.sizes: with ops.Graph().as_default(), \ session.Session(config=benchmark.benchmark_config()) as sess, \ ops.device(device_id): upper, diag, lower, vec = self._generateData( batch_size, matrix_size) x1 = self.baseline(upper, diag, lower, vec) x2 = linalg_impl.tridiagonal_matmul( (upper, diag, lower), vec) variables.global_variables_initializer().run() self.run_op_benchmark( sess, control_flow_ops.group(x1), min_iters=10, store_memory_usage=False, name=('tridiagonal_matmul_baseline_%s' '_batch_size_%d_matrix_size_%d' % (device_name, batch_size, matrix_size))) self.run_op_benchmark( sess, control_flow_ops.group(x2), min_iters=10, store_memory_usage=False, name= ('tridiagonal_matmul_%s_batch_size_%d_matrix_size_%d' % (device_name, batch_size, matrix_size)))
def benchmarkTridiagonalMulOp(self): devices = [('/cpu:0', 'cpu')] if test.is_gpu_available(cuda_only=True): devices += [('/gpu:0', 'gpu')] for device_option, size_option in itertools.product(devices, self.sizes): device_id, device_name = device_option m, batch_size, n = size_option with ops.Graph().as_default(), \ session.Session(config=benchmark.benchmark_config()) as sess, \ ops.device(device_id): upper, diag, lower, vec = self._generateData(batch_size, m, n) x1 = self.baseline(upper, diag, lower, vec) x2 = linalg_impl.tridiagonal_matmul((upper, diag, lower), vec, diagonals_format='sequence') variables.global_variables_initializer().run() self.run_op_benchmark( sess, control_flow_ops.group(x1), min_iters=10, store_memory_usage=False, name=('tridiagonal_matmul_baseline_%s' '_batch_size_%d_m_%d_n_%d' % (device_name, batch_size, m, n))) self.run_op_benchmark( sess, control_flow_ops.group(x2), min_iters=10, store_memory_usage=False, name=('tridiagonal_matmul_%s_batch_size_%d_m_%d_n_%d' % (device_name, batch_size, m, n)))
def _matmul(self, x, adjoint=False, adjoint_arg=False): diagonals = self.diagonals if adjoint: diagonals = self._construct_adjoint_diagonals(diagonals) x = linalg.adjoint(x) if adjoint_arg else x return linalg.tridiagonal_matmul( diagonals, x, diagonals_format=self.diagonals_format)
def _testAllFormats(self, superdiag, maindiag, subdiag, rhs, expected, dtype=dtypes.float64): superdiag_extended = np.pad(superdiag, [0, 1], 'constant') subdiag_extended = np.pad(subdiag, [1, 0], 'constant') diags_compact = np.stack([superdiag_extended, maindiag, subdiag_extended]) diags_matrix = np.diag(superdiag, 1) + np.diag(maindiag, 0) + np.diag( subdiag, -1) diags_sequence = (constant_op.constant(superdiag_extended, dtype), constant_op.constant(maindiag, dtype), constant_op.constant(subdiag_extended, dtype)) diags_compact = constant_op.constant(diags_compact, dtype) diags_matrix = constant_op.constant(diags_matrix, dtype) rhs = constant_op.constant(rhs, dtype) rhs_batch = array_ops.stack([rhs, 2 * rhs]) diags_compact_batch = array_ops.stack([diags_compact, 2 * diags_compact]) diags_matrix_batch = array_ops.stack([diags_matrix, 2 * diags_matrix]) diags_sequence_batch = [array_ops.stack([x, 2 * x]) for x in diags_sequence] results = [ linalg_impl.tridiagonal_matmul( diags_sequence, rhs, diagonals_format='sequence'), linalg_impl.tridiagonal_matmul( diags_compact, rhs, diagonals_format='compact'), linalg_impl.tridiagonal_matmul( diags_matrix, rhs, diagonals_format='matrix') ] results_batch = [ linalg_impl.tridiagonal_matmul( diags_sequence_batch, rhs_batch, diagonals_format='sequence'), linalg_impl.tridiagonal_matmul( diags_compact_batch, rhs_batch, diagonals_format='compact'), linalg_impl.tridiagonal_matmul( diags_matrix_batch, rhs_batch, diagonals_format='matrix') ] with self.cached_session(use_gpu=True): results = self.evaluate(results) results_batch = self.evaluate(results_batch) expected = np.array(expected) expected_batch = np.stack([expected, 4 * expected]) for result in results: self.assertAllClose(result, expected) for result in results_batch: self.assertAllClose(result, expected_batch)
def testBatch(self): b = 20 m = 10 n = 15 superdiag = self._randomComplexArray((b, m - 1)) maindiag = self._randomComplexArray((b, m)) subdiag = self._randomComplexArray((b, m - 1)) rhs = self._randomComplexArray((b, m, n)) matrix = np.stack([np.diag(superdiag[i], 1) + \ np.diag(maindiag[i], 0) + \ np.diag(subdiag[i], -1) for i in range(b)]) expected_result = np.matmul(matrix, rhs) result = linalg_impl.tridiagonal_matmul( constant_op.constant(matrix, dtype=dtypes.complex128), constant_op.constant(rhs, dtype=dtypes.complex128), diagonals_format='matrix') with self.cached_session(): result = self.evaluate(result) self.assertAllClose(result, expected_result)
def testBatch(self): b = 20 m = 10 n = 15 superdiag = self._randomComplexArray((b, m - 1)) maindiag = self._randomComplexArray((b, m)) subdiag = self._randomComplexArray((b, m - 1)) rhs = self._randomComplexArray((b, m, n)) matrix = np.stack([np.diag(superdiag[i], 1) + \ np.diag(maindiag[i], 0) + \ np.diag(subdiag[i], -1) for i in range(b)]) expected_result = np.matmul(matrix, rhs) result = linalg_impl.tridiagonal_matmul( constant_op.constant(matrix, dtype=dtypes.complex128), constant_op.constant(rhs, dtype=dtypes.complex128), diagonals_format='matrix') with self.cached_session(use_gpu=True): result = self.evaluate(result) self.assertAllClose(result, expected_result)
def benchmarkTridiagonalMulOp(self): devices = [('/cpu:0', 'cpu')] if test.is_gpu_available(cuda_only=True): devices += [('/gpu:0', 'gpu')] for device_option, size_option in itertools.product( devices, self.sizes): device_id, device_name = device_option m, batch_size, n = size_option with ops.Graph().as_default(), \ session.Session(config=benchmark.benchmark_config()) as sess, \ ops.device(device_id): upper, diag, lower, vec = self._generateData( batch_size, m, n) x1 = self.baseline(upper, diag, lower, vec) x2 = linalg_impl.tridiagonal_matmul( (upper, diag, lower), vec, diagonals_format='sequence') self.evaluate(variables.global_variables_initializer()) self.run_op_benchmark( sess, control_flow_ops.group(x1), min_iters=10, store_memory_usage=False, name=('tridiagonal_matmul_baseline_%s' '_batch_size_%d_m_%d_n_%d' % (device_name, batch_size, m, n))) self.run_op_benchmark( sess, control_flow_ops.group(x2), min_iters=10, store_memory_usage=False, name=('tridiagonal_matmul_%s_batch_size_%d_m_%d_n_%d' % (device_name, batch_size, m, n)))
def _testAllFormats(self, superdiag, maindiag, subdiag, rhs, expected, dtype=dtypes.float64): superdiag_extended = np.pad(superdiag, [0, 1], 'constant') subdiag_extended = np.pad(subdiag, [1, 0], 'constant') diags_compact = np.stack( [superdiag_extended, maindiag, subdiag_extended]) diags_matrix = np.diag(superdiag, 1) + np.diag(maindiag, 0) + np.diag( subdiag, -1) diags_sequence = (constant_op.constant(superdiag_extended, dtype), constant_op.constant(maindiag, dtype), constant_op.constant(subdiag_extended, dtype)) diags_compact = constant_op.constant(diags_compact, dtype) diags_matrix = constant_op.constant(diags_matrix, dtype) rhs = constant_op.constant(rhs, dtype) rhs_batch = array_ops.stack([rhs, 2 * rhs]) diags_compact_batch = array_ops.stack( [diags_compact, 2 * diags_compact]) diags_matrix_batch = array_ops.stack([diags_matrix, 2 * diags_matrix]) diags_sequence_batch = [ array_ops.stack([x, 2 * x]) for x in diags_sequence ] results = [ linalg_impl.tridiagonal_matmul(diags_sequence, rhs, diagonals_format='sequence'), linalg_impl.tridiagonal_matmul(diags_compact, rhs, diagonals_format='compact'), linalg_impl.tridiagonal_matmul(diags_matrix, rhs, diagonals_format='matrix') ] results_batch = [ linalg_impl.tridiagonal_matmul(diags_sequence_batch, rhs_batch, diagonals_format='sequence'), linalg_impl.tridiagonal_matmul(diags_compact_batch, rhs_batch, diagonals_format='compact'), linalg_impl.tridiagonal_matmul(diags_matrix_batch, rhs_batch, diagonals_format='matrix') ] with self.cached_session(use_gpu=False): results = self.evaluate(results) results_batch = self.evaluate(results_batch) expected = np.array(expected) expected_batch = np.stack([expected, 4 * expected]) for result in results: self.assertAllClose(result, expected) for result in results_batch: self.assertAllClose(result, expected_batch)
def _tridiagonal_matmul(self, diagonals, rhs, diagonals_format): return linalg_impl.tridiagonal_matmul( diagonals, rhs, diagonals_format=diagonals_format)