code: # Copyright 2017 The TensorFlow Authors. All Rights Reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  # ==============================================================================
  r"""Multivariate autoregressive model (vector autoregression).
  
  Implements the following model (num_blocks = max(ar_order, ma_order + 1)):
  
    y(t, 1) = \sum_{i=1}^{ar_order} ar_coefs[i] * y(t - 1, i)
    y(t, i) = y(t - 1, i - 1) + ma_coefs[i - 1] * e(t) for 1 < i < num_blocks
    y(t, num_blocks) = y(t - 1, num_blocks - 1) + e(t)
  
  Where e(t) are Gaussian with zero mean and learned covariance.
  
  Each element of ar_coefs and ma_coefs is a [num_features x num_features]
  matrix. Each y(t, i) is a vector of length num_features. Indices in the above
  equations are one-based. Initial conditions y(0, i) come from prior state (which
  may either be learned or left as a constant with high prior covariance).
  
  If ar_order > ma_order, the observation model is:
    y(t, 1) + observation_noise(t)
  
  If ma_order >= ar_order, it is (to observe the moving average component):
    y(t, 1) + y(t, num_blocks) + observation_noise(t)
  
  Where observation_noise(t) are Gaussian with zero mean and learned covariance.
  
  This implementation uses a formulation which puts all of the autoregressive
  coefficients in the transition equation for the observed component, which
  enables learning using truncated backpropagation. Noise is not applied directly
  to the observed component (with the exception of standard observation noise),
  which further aids learning of the autoregressive coefficients when VARMA is in
  an ensemble with other models (in which case having an observation noise term is
  usually unavoidable).
  """
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  from tensorflow.contrib.timeseries.python.timeseries import math_utils
  from tensorflow.contrib.timeseries.python.timeseries.state_space_models import state_space_model
  
  from tensorflow.python.framework import ops
  from tensorflow.python.framework import tensor_shape
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import init_ops
  from tensorflow.python.ops import linalg_ops
  from tensorflow.python.ops import math_ops
  from tensorflow.python.ops import variable_scope
  
  
  class VARMA(state_space_model.StateSpaceModel):
    """A VARMA model implementation as a special case of the state space model."""
  
    def __init__(self,
                 autoregressive_order,
                 moving_average_order,
                 configuration=state_space_model.StateSpaceModelConfiguration()):
      """Construct a VARMA model.
  
      The size of the latent state for this model is:
        num_features * max(autoregressive_order, moving_average_order + 1)
      Square matrices of this size are constructed and multiplied.
  
      Args:
        autoregressive_order: The maximum autoregressive lag.
        moving_average_order: The maximum moving average lag, after which
          transient deviations are expected to return to their long-term mean.
        configuration: A StateSpaceModelConfiguration object.
      """
      self.ar_order = autoregressive_order
      self.ma_order = moving_average_order
      self.state_num_blocks = max(autoregressive_order, moving_average_order + 1)
      super(VARMA, self).__init__(configuration=configuration)
      self.state_dimension = self.state_num_blocks * self.num_features
  
    def _define_parameters(self, observation_transition_tradeoff_log=None):
      with variable_scope.variable_scope(self._variable_scope):
        # TODO(allenl): Evaluate parameter transformations for AR/MA coefficients
        # which improve interpretability/stability.
        self.ar_coefs = variable_scope.get_variable(
            name="ar_coefs",
            shape=[self.num_features, self.num_features, self.ar_order],
            dtype=self.dtype,
            initializer=init_ops.zeros_initializer())
        self.ma_coefs = variable_scope.get_variable(
            name="ma_coefs",
            initializer=array_ops.tile(
                linalg_ops.eye(self.num_features, dtype=self.dtype)[None, :, :],
                [self.ma_order, 1, 1]),
            dtype=self.dtype)
      super(VARMA, self)._define_parameters(
          observation_transition_tradeoff_log=observation_transition_tradeoff_log)
  
    def get_state_transition(self):
      """Construct state transition matrix from VARMA parameters.
  
      Returns:
        the state transition matrix. It has shape
          [self.state_dimension, self.state_dimension].
      """
      # Pad any unused AR blocks with zeros. The extra state is necessary if
      # ma_order >= ar_order.
      ar_coefs_padded = array_ops.reshape(
          array_ops.pad(self.ar_coefs,
                        [[0, 0], [0, 0],
                         [0, self.state_num_blocks - self.ar_order]]),
          [self.num_features, self.state_dimension])
      shift_matrix = array_ops.pad(
          linalg_ops.eye(
              (self.state_num_blocks - 1) * self.num_features, dtype=self.dtype),
          [[0, 0], [0, self.num_features]])
      return array_ops.concat([ar_coefs_padded, shift_matrix], axis=0)
  
    def get_noise_transform(self):
      """Construct state noise transform matrix from VARMA parameters.
  
      Returns:
        the state noise transform matrix. It has shape
          [self.state_dimension, self.num_features].
      """
      # Noise is broadcast, through the moving average coefficients, to
      # un-observed parts of the latent state.
      ma_coefs_padded = array_ops.reshape(
          array_ops.pad(self.ma_coefs,
                        [[self.state_num_blocks - 1 - self.ma_order, 0], [0, 0],
                         [0, 0]]),
          [(self.state_num_blocks - 1) * self.num_features, self.num_features],
          name="noise_transform")
      # Deterministically apply noise to the oldest component.
      return array_ops.concat(
          [ma_coefs_padded,
           linalg_ops.eye(self.num_features, dtype=self.dtype)],
          axis=0)
  
    def get_observation_model(self, times):
      """Construct observation model matrix from VARMA parameters.
  
      Args:
        times: A [batch size] vector indicating the times observation models are
            requested for. Unused.
      Returns:
        the observation model matrix. It has shape
          [self.num_features, self.state_dimension].
      """
      del times  # StateSpaceModel will broadcast along the batch dimension
      if self.ar_order > self.ma_order or self.state_num_blocks < 2:
        return array_ops.pad(
            linalg_ops.eye(self.num_features, dtype=self.dtype),
            [[0, 0], [0, self.num_features * (self.state_num_blocks - 1)]],
            name="observation_model")
      else:
        # Add a second observed component which "catches" the accumulated moving
        # average errors as they reach the end of the state. If ar_order >
        # ma_order, this is unnecessary, since accumulated errors cycle naturally.
        return array_ops.concat(
            [
                array_ops.pad(
                    linalg_ops.eye(self.num_features, dtype=self.dtype),
                    [[0, 0], [0,
                              self.num_features * (self.state_num_blocks - 2)]]),
                linalg_ops.eye(self.num_features, dtype=self.dtype)
            ],
            axis=1,
            name="observation_model")
  
    def get_state_transition_noise_covariance(
        self, minimum_initial_variance=1e-5):
      # Most state space models use only an explicit observation noise term to
      # model deviations from expectations, and so a low initial transition noise
      # parameter is helpful there. Since deviations from expectations are also
      # modeled as transition noise in VARMA, we set its initial value based on a
      # slight over-estimate empirical observation noise.
      if self._input_statistics is not None:
        feature_variance = self._scale_variance(
            self._input_statistics.series_start_moments.variance)
        initial_transition_noise_scale = math_ops.log(
            math_ops.maximum(
                math_ops.reduce_mean(feature_variance), minimum_initial_variance))
      else:
        initial_transition_noise_scale = 0.
      state_noise_transform = ops.convert_to_tensor(
          self.get_noise_transform(), dtype=self.dtype)
      state_noise_dimension = tensor_shape.dimension_value(
          state_noise_transform.shape[1])
      return math_utils.variable_covariance_matrix(
          state_noise_dimension, "state_transition_noise",
          dtype=self.dtype,
          initial_overall_scale_log=initial_transition_noise_scale)
   # Copyright 2015 The TensorFlow Authors. All Rights Reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  # ==============================================================================
  """Tests for fft operations."""
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import numpy as np
  from six.moves import xrange  # pylint: disable=redefined-builtin
  
  from tensorflow.core.protobuf import config_pb2
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import errors
  from tensorflow.python.framework import ops
  from tensorflow.python.framework import test_util
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import gen_spectral_ops
  from tensorflow.python.ops import gradient_checker
  from tensorflow.python.ops import math_ops
  from tensorflow.python.ops import spectral_ops_test_util
  from tensorflow.python.ops.signal import fft_ops
  from tensorflow.python.platform import test
  
  VALID_FFT_RANKS = (1, 2, 3)
  
  
  class BaseFFTOpsTest(test.TestCase):
  
    def _compare(self, x, rank, fft_length=None, use_placeholder=False,
                 rtol=1e-4, atol=1e-4):
      self._compareForward(x, rank, fft_length, use_placeholder, rtol, atol)
      self._compareBackward(x, rank, fft_length, use_placeholder, rtol, atol)
  
    def _compareForward(self, x, rank, fft_length=None, use_placeholder=False,
                        rtol=1e-4, atol=1e-4):
      x_np = self._npFFT(x, rank, fft_length)
      if use_placeholder:
        x_ph = array_ops.placeholder(dtype=dtypes.as_dtype(x.dtype))
        x_tf = self._tfFFT(x_ph, rank, fft_length, feed_dict={x_ph: x})
      else:
        x_tf = self._tfFFT(x, rank, fft_length)
  
      self.assertAllClose(x_np, x_tf, rtol=rtol, atol=atol)
  
    def _compareBackward(self, x, rank, fft_length=None, use_placeholder=False,
                         rtol=1e-4, atol=1e-4):
      x_np = self._npIFFT(x, rank, fft_length)
      if use_placeholder:
        x_ph = array_ops.placeholder(dtype=dtypes.as_dtype(x.dtype))
        x_tf = self._tfIFFT(x_ph, rank, fft_length, feed_dict={x_ph: x})
      else:
        x_tf = self._tfIFFT(x, rank, fft_length)
  
      self.assertAllClose(x_np, x_tf, rtol=rtol, atol=atol)
  
    def _checkMemoryFail(self, x, rank):
      config = config_pb2.ConfigProto()
      config.gpu_options.per_process_gpu_memory_fraction = 1e-2
      with self.cached_session(config=config, force_gpu=True):
        self._tfFFT(x, rank, fft_length=None)
  
    def _checkGradComplex(self, func, x, y, result_is_complex=True,
                          rtol=1e-2, atol=1e-2):
      with self.cached_session(use_gpu=True):
        inx = ops.convert_to_tensor(x)
        iny = ops.convert_to_tensor(y)
        # func is a forward or inverse, real or complex, batched or unbatched FFT
        # function with a complex input.
        z = func(math_ops.complex(inx, iny))
        # loss = sum(|z|^2)
        loss = math_ops.reduce_sum(math_ops.real(z * math_ops.conj(z)))
  
        ((x_jacob_t, x_jacob_n),
         (y_jacob_t, y_jacob_n)) = gradient_checker.compute_gradient(
             [inx, iny], [list(x.shape), list(y.shape)],
             loss, [1],
             x_init_value=[x, y],
             delta=1e-2)
  
      self.assertAllClose(x_jacob_t, x_jacob_n, rtol=rtol, atol=atol)
      self.assertAllClose(y_jacob_t, y_jacob_n, rtol=rtol, atol=atol)
  
    def _checkGradReal(self, func, x, rtol=1e-2, atol=1e-2):
      with self.cached_session(use_gpu=True):
        inx = ops.convert_to_tensor(x)
        # func is a forward RFFT function (batched or unbatched).
        z = func(inx)
        # loss = sum(|z|^2)
        loss = math_ops.reduce_sum(math_ops.real(z * math_ops.conj(z)))
        x_jacob_t, x_jacob_n = test.compute_gradient(
            inx, list(x.shape), loss, [1], x_init_value=x, delta=1e-2)
  
      self.assertAllClose(x_jacob_t, x_jacob_n, rtol=rtol, atol=atol)
  
  
  class FFTOpsTest(BaseFFTOpsTest):
  
    def _tfFFT(self, x, rank, fft_length=None, feed_dict=None):
      # fft_length unused for complex FFTs.
      with self.cached_session(use_gpu=True) as sess:
        return sess.run(self._tfFFTForRank(rank)(x), feed_dict=feed_dict)
  
    def _tfIFFT(self, x, rank, fft_length=None, feed_dict=None):
      # fft_length unused for complex FFTs.
      with self.cached_session(use_gpu=True) as sess:
        return sess.run(self._tfIFFTForRank(rank)(x), feed_dict=feed_dict)
  
    def _npFFT(self, x, rank, fft_length=None):
      if rank == 1:
        return np.fft.fft2(x, s=fft_length, axes=(-1,))
      elif rank == 2:
        return np.fft.fft2(x, s=fft_length, axes=(-2, -1))
      elif rank == 3:
        return np.fft.fft2(x, s=fft_length, axes=(-3, -2, -1))
      else:
        raise ValueError("invalid rank")
  
    def _npIFFT(self, x, rank, fft_length=None):
      if rank == 1:
        return np.fft.ifft2(x, s=fft_length, axes=(-1,))
      elif rank == 2:
        return np.fft.ifft2(x, s=fft_length, axes=(-2, -1))
      elif rank == 3:
        return np.fft.ifft2(x, s=fft_length, axes=(-3, -2, -1))
      else:
        raise ValueError("invalid rank")
  
    def _tfFFTForRank(self, rank):
      if rank == 1:
        return fft_ops.fft
      elif rank == 2:
        return fft_ops.fft2d
      elif rank == 3:
        return fft_ops.fft3d
      else:
        raise ValueError("invalid rank")
  
    def _tfIFFTForRank(self, rank):
      if rank == 1:
        return fft_ops.ifft
      elif rank == 2:
        return fft_ops.ifft2d
      elif rank == 3:
        return fft_ops.ifft3d
      else:
        raise ValueError("invalid rank")
  
    @test_util.run_deprecated_v1
    def testEmpty(self):
      with spectral_ops_test_util.fft_kernel_label_map():
        for np_type in (np.complex64, np.complex128):
          for rank in VALID_FFT_RANKS:
            for dims in xrange(rank, rank + 3):
              x = np.zeros((0,) * dims).astype(np_type)
              self.assertEqual(x.shape, self._tfFFT(x, rank).shape)
              self.assertEqual(x.shape, self._tfIFFT(x, rank).shape)
  
    @test_util.run_deprecated_v1
    def testBasic(self):
      with spectral_ops_test_util.fft_kernel_label_map():
        for np_type, tol in ((np.complex64, 1e-4), (np.complex128, 1e-8)):
          for rank in VALID_FFT_RANKS:
            for dims in xrange(rank, rank + 3):
              self._compare(
                  np.mod(np.arange(np.power(4, dims)), 10).reshape(
                      (4,) * dims).astype(np_type), rank, rtol=tol, atol=tol)
  
    def testLargeBatch(self):
      if test.is_gpu_available(cuda_only=True):
        rank = 1
        for dims in xrange(rank, rank + 3):
          for np_type, tol in ((np.complex64, 1e-4), (np.complex128, 1e-5)):
            self._compare(
                np.mod(np.arange(np.power(128, dims)), 10).reshape(
                    (128,) * dims).astype(np_type), rank, rtol=tol, atol=tol)
  
    # TODO(yangzihao): Disable before we can figure out a way to
    # properly test memory fail for large batch fft.
    # def testLargeBatchMemoryFail(self):
    #   if test.is_gpu_available(cuda_only=True):
    #     rank = 1
    #     for dims in xrange(rank, rank + 3):
    #       self._checkMemoryFail(
    #           np.mod(np.arange(np.power(128, dims)), 64).reshape(
    #               (128,) * dims).astype(np.complex64), rank)
  
    @test_util.run_deprecated_v1
    def testBasicPlaceholder(self):
      with spectral_ops_test_util.fft_kernel_label_map():
        for np_type, tol in ((np.complex64, 1e-4), (np.complex128, 1e-8)):
          for rank in VALID_FFT_RANKS:
            for dims in xrange(rank, rank + 3):
              self._compare(
                  np.mod(np.arange(np.power(4, dims)), 10).reshape(
                      (4,) * dims).astype(np_type),
                  rank, use_placeholder=True, rtol=tol, atol=tol)
  
    @test_util.run_deprecated_v1
    def testRandom(self):
      with spectral_ops_test_util.fft_kernel_label_map():
        for np_type, tol in ((np.complex64, 1e-4), (np.complex128, 5e-6)):
          def gen(shape):
            n = np.prod(shape)
            re = np.random.uniform(size=n)
            im = np.random.uniform(size=n)
            return (re + im * 1j).reshape(shape)
  
          for rank in VALID_FFT_RANKS:
            for dims in xrange(rank, rank + 3):
              self._compare(gen((4,) * dims).astype(np_type), rank,
                            rtol=tol, atol=tol)
  
    @test_util.run_deprecated_v1
    def testRandom1D(self):
      with spectral_ops_test_util.fft_kernel_label_map():
        for np_type in (np.complex64, np.complex128):
          has_gpu = test.is_gpu_available(cuda_only=True)
          tol = {(np.complex64, True): 1e-4,
                 (np.complex64, False): 1e-2,
                 (np.complex128, True): 1e-4,
                 (np.complex128, False): 1e-2}[(np_type, has_gpu)]
          def gen(shape):
            n = np.prod(shape)
            re = np.random.uniform(size=n)
            im = np.random.uniform(size=n)
            return (re + im * 1j).reshape(shape)
  
          # Check a variety of power-of-2 FFT sizes.
          for dim in (128, 256, 512, 1024):
            self._compare(gen((dim,)).astype(np_type), 1, rtol=tol, atol=tol)
  
          # Check a variety of non-power-of-2 FFT sizes.
          for dim in (127, 255, 511, 1023):
            self._compare(gen((dim,)).astype(np_type), 1, rtol=tol, atol=tol)
  
    @test_util.run_deprecated_v1
    def testError(self):
      for rank in VALID_FFT_RANKS:
        for dims in xrange(0, rank):
          x = np.zeros((1,) * dims).astype(np.complex64)
          with self.assertRaisesWithPredicateMatch(
              ValueError, "Shape must be .*rank {}.*".format(rank)):
            self._tfFFT(x, rank)
          with self.assertRaisesWithPredicateMatch(
              ValueError, "Shape must be .*rank {}.*".format(rank)):
            self._tfIFFT(x, rank)
  
    @test_util.run_deprecated_v1
    def testGrad_Simple(self):
      with spectral_ops_test_util.fft_kernel_label_map():
        for np_type, tol in ((np.float32, 1e-4), (np.float64, 1e-10)):
          for rank in VALID_FFT_RANKS:
            for dims in xrange(rank, rank + 2):
              re = np.ones(shape=(4,) * dims, dtype=np_type) / 10.0
              im = np.zeros(shape=(4,) * dims, dtype=np_type)
              self._checkGradComplex(self._tfFFTForRank(rank), re, im,
                                     rtol=tol, atol=tol)
              self._checkGradComplex(self._tfIFFTForRank(rank), re, im,
                                     rtol=tol, atol=tol)
  
    @test_util.run_deprecated_v1
    def testGrad_Random(self):
      with spectral_ops_test_util.fft_kernel_label_map():
        for np_type, tol in ((np.float32, 1e-2), (np.float64, 1e-10)):
          for rank in VALID_FFT_RANKS:
            for dims in xrange(rank, rank + 2):
              re = np.random.rand(*((3,) * dims)).astype(np_type) * 2 - 1
              im = np.random.rand(*((3,) * dims)).astype(np_type) * 2 - 1
              self._checkGradComplex(self._tfFFTForRank(rank), re, im,
                                     rtol=tol, atol=tol)
              self._checkGradComplex(self._tfIFFTForRank(rank), re, im,
                                     rtol=tol, atol=tol)
  
  
  class RFFTOpsTest(BaseFFTOpsTest):
  
    def _compareBackward(self, x, rank, fft_length=None, use_placeholder=False):
      super(RFFTOpsTest, self)._compareBackward(x, rank, fft_length,
                                                use_placeholder)
  
    def _tfFFT(self, x, rank, fft_length=None, feed_dict=None):
      with self.cached_session(use_gpu=True) as sess:
        return sess.run(
            self._tfFFTForRank(rank)(x, fft_length), feed_dict=feed_dict)
  
    def _tfIFFT(self, x, rank, fft_length=None, feed_dict=None):
      with self.cached_session(use_gpu=True) as sess:
        return sess.run(
            self._tfIFFTForRank(rank)(x, fft_length), feed_dict=feed_dict)
  
    def _npFFT(self, x, rank, fft_length=None):
      if rank == 1:
        return np.fft.rfft2(x, s=fft_length, axes=(-1,))
      elif rank == 2:
        return np.fft.rfft2(x, s=fft_length, axes=(-2, -1))
      elif rank == 3:
        return np.fft.rfft2(x, s=fft_length, axes=(-3, -2, -1))
      else:
        raise ValueError("invalid rank")
  
    def _npIFFT(self, x, rank, fft_length=None):
      if rank == 1:
        return np.fft.irfft2(x, s=fft_length, axes=(-1,))
      elif rank == 2:
        return np.fft.irfft2(x, s=fft_length, axes=(-2, -1))
      elif rank == 3:
        return np.fft.irfft2(x, s=fft_length, axes=(-3, -2, -1))
      else:
        raise ValueError("invalid rank")
  
    def _tfFFTForRank(self, rank):
      if rank == 1:
        return fft_ops.rfft
      elif rank == 2:
        return fft_ops.rfft2d
      elif rank == 3:
        return fft_ops.rfft3d
      else:
        raise ValueError("invalid rank")
  
    def _tfIFFTForRank(self, rank):
      if rank == 1:
        return fft_ops.irfft
      elif rank == 2:
        return fft_ops.irfft2d
      elif rank == 3:
        return fft_ops.irfft3d
      else:
        raise ValueError("invalid rank")
  
    @test_util.run_deprecated_v1
    def testEmpty(self):
      with spectral_ops_test_util.fft_kernel_label_map():
        for rank in VALID_FFT_RANKS:
          for dims in xrange(rank, rank + 3):
            x = np.zeros((0,) * dims).astype(np.float32)
            self.assertEqual(x.shape, self._tfFFT(x, rank).shape)
            x = np.zeros((0,) * dims).astype(np.complex64)
            self.assertEqual(x.shape, self._tfIFFT(x, rank).shape)
  
    @test_util.run_deprecated_v1
    def testBasic(self):
      with spectral_ops_test_util.fft_kernel_label_map():
        for rank in VALID_FFT_RANKS:
          for dims in xrange(rank, rank + 3):
            for size in (5, 6):
              inner_dim = size // 2 + 1
              r2c = np.mod(np.arange(np.power(size, dims)), 10).reshape(
                  (size,) * dims)
              self._compareForward(r2c.astype(np.float32), rank, (size,) * rank)
              c2r = np.mod(np.arange(np.power(size, dims - 1) * inner_dim),
                           10).reshape((size,) * (dims - 1) + (inner_dim,))
              self._compareBackward(
                  c2r.astype(np.complex64), rank, (size,) * rank)
  
    def testLargeBatch(self):
      if test.is_gpu_available(cuda_only=True):
        rank = 1
        for dims in xrange(rank, rank + 3):
          for size in (64, 128):
            inner_dim = size // 2 + 1
            r2c = np.mod(np.arange(np.power(size, dims)), 10).reshape(
                (size,) * dims)
            self._compareForward(r2c.astype(np.float32), rank, (size,) * rank)
            c2r = np.mod(np.arange(np.power(size, dims - 1) * inner_dim),
                         10).reshape((size,) * (dims - 1) + (inner_dim,))
            self._compareBackward(c2r.astype(np.complex64), rank, (size,) * rank)
  
    @test_util.run_deprecated_v1
    def testBasicPlaceholder(self):
      with spectral_ops_test_util.fft_kernel_label_map():
        for rank in VALID_FFT_RANKS:
          for dims in xrange(rank, rank + 3):
            for size in (5, 6):
              inner_dim = size // 2 + 1
              r2c = np.mod(np.arange(np.power(size, dims)), 10).reshape(
                  (size,) * dims)
              self._compareForward(
                  r2c.astype(np.float32),
                  rank, (size,) * rank,
                  use_placeholder=True)
              c2r = np.mod(np.arange(np.power(size, dims - 1) * inner_dim),
                           10).reshape((size,) * (dims - 1) + (inner_dim,))
              self._compareBackward(
                  c2r.astype(np.complex64),
                  rank, (size,) * rank,
                  use_placeholder=True)
  
    @test_util.run_deprecated_v1
    def testFftLength(self):
      if test.is_gpu_available(cuda_only=True):
        with spectral_ops_test_util.fft_kernel_label_map():
          for rank in VALID_FFT_RANKS:
            for dims in xrange(rank, rank + 3):
              for size in (5, 6):
                inner_dim = size // 2 + 1
                r2c = np.mod(np.arange(np.power(size, dims)), 10).reshape(
                    (size,) * dims)
                c2r = np.mod(np.arange(np.power(size, dims - 1) * inner_dim),
                             10).reshape((size,) * (dims - 1) + (inner_dim,))
                # Test truncation (FFT size < dimensions).
                fft_length = (size - 2,) * rank
                self._compareForward(r2c.astype(np.float32), rank, fft_length)
                self._compareBackward(c2r.astype(np.complex64), rank, fft_length)
                # Confirm it works with unknown shapes as well.
                self._compareForward(
                    r2c.astype(np.float32),
                    rank,
                    fft_length,
                    use_placeholder=True)
                self._compareBackward(
                    c2r.astype(np.complex64),
                    rank,
                    fft_length,
                    use_placeholder=True)
                # Test padding (FFT size > dimensions).
                fft_length = (size + 2,) * rank
                self._compareForward(r2c.astype(np.float32), rank, fft_length)
                self._compareBackward(c2r.astype(np.complex64), rank, fft_length)
                # Confirm it works with unknown shapes as well.
                self._compareForward(
                    r2c.astype(np.float32),
                    rank,
                    fft_length,
                    use_placeholder=True)
                self._compareBackward(
                    c2r.astype(np.complex64),
                    rank,
                    fft_length,
                    use_placeholder=True)
  
    @test_util.run_deprecated_v1
    def testRandom(self):
      with spectral_ops_test_util.fft_kernel_label_map():
        def gen_real(shape):
          n = np.prod(shape)
          re = np.random.uniform(size=n)
          ret = re.reshape(shape)
          return ret
  
        def gen_complex(shape):
          n = np.prod(shape)
          re = np.random.uniform(size=n)
          im = np.random.uniform(size=n)
          ret = (re + im * 1j).reshape(shape)
          return ret
  
        for rank in VALID_FFT_RANKS:
          for dims in xrange(rank, rank + 3):
            for size in (5, 6):
              inner_dim = size // 2 + 1
              self._compareForward(gen_real((size,) * dims), rank, (size,) * rank)
              complex_dims = (size,) * (dims - 1) + (inner_dim,)
              self._compareBackward(
                  gen_complex(complex_dims), rank, (size,) * rank)
  
    @test_util.run_deprecated_v1
    def testError(self):
      with spectral_ops_test_util.fft_kernel_label_map():
        for rank in VALID_FFT_RANKS:
          for dims in xrange(0, rank):
            x = np.zeros((1,) * dims).astype(np.complex64)
            with self.assertRaisesWithPredicateMatch(
                ValueError, "Shape .* must have rank at least {}".format(rank)):
              self._tfFFT(x, rank)
            with self.assertRaisesWithPredicateMatch(
                ValueError, "Shape .* must have rank at least {}".format(rank)):
              self._tfIFFT(x, rank)
          for dims in xrange(rank, rank + 2):
            x = np.zeros((1,) * rank)
  
            # Test non-rank-1 fft_length produces an error.
            fft_length = np.zeros((1, 1)).astype(np.int32)
            with self.assertRaisesWithPredicateMatch(ValueError,
                                                     "Shape .* must have rank 1"):
              self._tfFFT(x, rank, fft_length)
            with self.assertRaisesWithPredicateMatch(ValueError,
                                                     "Shape .* must have rank 1"):
              self._tfIFFT(x, rank, fft_length)
  
            # Test wrong fft_length length.
            fft_length = np.zeros((rank + 1,)).astype(np.int32)
            with self.assertRaisesWithPredicateMatch(
                ValueError, "Dimension must be .*but is {}.*".format(rank + 1)):
              self._tfFFT(x, rank, fft_length)
            with self.assertRaisesWithPredicateMatch(
                ValueError, "Dimension must be .*but is {}.*".format(rank + 1)):
              self._tfIFFT(x, rank, fft_length)
  
          # Test that calling the kernel directly without padding to fft_length
          # produces an error.
          rffts_for_rank = {
              1: [gen_spectral_ops.rfft, gen_spectral_ops.irfft],
              2: [gen_spectral_ops.rfft2d, gen_spectral_ops.irfft2d],
              3: [gen_spectral_ops.rfft3d, gen_spectral_ops.irfft3d]
          }
          rfft_fn, irfft_fn = rffts_for_rank[rank]
          with self.assertRaisesWithPredicateMatch(
              errors.InvalidArgumentError,
              "Input dimension .* must have length of at least 6 but got: 5"):
            x = np.zeros((5,) * rank).astype(np.float32)
            fft_length = [6] * rank
            with self.cached_session():
              self.evaluate(rfft_fn(x, fft_length))
  
          with self.assertRaisesWithPredicateMatch(
              errors.InvalidArgumentError,
              "Input dimension .* must have length of at least .* but got: 3"):
            x = np.zeros((3,) * rank).astype(np.complex64)
            fft_length = [6] * rank
            with self.cached_session():
              self.evaluate(irfft_fn(x, fft_length))
  
    @test_util.run_deprecated_v1
    def testGrad_Simple(self):
      with spectral_ops_test_util.fft_kernel_label_map():
        for rank in VALID_FFT_RANKS:
          # rfft3d/irfft3d do not have gradients yet.
          if rank == 3:
            continue
          for dims in xrange(rank, rank + 2):
            for size in (5, 6):
              re = np.ones(shape=(size,) * dims, dtype=np.float32)
              im = -np.ones(shape=(size,) * dims, dtype=np.float32)
              self._checkGradReal(self._tfFFTForRank(rank), re)
              self._checkGradComplex(
                  self._tfIFFTForRank(rank), re, im, result_is_complex=False)
  
    @test_util.run_deprecated_v1
    def testGrad_Random(self):
      with spectral_ops_test_util.fft_kernel_label_map():
        for rank in VALID_FFT_RANKS:
          # rfft3d/irfft3d do not have gradients yet.
          if rank == 3:
            continue
          for dims in xrange(rank, rank + 2):
            for size in (5, 6):
              re = np.random.rand(*((size,) * dims)).astype(np.float32) * 2 - 1
              im = np.random.rand(*((size,) * dims)).astype(np.float32) * 2 - 1
              self._checkGradReal(self._tfFFTForRank(rank), re)
              self._checkGradComplex(
                  self._tfIFFTForRank(rank), re, im, result_is_complex=False)
  
  
  class FFTShiftTest(test.TestCase):
  
    @test_util.run_deprecated_v1
    def testDefinition(self):
      with self.session():
        x = [0, 1, 2, 3, 4, -4, -3, -2, -1]
        y = [-4, -3, -2, -1, 0, 1, 2, 3, 4]
        self.assertAllEqual(fft_ops.fftshift(x).eval(), y)
        self.assertAllEqual(fft_ops.ifftshift(y).eval(), x)
        x = [0, 1, 2, 3, 4, -5, -4, -3, -2, -1]
        y = [-5, -4, -3, -2, -1, 0, 1, 2, 3, 4]
        self.assertAllEqual(fft_ops.fftshift(x).eval(), y)
        self.assertAllEqual(fft_ops.ifftshift(y).eval(), x)
  
    @test_util.run_deprecated_v1
    def testAxesKeyword(self):
      with self.session():
        freqs = [[0, 1, 2], [3, 4, -4], [-3, -2, -1]]
        shifted = [[-1, -3, -2], [2, 0, 1], [-4, 3, 4]]
        self.assertAllEqual(fft_ops.fftshift(freqs, axes=(0, 1)).eval(), shifted)
        self.assertAllEqual(
            fft_ops.fftshift(freqs, axes=0).eval(),
            fft_ops.fftshift(freqs, axes=(0,)).eval())
        self.assertAllEqual(fft_ops.ifftshift(shifted, axes=(0, 1)).eval(), freqs)
        self.assertAllEqual(
            fft_ops.ifftshift(shifted, axes=0).eval(),
            fft_ops.ifftshift(shifted, axes=(0,)).eval())
        self.assertAllEqual(fft_ops.fftshift(freqs).eval(), shifted)
        self.assertAllEqual(fft_ops.ifftshift(shifted).eval(), freqs)
  
    @test_util.run_deprecated_v1
    def testNumpyCompatibility(self):
      with self.session():
        x = [0, 1, 2, 3, 4, -4, -3, -2, -1]
        y = [-4, -3, -2, -1, 0, 1, 2, 3, 4]
        self.assertAllEqual(fft_ops.fftshift(x).eval(), np.fft.fftshift(x))
        self.assertAllEqual(fft_ops.ifftshift(y).eval(), np.fft.ifftshift(y))
        x = [0, 1, 2, 3, 4, -5, -4, -3, -2, -1]
        y = [-5, -4, -3, -2, -1, 0, 1, 2, 3, 4]
        self.assertAllEqual(fft_ops.fftshift(x).eval(), np.fft.fftshift(x))
        self.assertAllEqual(fft_ops.ifftshift(y).eval(), np.fft.ifftshift(y))
        freqs = [[0, 1, 2], [3, 4, -4], [-3, -2, -1]]
        shifted = [[-1, -3, -2], [2, 0, 1], [-4, 3, 4]]
        self.assertAllEqual(
            fft_ops.fftshift(freqs, axes=(0, 1)).eval(),
            np.fft.fftshift(freqs, axes=(0, 1)))
        self.assertAllEqual(
            fft_ops.ifftshift(shifted, axes=(0, 1)).eval(),
            np.fft.ifftshift(shifted, axes=(0, 1)))
  
  
  if __name__ == "__main__":
    test.main()
   # Copyright 2018 The TensorFlow Authors. All Rights Reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  # ==============================================================================
  """Tests for the `LatencyAllEdges` optimization."""
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  from tensorflow.python.data.experimental.kernel_tests import stats_dataset_test_base
  from tensorflow.python.data.experimental.ops import optimization
  from tensorflow.python.data.experimental.ops import stats_aggregator
  from tensorflow.python.data.ops import dataset_ops
  from tensorflow.python.platform import test
  
  
  class LatencyAllEdgesTest(stats_dataset_test_base.StatsDatasetTestBase):
  
    def testLatencyStatsOptimization(self):
      aggregator = stats_aggregator.StatsAggregator()
      dataset = dataset_ops.Dataset.from_tensors(1).apply(
          optimization.assert_next(
              ["LatencyStats", "Map", "LatencyStats", "Prefetch",
               "LatencyStats"])).map(lambda x: x * x).prefetch(1)
      options = dataset_ops.Options()
      options.experimental_optimization.apply_default_optimizations = False
      options.experimental_stats.latency_all_edges = True
      options.experimental_stats.aggregator = aggregator
      dataset = dataset.with_options(options)
      self.assertDatasetProduces(
          dataset,
          expected_output=[1],
          requires_initialization=True,
          num_test_iterations=1)
      handle = self.getHandle(aggregator)
      self.assertStatisticsHasCount(
          handle, self.regexForNodeName("record_latency::TensorDataset"), 1)
      self.assertStatisticsHasCount(
          handle, self.regexForNodeName("record_latency::MapDataset"), 1)
      self.assertStatisticsHasCount(
          handle, self.regexForNodeName("record_latency::PrefetchDataset"), 1)
  
  
  if __name__ == "__main__":
    test.main()
   # Copyright 2017 The TensorFlow Authors. All Rights Reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  # ==============================================================================
  """Core module for TensorFlow distribution objects and helpers."""
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  from tensorflow.python.util import deprecation
  
  
  # pylint: disable=wildcard-import,unused-import,g-import-not-at-top
  with deprecation.silence():
    from tensorflow.python.ops.distributions.bernoulli import Bernoulli
    from tensorflow.python.ops.distributions.beta import Beta
    from tensorflow.python.ops.distributions.categorical import Categorical
    from tensorflow.python.ops.distributions.dirichlet import Dirichlet
    from tensorflow.python.ops.distributions.dirichlet_multinomial import DirichletMultinomial
    from tensorflow.python.ops.distributions.distribution import *
    from tensorflow.python.ops.distributions.exponential import Exponential
    from tensorflow.python.ops.distributions.gamma import Gamma
    from tensorflow.python.ops.distributions.kullback_leibler import *
    from tensorflow.python.ops.distributions.laplace import Laplace
    from tensorflow.python.ops.distributions.multinomial import Multinomial
    from tensorflow.python.ops.distributions.normal import Normal
    from tensorflow.python.ops.distributions.student_t import StudentT
    from tensorflow.python.ops.distributions.uniform import Uniform
  # pylint: enable=wildcard-import,unused-import
  del deprecation
   # Copyright 2017 The TensorFlow Authors. All Rights Reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  # ==============================================================================
  """Tests for MultivariateNormalFullCovariance."""
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import numpy as np
  from scipy import stats
  from tensorflow.contrib import distributions
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import math_ops
  from tensorflow.python.ops import nn_ops
  from tensorflow.python.platform import test
  
  
  ds = distributions
  rng = np.random.RandomState(42)
  
  
  class MultivariateNormalFullCovarianceTest(test.TestCase):
  
    def _random_pd_matrix(self, *shape):
      mat = rng.rand(*shape)
      chol = ds.matrix_diag_transform(mat, transform=nn_ops.softplus)
      chol = array_ops.matrix_band_part(chol, -1, 0)
      return math_ops.matmul(chol, chol, adjoint_b=True).eval()
  
    def testRaisesIfInitializedWithNonSymmetricMatrix(self):
      with self.cached_session():
        mu = [1., 2.]
        sigma = [[1., 0.], [1., 1.]]  # Nonsingular, but not symmetric
        mvn = ds.MultivariateNormalFullCovariance(mu, sigma, validate_args=True)
        with self.assertRaisesOpError("not symmetric"):
          mvn.covariance().eval()
  
    def testNamePropertyIsSetByInitArg(self):
      with self.cached_session():
        mu = [1., 2.]
        sigma = [[1., 0.], [0., 1.]]
        mvn = ds.MultivariateNormalFullCovariance(mu, sigma, name="Billy")
        self.assertEqual(mvn.name, "Billy/")
  
    def testDoesNotRaiseIfInitializedWithSymmetricMatrix(self):
      with self.cached_session():
        mu = rng.rand(10)
        sigma = self._random_pd_matrix(10, 10)
        mvn = ds.MultivariateNormalFullCovariance(mu, sigma, validate_args=True)
        # Should not raise
        mvn.covariance().eval()
  
    def testLogPDFScalarBatch(self):
      with self.cached_session():
        mu = rng.rand(2)
        sigma = self._random_pd_matrix(2, 2)
        mvn = ds.MultivariateNormalFullCovariance(mu, sigma, validate_args=True)
        x = rng.rand(2)
  
        log_pdf = mvn.log_prob(x)
        pdf = mvn.prob(x)
  
        scipy_mvn = stats.multivariate_normal(mean=mu, cov=sigma)
  
        expected_log_pdf = scipy_mvn.logpdf(x)
        expected_pdf = scipy_mvn.pdf(x)
        self.assertEqual((), log_pdf.get_shape())
        self.assertEqual((), pdf.get_shape())
        self.assertAllClose(expected_log_pdf, log_pdf.eval())
        self.assertAllClose(expected_pdf, pdf.eval())
  
    def testLogPDFScalarBatchCovarianceNotProvided(self):
      with self.cached_session():
        mu = rng.rand(2)
        mvn = ds.MultivariateNormalFullCovariance(
            mu, covariance_matrix=None, validate_args=True)
        x = rng.rand(2)
  
        log_pdf = mvn.log_prob(x)
        pdf = mvn.prob(x)
  
        # Initialize a scipy_mvn with the default covariance.
        scipy_mvn = stats.multivariate_normal(mean=mu, cov=np.eye(2))
  
        expected_log_pdf = scipy_mvn.logpdf(x)
        expected_pdf = scipy_mvn.pdf(x)
        self.assertEqual((), log_pdf.get_shape())
        self.assertEqual((), pdf.get_shape())
        self.assertAllClose(expected_log_pdf, log_pdf.eval())
        self.assertAllClose(expected_pdf, pdf.eval())
  
    def testShapes(self):
      with self.cached_session():
        mu = rng.rand(3, 5, 2)
        covariance = self._random_pd_matrix(3, 5, 2, 2)
  
        mvn = ds.MultivariateNormalFullCovariance(
            mu, covariance, validate_args=True)
  
        # Shapes known at graph construction time.
        self.assertEqual((2,), tuple(mvn.event_shape.as_list()))
        self.assertEqual((3, 5), tuple(mvn.batch_shape.as_list()))
  
        # Shapes known at runtime.
        self.assertEqual((2,), tuple(mvn.event_shape_tensor().eval()))
        self.assertEqual((3, 5), tuple(mvn.batch_shape_tensor().eval()))
  
    def _random_mu_and_sigma(self, batch_shape, event_shape):
      # This ensures sigma is positive def.
      mat_shape = batch_shape + event_shape + event_shape
      mat = rng.randn(*mat_shape)
      perm = np.arange(mat.ndim)
      perm[-2:] = [perm[-1], perm[-2]]
      sigma = np.matmul(mat, np.transpose(mat, perm))
  
      mu_shape = batch_shape + event_shape
      mu = rng.randn(*mu_shape)
  
      return mu, sigma
  
    def testKLBatch(self):
      batch_shape = [2]
      event_shape = [3]
      with self.cached_session():
        mu_a, sigma_a = self._random_mu_and_sigma(batch_shape, event_shape)
        mu_b, sigma_b = self._random_mu_and_sigma(batch_shape, event_shape)
        mvn_a = ds.MultivariateNormalFullCovariance(
            loc=mu_a,
            covariance_matrix=sigma_a,
            validate_args=True)
        mvn_b = ds.MultivariateNormalFullCovariance(
            loc=mu_b,
            covariance_matrix=sigma_b,
            validate_args=True)
  
        kl = ds.kl_divergence(mvn_a, mvn_b)
        self.assertEqual(batch_shape, kl.get_shape())
  
        kl_v = kl.eval()
        expected_kl_0 = _compute_non_batch_kl(mu_a[0, :], sigma_a[0, :, :],
                                              mu_b[0, :], sigma_b[0, :])
        expected_kl_1 = _compute_non_batch_kl(mu_a[1, :], sigma_a[1, :, :],
                                              mu_b[1, :], sigma_b[1, :])
        self.assertAllClose(expected_kl_0, kl_v[0])
        self.assertAllClose(expected_kl_1, kl_v[1])
  
    def testKLBatchBroadcast(self):
      batch_shape = [2]
      event_shape = [3]
      with self.cached_session():
        mu_a, sigma_a = self._random_mu_and_sigma(batch_shape, event_shape)
        # No batch shape.
        mu_b, sigma_b = self._random_mu_and_sigma([], event_shape)
        mvn_a = ds.MultivariateNormalFullCovariance(
            loc=mu_a,
            covariance_matrix=sigma_a,
            validate_args=True)
        mvn_b = ds.MultivariateNormalFullCovariance(
            loc=mu_b,
            covariance_matrix=sigma_b,
            validate_args=True)
  
        kl = ds.kl_divergence(mvn_a, mvn_b)
        self.assertEqual(batch_shape, kl.get_shape())
  
        kl_v = kl.eval()
        expected_kl_0 = _compute_non_batch_kl(mu_a[0, :], sigma_a[0, :, :],
                                              mu_b, sigma_b)
        expected_kl_1 = _compute_non_batch_kl(mu_a[1, :], sigma_a[1, :, :],
                                              mu_b, sigma_b)
        self.assertAllClose(expected_kl_0, kl_v[0])
        self.assertAllClose(expected_kl_1, kl_v[1])
  
  
  def _compute_non_batch_kl(mu_a, sigma_a, mu_b, sigma_b):
    """Non-batch KL for N(mu_a, sigma_a), N(mu_b, sigma_b)."""
    # Check using numpy operations
    # This mostly repeats the tensorflow code _kl_mvn_mvn(), but in numpy.
    # So it is important to also check that KL(mvn, mvn) = 0.
    sigma_b_inv = np.linalg.inv(sigma_b)
  
    t = np.trace(sigma_b_inv.dot(sigma_a))
    q = (mu_b - mu_a).dot(sigma_b_inv).dot(mu_b - mu_a)
    k = mu_a.shape[0]
    l = np.log(np.linalg.det(sigma_b) / np.linalg.det(sigma_a))
  
    return 0.5 * (t + q - k + l)
  
  
  if __name__ == "__main__":
    test.main()
   # Copyright 2016 The TensorFlow Authors. All Rights Reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  # ==============================================================================
  """Tests for Keras initializers."""
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import numpy as np
  
  from tensorflow.python import keras
  from tensorflow.python import tf2
  from tensorflow.python.framework import test_util
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import init_ops
  from tensorflow.python.platform import test
  
  
  @test_util.run_all_in_graph_and_eager_modes
  class KerasInitializersTest(test.TestCase):
  
    def _runner(self, init, shape, target_mean=None, target_std=None,
                target_max=None, target_min=None):
      variable = keras.backend.variable(init(shape))
      output = keras.backend.get_value(variable)
      # Test serialization (assumes deterministic behavior).
      config = init.get_config()
      reconstructed_init = init.__class__.from_config(config)
      variable = keras.backend.variable(reconstructed_init(shape))
      output_2 = keras.backend.get_value(variable)
      self.assertAllClose(output, output_2, atol=1e-4)
  
    def test_uniform(self):
      tensor_shape = (9, 6, 7)
      with self.cached_session():
        self._runner(
            keras.initializers.RandomUniformV2(minval=-1, maxval=1, seed=124),
            tensor_shape,
            target_mean=0.,
            target_max=1,
            target_min=-1)
  
    def test_normal(self):
      tensor_shape = (8, 12, 99)
      with self.cached_session():
        self._runner(
            keras.initializers.RandomNormalV2(mean=0, stddev=1, seed=153),
            tensor_shape,
            target_mean=0.,
            target_std=1)
  
    def test_truncated_normal(self):
      tensor_shape = (12, 99, 7)
      with self.cached_session():
        self._runner(
            keras.initializers.TruncatedNormalV2(mean=0, stddev=1, seed=126),
            tensor_shape,
            target_mean=0.,
            target_max=2,
            target_min=-2)
  
    def test_constant(self):
      tensor_shape = (5, 6, 4)
      with self.cached_session():
        self._runner(
            keras.initializers.ConstantV2(2.),
            tensor_shape,
            target_mean=2,
            target_max=2,
            target_min=2)
  
    def test_lecun_uniform(self):
      tensor_shape = (5, 6, 4, 2)
      with self.cached_session():
        fan_in, _ = init_ops._compute_fans(tensor_shape)
        std = np.sqrt(1. / fan_in)
        self._runner(
            keras.initializers.lecun_uniformV2(seed=123),
            tensor_shape,
            target_mean=0.,
            target_std=std)
  
    def test_glorot_uniform(self):
      tensor_shape = (5, 6, 4, 2)
      with self.cached_session():
        fan_in, fan_out = init_ops._compute_fans(tensor_shape)
        std = np.sqrt(2. / (fan_in + fan_out))
        self._runner(
            keras.initializers.GlorotUniformV2(seed=123),
            tensor_shape,
            target_mean=0.,
            target_std=std)
  
    def test_he_uniform(self):
      tensor_shape = (5, 6, 4, 2)
      with self.cached_session():
        fan_in, _ = init_ops._compute_fans(tensor_shape)
        std = np.sqrt(2. / fan_in)
        self._runner(
            keras.initializers.he_uniformV2(seed=123),
            tensor_shape,
            target_mean=0.,
            target_std=std)
  
    def test_lecun_normal(self):
      tensor_shape = (5, 6, 4, 2)
      with self.cached_session():
        fan_in, _ = init_ops._compute_fans(tensor_shape)
        std = np.sqrt(1. / fan_in)
        self._runner(
            keras.initializers.lecun_normalV2(seed=123),
            tensor_shape,
            target_mean=0.,
            target_std=std)
  
    def test_glorot_normal(self):
      tensor_shape = (5, 6, 4, 2)
      with self.cached_session():
        fan_in, fan_out = init_ops._compute_fans(tensor_shape)
        std = np.sqrt(2. / (fan_in + fan_out))
        self._runner(
            keras.initializers.GlorotNormalV2(seed=123),
            tensor_shape,
            target_mean=0.,
            target_std=std)
  
    def test_he_normal(self):
      tensor_shape = (5, 6, 4, 2)
      with self.cached_session():
        fan_in, _ = init_ops._compute_fans(tensor_shape)
        std = np.sqrt(2. / fan_in)
        self._runner(
            keras.initializers.he_normalV2(seed=123),
            tensor_shape,
            target_mean=0.,
            target_std=std)
  
    def test_orthogonal(self):
      tensor_shape = (20, 20)
      with self.cached_session():
        self._runner(
            keras.initializers.OrthogonalV2(seed=123),
            tensor_shape,
            target_mean=0.)
  
    def test_identity(self):
      with self.cached_session():
        tensor_shape = (3, 4, 5)
        with self.assertRaises(ValueError):
          self._runner(
              keras.initializers.IdentityV2(),
              tensor_shape,
              target_mean=1. / tensor_shape[0],
              target_max=1.)
  
        tensor_shape = (3, 3)
        self._runner(
            keras.initializers.IdentityV2(),
            tensor_shape,
            target_mean=1. / tensor_shape[0],
            target_max=1.)
  
    def test_zero(self):
      tensor_shape = (4, 5)
      with self.cached_session():
        self._runner(
            keras.initializers.ZerosV2(),
            tensor_shape,
            target_mean=0.,
            target_max=0.)
  
    def test_one(self):
      tensor_shape = (4, 5)
      with self.cached_session():
        self._runner(
            keras.initializers.OnesV2(),
            tensor_shape,
            target_mean=1.,
            target_max=1.)
  
    def test_default_random_uniform(self):
      ru = keras.initializers.get('uniform')
      self.assertEqual(ru.minval, -0.05)
      self.assertEqual(ru.maxval, 0.05)
  
    def test_default_random_normal(self):
      rn = keras.initializers.get('normal')
      self.assertEqual(rn.mean, 0.0)
      self.assertEqual(rn.stddev, 0.05)
  
    def test_default_truncated_normal(self):
      tn = keras.initializers.get('truncated_normal')
      self.assertEqual(tn.mean, 0.0)
      self.assertEqual(tn.stddev, 0.05)
  
    def test_initializer_v2_get(self):
      tf2_force_enabled = tf2._force_enable  # pylint: disable=protected-access
      try:
        tf2.enable()
        rn = keras.initializers.get('random_normal')
        self.assertIn('init_ops_v2', rn.__class__.__module__)
      finally:
        tf2._force_enable = tf2_force_enabled  # pylint: disable=protected-access
  
    def test_custom_initializer_saving(self):
  
      def my_initializer(shape, dtype=None):
        return array_ops.ones(shape, dtype=dtype)
  
      inputs = keras.Input((10,))
      outputs = keras.layers.Dense(1, kernel_initializer=my_initializer)(inputs)
      model = keras.Model(inputs, outputs)
      model2 = model.from_config(
          model.get_config(), custom_objects={'my_initializer': my_initializer})
      self.assertEqual(model2.layers[1].kernel_initializer, my_initializer)
  
    @test_util.run_v2_only
    def test_load_external_variance_scaling_v2(self):
      external_serialized_json = {
          'class_name': 'VarianceScaling',
          'config': {
              'distribution': 'normal',
              'mode': 'fan_avg',
              'scale': 1.0,
              'seed': None
          }
      }
      initializer = keras.initializers.deserialize(external_serialized_json)
      self.assertEqual(initializer.distribution, 'truncated_normal')
  
  
  if __name__ == '__main__':
    test.main()
   # Copyright 2016 The TensorFlow Authors. All Rights Reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  # ==============================================================================
  """LSTM Block Cell ops."""
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  from absl.testing import parameterized
  import numpy as np
  
  from tensorflow.contrib.rnn.python.kernel_tests import benchmarking
  from tensorflow.contrib.rnn.python.ops import lstm_ops
  from tensorflow.python.client import session
  from tensorflow.python.framework import constant_op
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import ops
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import gen_array_ops
  from tensorflow.python.ops import gen_bitwise_ops
  from tensorflow.python.ops import gradients_impl
  from tensorflow.python.ops import init_ops
  from tensorflow.python.ops import rnn
  from tensorflow.python.ops import rnn_cell
  from tensorflow.python.ops import variable_scope
  from tensorflow.python.ops import variables
  from tensorflow.python.platform import test
  
  block_lstm = lstm_ops._block_lstm  # pylint: disable=protected-access
  
  
  class _MaskedRandomUniformInitializer(init_ops.RandomUniform):
    """Initializer for uniform dist tensors with trailing bits zeroed-out.
  
    Allow returning tensors with last few mantissa bits set to 0. This potentially
    helps avoid getting into precision issues when testing low precision (float16)
    computation.
    """
  
    def __init__(self,
                 minval=0,
                 maxval=None,
                 seed=None,
                 dtype=dtypes.float16,
                 num_valid_mantissa_bits=4):
      """Constructor.
  
      Args:
        minval: A python scalar or a scalar tensor. Lower bound of the range of
          random values to generate.
        maxval: A python scalar or a scalar tensor. Upper bound of the range of
          random values to generate.  Defaults to 1 for float types.
        seed: A Python integer. Used to create random seeds. See
          `tf.compat.v1.set_random_seed` for behavior.
        dtype: The data type. Only supports tf.float16 for now.
        num_valid_mantissa_bits: number of non-zero mantissa bits, default to 4.
  
      Raises:
        ValueError: An error if `dtype` is not tf.float16.
      """
      if dtype not in (dtypes.float16,):
        raise ValueError("dtype: %s not supported" % dtype.name)
  
      super(_MaskedRandomUniformInitializer, self).__init__(
          minval=minval, maxval=maxval, seed=seed, dtype=dtype)
      self._num_mantissa_bits = 10
      self._num_valid_mantissa_bits = num_valid_mantissa_bits
  
    def __call__(self, shape, dtype=dtypes.float16, partition_info=None):
      if dtype and dtype != dtypes.float16:
        raise ValueError("dtype: %s not supported" % dtype.name)
      res = super(_MaskedRandomUniformInitializer, self).__call__(
          shape, dtype, partition_info)
      # get uint16 view of the underlying buffer.
      res = gen_array_ops.bitcast(res, dtypes.uint16)
  
      # mask the last `shift` mantissa bits.
      shift = self._num_mantissa_bits - self._num_valid_mantissa_bits
      mask = (0xffff >> shift) << shift
      res = gen_bitwise_ops.bitwise_and(res, mask)
  
      # restore float16 view.
      return gen_array_ops.bitcast(res, dtype)
  
  
  def _get_initializer(init_bound, dtype, seed):
    if dtype == dtypes.float16:
      return _MaskedRandomUniformInitializer(
          -init_bound, init_bound, dtype=dtype, seed=seed)
    else:
      return init_ops.random_uniform_initializer(
          -init_bound, init_bound, dtype=dtype, seed=seed)
  
  
  def blocks_match(sess, use_peephole, dtype=dtypes.float32, cell_clip=None):
    batch_size = 2
    input_size = 3
    cell_size = 4
    sequence_length = 4
  
    inputs = []
    for _ in range(sequence_length):
      inp = ops.convert_to_tensor(
          np.random.randn(batch_size, input_size), dtype=dtype)
      inputs.append(inp)
    stacked_inputs = array_ops.stack(inputs)
  
    init_bound = 1e-1 if dtype == dtypes.float16 else 1e-2
    initializer = _get_initializer(init_bound, dtype=dtype, seed=19890212)
  
    with variable_scope.variable_scope("test", initializer=initializer):
      # magic naming so that the cells pick up these variables and reuse them
      if use_peephole:
        wci = variable_scope.get_variable(
            "rnn/lstm_cell/w_i_diag", shape=[cell_size], dtype=dtype)
        wcf = variable_scope.get_variable(
            "rnn/lstm_cell/w_f_diag", shape=[cell_size], dtype=dtype)
        wco = variable_scope.get_variable(
            "rnn/lstm_cell/w_o_diag", shape=[cell_size], dtype=dtype)
  
      w = variable_scope.get_variable(
          "rnn/lstm_cell/kernel",
          shape=[input_size + cell_size, cell_size * 4],
          dtype=dtype)
      b = variable_scope.get_variable(
          "rnn/lstm_cell/bias",
          shape=[cell_size * 4],
          dtype=dtype,
          initializer=init_ops.zeros_initializer())
  
      basic_cell = rnn_cell.LSTMCell(
          cell_size,
          use_peepholes=use_peephole,
          cell_clip=cell_clip,
          dtype=dtype,
          state_is_tuple=True,
          reuse=True)
      basic_outputs_op, basic_state_op = rnn.static_rnn(
          basic_cell, inputs, dtype=dtype)
  
      if use_peephole:
        _, _, _, _, _, _, block_outputs_op = block_lstm(
            ops.convert_to_tensor(sequence_length, dtype=dtypes.int64),
            inputs,
            w,
            b,
            wci=wci,
            wcf=wcf,
            wco=wco,
            cell_clip=cell_clip,
            use_peephole=True)
      else:
        _, _, _, _, _, _, block_outputs_op = block_lstm(
            ops.convert_to_tensor(sequence_length, dtype=dtypes.int64),
            inputs,
            w,
            b,
            cell_clip=cell_clip)
  
      fused_cell = lstm_ops.LSTMBlockFusedCell(
          cell_size,
          cell_clip=cell_clip,
          use_peephole=use_peephole,
          reuse=True,
          name="rnn/lstm_cell")
      fused_outputs_op, fused_state_op = fused_cell(stacked_inputs, dtype=dtype)
  
      sess.run([variables.global_variables_initializer()])
      basic_outputs, basic_state = sess.run([basic_outputs_op, basic_state_op[0]])
      basic_grads = sess.run(gradients_impl.gradients(basic_outputs_op, inputs))
      xs = [w, b]
      if use_peephole:
        xs += [wci, wcf, wco]
      basic_wgrads = sess.run(gradients_impl.gradients(basic_outputs_op, xs))
  
      block_outputs = sess.run(block_outputs_op)
      block_grads = sess.run(gradients_impl.gradients(block_outputs_op, inputs))
      block_wgrads = sess.run(gradients_impl.gradients(block_outputs_op, xs))
  
      xs = [w, b]
      if use_peephole:
        xs += [wci, wcf, wco]
      fused_outputs, fused_state = sess.run([fused_outputs_op, fused_state_op[0]])
      fused_grads = sess.run(gradients_impl.gradients(fused_outputs_op, inputs))
      fused_wgrads = sess.run(gradients_impl.gradients(fused_outputs_op, xs))
  
      return (basic_state, fused_state, basic_outputs, block_outputs,
              fused_outputs, basic_grads, block_grads, fused_grads, basic_wgrads,
              block_wgrads, fused_wgrads)
  
  
  class LSTMBlockCellTest(test.TestCase, parameterized.TestCase):
  
    TEST_CASES = ({
        "testcase_name": "Fp32",
        "dtype": dtypes.float32,
        "rtol": 1e-6,
        "atol": 1e-6
    }, {
        "testcase_name": "Fp16",
        "dtype": dtypes.float16,
        "rtol": 8e-3,
        "atol": 8e-4
    })
  
    def testNoneDimsWithDynamicRNN(self):
      with self.session(use_gpu=True, graph=ops.Graph()) as sess:
        batch_size = 4
        num_steps = 5
        input_dim = 6
        cell_size = 7
  
        cell = lstm_ops.LSTMBlockCell(cell_size)
        x = array_ops.placeholder(dtypes.float32, shape=(None, None, input_dim))
  
        output, _ = rnn.dynamic_rnn(
            cell, x, time_major=True, dtype=dtypes.float32)
        sess.run(variables.global_variables_initializer())
        feed = {}
        feed[x] = np.random.randn(num_steps, batch_size, input_dim)
        sess.run(output, feed)
  
    def testLSTMBlockCell(self):
      with self.session(use_gpu=True, graph=ops.Graph()) as sess:
        with variable_scope.variable_scope(
            "root", initializer=init_ops.constant_initializer(0.5)):
          x = array_ops.zeros([1, 2])
          m0 = array_ops.zeros([1, 2])
          m1 = array_ops.zeros([1, 2])
          m2 = array_ops.zeros([1, 2])
          m3 = array_ops.zeros([1, 2])
          g, ((out_m0, out_m1), (out_m2, out_m3)) = rnn_cell.MultiRNNCell(
              [lstm_ops.LSTMBlockCell(2)
               for _ in range(2)], state_is_tuple=True)(x, ((m0, m1), (m2, m3)))
          sess.run([variables.global_variables_initializer()])
          res = sess.run([g, out_m0, out_m1, out_m2, out_m3], {
              x.name: np.array([[1., 1.]]),
              m0.name: 0.1 * np.ones([1, 2]),
              m1.name: 0.1 * np.ones([1, 2]),
              m2.name: 0.1 * np.ones([1, 2]),
              m3.name: 0.1 * np.ones([1, 2])
          })
          self.assertEqual(len(res), 5)
          self.assertAllClose(res[0], [[0.24024698, 0.24024698]])
          # These numbers are from testBasicLSTMCell and only test c/h.
          self.assertAllClose(res[1], [[0.68967271, 0.68967271]])
          self.assertAllClose(res[2], [[0.44848421, 0.44848421]])
          self.assertAllClose(res[3], [[0.39897051, 0.39897051]])
          self.assertAllClose(res[4], [[0.24024698, 0.24024698]])
  
    def testCompatibleNames(self):
      with self.session(use_gpu=True, graph=ops.Graph()):
        cell = rnn_cell.LSTMCell(10)
        pcell = rnn_cell.LSTMCell(10, use_peepholes=True)
        inputs = [array_ops.zeros([4, 5])] * 6
        rnn.static_rnn(cell, inputs, dtype=dtypes.float32, scope="basic")
        rnn.static_rnn(pcell, inputs, dtype=dtypes.float32, scope="peephole")
        basic_names = {
            v.name: v.get_shape()
            for v in variables.trainable_variables()
        }
  
      with self.session(use_gpu=True, graph=ops.Graph()):
        cell = lstm_ops.LSTMBlockCell(10)
        pcell = lstm_ops.LSTMBlockCell(10, use_peephole=True)
        inputs = [array_ops.zeros([4, 5])] * 6
        rnn.static_rnn(cell, inputs, dtype=dtypes.float32, scope="basic")
        rnn.static_rnn(pcell, inputs, dtype=dtypes.float32, scope="peephole")
        block_names = {
            v.name: v.get_shape()
            for v in variables.trainable_variables()
        }
  
      with self.session(use_gpu=True, graph=ops.Graph()):
        cell = lstm_ops.LSTMBlockFusedCell(10)
        pcell = lstm_ops.LSTMBlockFusedCell(10, use_peephole=True)
        inputs = array_ops.stack([array_ops.zeros([4, 5])] * 6)
        cell(inputs, dtype=dtypes.float32, scope="basic/lstm_cell")
        pcell(inputs, dtype=dtypes.float32, scope="peephole/lstm_cell")
        fused_names = {
            v.name: v.get_shape()
            for v in variables.trainable_variables()
        }
  
      self.assertEqual(basic_names, block_names)
      self.assertEqual(basic_names, fused_names)
  
    def testLSTMBasicToBlockCell(self):
      with self.session(use_gpu=True) as sess:
        x = array_ops.zeros([1, 2])
        x_values = np.random.randn(1, 2)
  
        m0_val = 0.1 * np.ones([1, 2])
        m1_val = -0.1 * np.ones([1, 2])
        m2_val = -0.2 * np.ones([1, 2])
        m3_val = 0.2 * np.ones([1, 2])
  
        initializer = init_ops.random_uniform_initializer(
            -0.01, 0.01, seed=19890212)
        with variable_scope.variable_scope("basic", initializer=initializer):
          m0 = array_ops.zeros([1, 2])
          m1 = array_ops.zeros([1, 2])
          m2 = array_ops.zeros([1, 2])
          m3 = array_ops.zeros([1, 2])
          g, ((out_m0, out_m1), (out_m2, out_m3)) = rnn_cell.MultiRNNCell(
              [rnn_cell.BasicLSTMCell(2, state_is_tuple=True) for _ in range(2)],
              state_is_tuple=True)(x, ((m0, m1), (m2, m3)))
          sess.run([variables.global_variables_initializer()])
          basic_res = sess.run([g, out_m0, out_m1, out_m2, out_m3], {
              x.name: x_values,
              m0.name: m0_val,
              m1.name: m1_val,
              m2.name: m2_val,
              m3.name: m3_val
          })
  
        with variable_scope.variable_scope("block", initializer=initializer):
          m0 = array_ops.zeros([1, 2])
          m1 = array_ops.zeros([1, 2])
          m2 = array_ops.zeros([1, 2])
          m3 = array_ops.zeros([1, 2])
          g, ((out_m0, out_m1), (out_m2, out_m3)) = rnn_cell.MultiRNNCell(
              [lstm_ops.LSTMBlockCell(2)
               for _ in range(2)], state_is_tuple=True)(x, ((m0, m1), (m2, m3)))
          sess.run([variables.global_variables_initializer()])
          block_res = sess.run([g, out_m0, out_m1, out_m2, out_m3], {
              x.name: x_values,
              m0.name: m0_val,
              m1.name: m1_val,
              m2.name: m2_val,
              m3.name: m3_val
          })
  
        self.assertEqual(len(basic_res), len(block_res))
        for basic, block in zip(basic_res, block_res):
          self.assertAllClose(basic, block)
  
    def testLSTMBasicToBlockCellPeeping(self):
      with self.session(use_gpu=True) as sess:
        x = array_ops.zeros([1, 2])
        x_values = np.random.randn(1, 2)
  
        m0_val = 0.1 * np.ones([1, 2])
        m1_val = -0.1 * np.ones([1, 2])
        m2_val = -0.2 * np.ones([1, 2])
        m3_val = 0.2 * np.ones([1, 2])
  
        initializer = init_ops.random_uniform_initializer(
            -0.01, 0.01, seed=19890212)
        with variable_scope.variable_scope("basic", initializer=initializer):
          m0 = array_ops.zeros([1, 2])
          m1 = array_ops.zeros([1, 2])
          m2 = array_ops.zeros([1, 2])
          m3 = array_ops.zeros([1, 2])
          g, ((out_m0, out_m1), (out_m2, out_m3)) = rnn_cell.MultiRNNCell(
              [
                  rnn_cell.LSTMCell(2, use_peepholes=True, state_is_tuple=True)
                  for _ in range(2)
              ],
              state_is_tuple=True)(x, ((m0, m1), (m2, m3)))
          sess.run([variables.global_variables_initializer()])
          basic_res = sess.run([g, out_m0, out_m1, out_m2, out_m3], {
              x.name: x_values,
              m0.name: m0_val,
              m1.name: m1_val,
              m2.name: m2_val,
              m3.name: m3_val
          })
  
        with variable_scope.variable_scope("block", initializer=initializer):
          m0 = array_ops.zeros([1, 2])
          m1 = array_ops.zeros([1, 2])
          m2 = array_ops.zeros([1, 2])
          m3 = array_ops.zeros([1, 2])
          g, ((out_m0, out_m1), (out_m2, out_m3)) = rnn_cell.MultiRNNCell(
              [lstm_ops.LSTMBlockCell(2, use_peephole=True) for _ in range(2)],
              state_is_tuple=True)(x, ((m0, m1), (m2, m3)))
          sess.run([variables.global_variables_initializer()])
          block_res = sess.run([g, out_m0, out_m1, out_m2, out_m3], {
              x.name: x_values,
              m0.name: m0_val,
              m1.name: m1_val,
              m2.name: m2_val,
              m3.name: m3_val
          })
  
        self.assertEqual(len(basic_res), len(block_res))
        for basic, block in zip(basic_res, block_res):
          self.assertAllClose(basic, block)
  
    def LSTMBasicToBlockTestHelper(self,
                                   dtype=dtypes.float32,
                                   use_peephole=False,
                                   cell_clip=None,
                                   rtol=1e-6,
                                   atol=1e-6):
      with self.session(use_gpu=True, graph=ops.Graph()) as sess:
        (basic_state, fused_state, basic_outputs, block_outputs, fused_outputs,
         basic_grads, block_grads, fused_grads, basic_wgrads, block_wgrads,
         fused_wgrads) = blocks_match(
             sess, use_peephole=use_peephole, dtype=dtype, cell_clip=cell_clip)
  
        self.assertAllClose(basic_outputs, block_outputs, rtol=rtol, atol=atol)
        self.assertAllClose(basic_grads, block_grads, rtol=rtol, atol=atol)
        for basic, block in zip(basic_wgrads, block_wgrads):
          self.assertAllClose(basic, block, rtol=rtol, atol=atol)
  
        self.assertAllClose(basic_outputs, fused_outputs, rtol=rtol, atol=atol)
        self.assertAllClose(basic_state, fused_state, rtol=rtol, atol=atol)
        self.assertAllClose(basic_grads, fused_grads, rtol=rtol, atol=atol)
        for basic, fused in zip(basic_wgrads, fused_wgrads):
          self.assertAllClose(basic, fused, rtol=rtol, atol=atol)
  
    @parameterized.named_parameters(*TEST_CASES)
    def testLSTMBasicToBlock(self, dtype, rtol, atol):
      self.LSTMBasicToBlockTestHelper(
          dtype, use_peephole=False, rtol=rtol, atol=atol)
  
    @parameterized.named_parameters(*TEST_CASES)
    def testLSTMBasicToBlockPeeping(self, dtype, rtol, atol):
      self.LSTMBasicToBlockTestHelper(
          dtype, use_peephole=True, rtol=rtol, atol=atol)
  
    @parameterized.named_parameters(*TEST_CASES)
    def testLSTMBasicToBlockCellClip(self, dtype, rtol, atol):
      self.LSTMBasicToBlockTestHelper(
          dtype, use_peephole=True, cell_clip=0.5, rtol=rtol, atol=atol)
  
    def testLSTMFusedSequenceLengths(self):
      """Verify proper support for sequence lengths in LSTMBlockFusedCell."""
      with self.session(use_gpu=True) as sess:
        batch_size = 3
        input_size = 4
        cell_size = 5
        max_sequence_length = 6
  
        inputs = []
        for _ in range(max_sequence_length):
          inp = ops.convert_to_tensor(
              np.random.randn(batch_size, input_size), dtype=dtypes.float32)
          inputs.append(inp)
        seq_lengths = constant_op.constant([3, 4, 5])
        cell_inputs = array_ops.stack(inputs)
  
        initializer = init_ops.random_uniform_initializer(
            -0.01, 0.01, seed=19890213)
  
        with variable_scope.variable_scope("lstm_cell", initializer=initializer):
          # magic naming so that the cells pick up these variables and reuse them
          variable_scope.get_variable(
              "kernel",
              shape=[input_size + cell_size, cell_size * 4],
              dtype=dtypes.float32)
  
          variable_scope.get_variable(
              "bias",
              shape=[cell_size * 4],
              dtype=dtypes.float32,
              initializer=init_ops.zeros_initializer())
  
        cell = lstm_ops.LSTMBlockFusedCell(
            cell_size, cell_clip=0, use_peephole=False, reuse=True,
            name="lstm_cell")
  
        fused_outputs_op, fused_state_op = cell(
            cell_inputs, dtype=dtypes.float32, sequence_length=seq_lengths)
  
        cell_vars = [
            v for v in variables.trainable_variables()
            if v.name.endswith("kernel") or v.name.endswith("bias")
        ]
  
        # Verify that state propagation works if we turn our sequence into
        # tiny (single-time) subsequences, i.e. unfuse the cell
        unfused_outputs_op = []
        state = None
        with variable_scope.variable_scope(
            variable_scope.get_variable_scope(), reuse=True):
          for i, inp in enumerate(inputs):
            lengths = [int(i < l) for l in seq_lengths.eval()]
            output, state = cell(
                array_ops.expand_dims(inp, 0),
                initial_state=state,
                dtype=dtypes.float32,
                sequence_length=lengths)
            unfused_outputs_op.append(output[0])
        unfused_outputs_op = array_ops.stack(unfused_outputs_op)
  
        sess.run([variables.global_variables_initializer()])
        unfused_outputs, unfused_state = sess.run([unfused_outputs_op, state[0]])
        unfused_grads = sess.run(
            gradients_impl.gradients(unfused_outputs_op, inputs))
        unfused_wgrads = sess.run(
            gradients_impl.gradients(unfused_outputs_op, cell_vars))
  
        fused_outputs, fused_state = sess.run(
            [fused_outputs_op, fused_state_op[0]])
        fused_grads = sess.run(gradients_impl.gradients(fused_outputs_op, inputs))
        fused_wgrads = sess.run(
            gradients_impl.gradients(fused_outputs_op, cell_vars))
  
        self.assertAllClose(fused_outputs, unfused_outputs)
        self.assertAllClose(fused_state, unfused_state)
        self.assertAllClose(fused_grads, unfused_grads)
        for fused, unfused in zip(fused_wgrads, unfused_wgrads):
          self.assertAllClose(fused, unfused, rtol=1e-6, atol=1e-6)
  
  #### Benchmarking.
  
  
  class BenchmarkLSTMBlock(test.Benchmark):
  
    def benchmarkLSTMBlockCellFpropWithDynamicRNN(self):
      print("BlockLSTMCell forward propagation via dynamic_rnn().")
      print("--------------------------------------------------------------")
      print("LSTMBlockCell Seconds per inference.")
      print("batch_size,cell_size,input_size,time_steps,use_gpu,wall_time")
      iters = 10
      for config in benchmarking.dict_product({
          "batch_size": [1, 8, 13, 32, 67, 128],
          "cell_size": [128, 250, 512, 650, 1024, 1350],
          "time_steps": [40],
          "use_gpu": [True, False],
          "dtype": ["float32", "float16"],
      }):
        dtype = dtypes.float32 if config["dtype"] == "float32" else dtypes.float16
        with ops.Graph().as_default():
          with benchmarking.device(use_gpu=config["use_gpu"]):
            inputs = variable_scope.get_variable(
                "x",
                dtype=dtype,
                shape=[
                    config["time_steps"], config["batch_size"],
                    config["cell_size"]
                ])
            cell = lstm_ops.LSTMBlockCell(config["cell_size"], dtype=dtype)
            outputs = rnn.dynamic_rnn(cell, inputs, time_major=True, dtype=dtype)
            init_op = variables.global_variables_initializer()
  
          with session.Session() as sess:
            sess.run(init_op)
            wall_time = benchmarking.seconds_per_run(outputs, sess, iters)
  
          # Print to stdout. If the TEST_REPORT_FILE_PREFIX environment variable
          # is set, this will produce a copy-paste-able CSV file.
          print(",".join(
              map(str, [
                  config["dtype"], config["batch_size"], config["cell_size"],
                  config["cell_size"], config["time_steps"], config["use_gpu"],
                  wall_time
              ])))
          benchmark_name_template = "_".join([
              "LSTMBlockCell_fprop", "DT_%(dtype)s", "BS%(batch_size)i",
              "CS%(cell_size)i", "IS%(cell_size)i", "TS%(time_steps)i",
              "gpu_%(use_gpu)s"
          ])
  
          self.report_benchmark(
              name=benchmark_name_template % config,
              iters=iters,
              wall_time=wall_time,
              extras=config)
  
    def benchmarkLSTMBlockCellBpropWithDynamicRNN(self):
      print("BlockLSTMCell backward propagation via dynamic_rnn().")
      print("--------------------------------------------------------------")
      print("LSTMBlockCell Seconds per inference.")
      print("batch_size,cell_size,input_size,time_steps,use_gpu,wall_time")
      iters = 10
      for config in benchmarking.dict_product({
          "batch_size": [1, 8, 13, 32, 67, 128],
          "cell_size": [128, 250, 512, 650, 1024, 1350],
          "time_steps": [40],
          "use_gpu": [True, False],
          "dtype": ["float32", "float16"],
      }):
        dtype = dtypes.float32 if config["dtype"] == "float32" else dtypes.float16
        with ops.Graph().as_default():
          with benchmarking.device(use_gpu=config["use_gpu"]):
            time_steps = config["time_steps"]
            batch_size = config["batch_size"]
            cell_size = input_size = config["cell_size"]
            inputs = variable_scope.get_variable(
                "x", [time_steps, batch_size, cell_size],
                trainable=False,
                dtype=dtype)
            with variable_scope.variable_scope(
                "rnn", reuse=variable_scope.AUTO_REUSE):
              w = variable_scope.get_variable(
                  "rnn/lstm_cell/kernel",
                  shape=[input_size + cell_size, cell_size * 4],
                  dtype=dtype)
              b = variable_scope.get_variable(
                  "rnn/lstm_cell/bias",
                  shape=[cell_size * 4],
                  dtype=dtype,
                  initializer=init_ops.zeros_initializer())
              cell = lstm_ops.LSTMBlockCell(cell_size, dtype=dtype)
              outputs = rnn.dynamic_rnn(
                  cell, inputs, time_major=True, dtype=dtype)
            grads = gradients_impl.gradients(outputs, [inputs, w, b])
            init_op = variables.global_variables_initializer()
  
          with session.Session() as sess:
            sess.run(init_op)
            wall_time = benchmarking.seconds_per_run(grads, sess, iters)
  
          # Print to stdout. If the TEST_REPORT_FILE_PREFIX environment variable
          # is set, this will produce a copy-paste-able CSV file.
          print(",".join(
              map(str, [
                  config["dtype"], batch_size, cell_size, cell_size, time_steps,
                  config["use_gpu"], wall_time
              ])))
          benchmark_name_template = "_".join([
              "LSTMBlockCell_bprop", "DT_%(dtype)s", "BS%(batch_size)i",
              "CS%(cell_size)i", "IS%(cell_size)i", "TS%(time_steps)i",
              "gpu_%(use_gpu)s"
          ])
  
          self.report_benchmark(
              name=benchmark_name_template % config,
              iters=iters,
              wall_time=wall_time,
              extras=config)
  
  
  if __name__ == "__main__":
    test.main()
   # Copyright 2019 The TensorFlow Authors. All Rights Reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  # ==============================================================================
  """Training related logic for Keras model in TF 2.0 context.
  
  Note that all the code under this module is under active development, please DO
  NOT use it unless you are really sure what you are doing.
  """
  
  # pylint: disable=protected-access
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import collections
  import functools
  
  from tensorflow.python.distribute import distribution_strategy_context
  from tensorflow.python.eager import def_function
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import tensor_util
  from tensorflow.python.framework.ops import composite_tensor
  from tensorflow.python.keras import backend
  from tensorflow.python.keras.distribute import distributed_training_utils as dist_utils
  from tensorflow.python.keras.engine import training_eager
  from tensorflow.python.keras.engine import training_utils
  from tensorflow.python.keras.utils.mode_keys import ModeKeys
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import math_ops
  from tensorflow.python.util import nest
  
  
  def _get_or_make_execution_function(model, mode):
    """Makes or reuses function to run one step of distributed model execution."""
    model._init_distributed_function_cache_if_not_compiled()
  
    # Use a key with 'v2' to distinguish from fall-back execution functions.
    key = (mode, 'v2')
    distributed_function = dist_utils.get_distributed_function(model, key)
    if distributed_function:
      return distributed_function
  
    distribution_function = _make_execution_function(model, mode)
    dist_utils.set_distributed_function(model, key, distribution_function)
    return distribution_function
  
  
  def _make_execution_function(model, mode):
    """Creates a function to run one step of distributed model execution."""
    per_replica_function = _make_replica_execution_function(mode)
  
    def distributed_function(input_iterator):
      """A single step of the distributed execution across replicas."""
      x, y, sample_weights = _prepare_feed_values(
          model, input_iterator, mode)
      # Call `Model.{train,test,predict}_on_batch` on every replica passing
      # PerReplicas as arguments.  On every replica inside this call, each
      # PerReplica object will return the value for that replica.  The outputs
      # are PerReplicas too.
      strategy = distribution_strategy_context.get_strategy()
      outputs = strategy.experimental_run_v2(
          per_replica_function, args=(model, x, y, sample_weights))
      # Out of PerReplica outputs reduce or pick values to return.
      all_outputs = dist_utils.unwrap_output_dict(
          strategy, outputs, mode)
      return all_outputs
  
    if not model.run_eagerly:
      distributed_function = def_function.function(
          distributed_function, autograph=False)
  
    def execution_function(input_fn):
      # `numpy` translates Tensors to values in Eager mode.
      return nest.map_structure(_non_none_constant_value,
                                distributed_function(input_fn))
  
    return execution_function
  
  
  def _non_none_constant_value(v):
    constant_value = tensor_util.constant_value(v)
    return constant_value if constant_value is not None else v
  
  
  def _prepare_feed_values(model, inputs, mode):
    """Prepare feed values to the model execution function.
  
    Arguments:
      model: Model to prepare feed values for.
      inputs: An iterator of model inputs, targets, and sample_weights.
        model inputs may be lists, single values, or dicts mapping input feed
        names to values.
      mode: One of ModeKeys.TRAIN/ModeKeys.TEST/ModeKeys.PREDICT.
  
    Returns:
      Feed values for the model in the given mode. This is a tuple of
      the structure (inputs, targets, sample_weights), where each of
      (tuple, targets, sample_weights) may be a python list. Single values
      for inputs will always be wrapped in lists.
    """
    inputs, targets, sample_weights = _get_input_from_iterator(inputs)
  
    # When the inputs are dict, then we want to flatten it in the same order as
    # the input layers, such that the data are fed into the input layers in the
    # correct order.
    if isinstance(inputs, dict):
      inputs = [inputs[key] for key in model._feed_input_names]
    else:
      inputs = training_utils.ModelInputs(inputs).as_list()
  
    if mode == ModeKeys.PREDICT:
      sample_weights = []
      targets = []
  
    ins = [inputs, targets, sample_weights]
    return tuple(ins)
  
  
  def _get_input_from_iterator(iterator):
    """Get elements from the iterator and verify the input shape and type."""
    next_element = next(iterator)
  
    if (tensor_util.is_tensor(next_element) or
        isinstance(next_element, (dict, composite_tensor.CompositeTensor))):
      next_element = [next_element]
    if len(next_element) == 1:
      x, = next_element
      y = None
      sample_weights = None
    elif len(next_element) == 2:
      x, y = next_element
      sample_weights = None
    else:
      x, y, sample_weights = next_element
  
    # Validate that all the elements in x and y are of the same type and shape.
    dist_utils.validate_distributed_dataset_inputs(
        distribution_strategy_context.get_strategy(), x, y, sample_weights)
    return x, y, sample_weights
  
  
  def _make_replica_execution_function(mode):
    """A single step of the distributed execution on a replica."""
    if mode == ModeKeys.TRAIN:
      func = train_on_batch
    elif mode == ModeKeys.TEST:
      func = test_on_batch
    else:
      def _predict_on_batch(model, x, y=None, sample_weights=None):
        del y, sample_weights
        return predict_on_batch(model, x)
  
      func = _predict_on_batch
  
    if mode != ModeKeys.PREDICT:
      # `reset_metrics` is set to False to maintain stateful metrics across
      # batch-level calls.
      func = functools.partial(func, reset_metrics=False)
  
    return func
  
  
  def _prepare_model_with_inputs(model, dataset):
    """Use the data from the adapter to config the model.
  
    Model need to be properly configured before training, eg build with inputs, or
    compile with inputs for subclass model.
  
    Args:
      model: a Keras model object.
      dataset: a eager dataset instance where the data will be extracted.
    """
    if not model.inputs:
      inputs, target, _ = model._build_model_with_inputs(dataset, targets=None)
    else:
      inputs, target, _ = _get_input_from_iterator(iter(dataset))
  
    if not model._is_compiled and model.optimizer:
      model._compile_from_inputs(inputs, target, dataset, None)
  
    if target is not None:
      training_utils.prepare_sample_weight_modes(model._training_endpoints,
                                                 model.sample_weight_mode)
  
  
  def train_on_batch(
      model,
      x,
      y=None,
      sample_weight=None,
      class_weight=None,
      reset_metrics=True):
    """Runs a single gradient update on a single batch of data.
  
    Arguments:
        model: The model to train.
        x: Input data. It could be:
          - A Numpy array (or array-like), or a list of arrays
              (in case the model has multiple inputs).
          - A TensorFlow tensor, or a list of tensors
              (in case the model has multiple inputs).
          - A dict mapping input names to the corresponding array/tensors,
              if the model has named inputs.
          - A `tf.data` dataset.
        y: Target data. Like the input data `x`, it could be either Numpy
          array(s) or TensorFlow tensor(s). It should be consistent with `x`
          (you cannot have Numpy inputs and tensor targets, or inversely). If
          `x` is a dataset `y` should not be specified
          (since targets will be obtained from the iterator).
        sample_weight: Optional array of the same length as x, containing
          weights to apply to the model's loss for each sample. In the case of
          temporal data, you can pass a 2D array with shape (samples,
          sequence_length), to apply a different weight to every timestep of
          every sample. In this case you should make sure to specify
          sample_weight_mode="temporal" in compile(). This argument is not
          supported when `x` is a dataset.
        class_weight: Optional dictionary mapping class indices (integers) to a
          weight (float) to apply to the model's loss for the samples from this
          class during training. This can be useful to tell the model to "pay
          more attention" to samples from an under-represented class.
        reset_metrics: If `True`, the metrics returned will be only for this
          batch. If `False`, the metrics will be statefully accumulated across
          batches.
  
    Returns:
        Scalar training loss
        (if the model has a single output and no metrics)
        or list of scalars (if the model has multiple outputs
        and/or metrics). The attribute `model.metrics_names` will give you
        the display labels for the scalar outputs.
  
    Raises:
      ValueError: In case of invalid user-provided arguments.
    """
    model._assert_compile_was_called()
  
    # TODO(scottzhu): Standardization should happen in the data handlers,
    ## not on a per batch basis in the *_on_batch methods
    # Validate and standardize user data.
    x, y, sample_weights = model._standardize_user_data(
        x, y, sample_weight=sample_weight, class_weight=class_weight,
        extract_tensors_from_dataset=True)
    batch_size = array_ops.shape(nest.flatten(x, expand_composites=True)[0])[0]
    # If `model._distribution_strategy` is True, then we are in a replica context
    # at this point because of the check above.  `train_on_batch` is being run
    # for each replica by `model._distribution_strategy` and the same code path
    # as Eager is expected to be taken.
    outputs = training_eager.train_on_batch(
        model,
        x,
        y,
        sample_weights=sample_weights,
        output_loss_metrics=model._output_loss_metrics)
  
    if reset_metrics:
      model.reset_metrics()
  
    outputs['batch_size'] = math_ops.cast(batch_size, dtypes.int64)
    return outputs
  
  
  def test_on_batch(model, x, y=None, sample_weight=None, reset_metrics=True):
    """Test the model on a single batch of samples.
  
    Arguments:
        model: The model to test.
        x: Input data. It could be:
          - A Numpy array (or array-like), or a list of arrays
            (in case the model has multiple inputs).
          - A TensorFlow tensor, or a list of tensors
            (in case the model has multiple inputs).
          - A dict mapping input names to the corresponding array/tensors,
            if the model has named inputs.
          - A `tf.data` dataset.
        y: Target data. Like the input data `x`,
          it could be either Numpy array(s) or TensorFlow tensor(s).
          It should be consistent with `x` (you cannot have Numpy inputs and
          tensor targets, or inversely). If `x` is a dataset,
          `y` should not be specified
          (since targets will be obtained from the iterator).
        sample_weight: Optional array of the same length as x, containing
            weights to apply to the model's loss for each sample.
            In the case of temporal data, you can pass a 2D array
            with shape (samples, sequence_length),
            to apply a different weight to every timestep of every sample.
            In this case you should make sure to specify
            sample_weight_mode="temporal" in compile(). This argument is not
            supported when `x` is a dataset.
        reset_metrics: If `True`, the metrics returned will be only for this
          batch. If `False`, the metrics will be statefully accumulated across
          batches.
  
    Returns:
        Scalar test loss (if the model has a single output and no metrics)
        or list of scalars (if the model has multiple outputs
        and/or metrics). The attribute `model.metrics_names` will give you
        the display labels for the scalar outputs.
  
    Raises:
        ValueError: In case of invalid user-provided arguments.
    """
    model._assert_compile_was_called()
  
    # TODO(scottzhu): Standardization should happen in the data handlers,
    ## not on a per batch basis in the *_on_batch methods
    # Validate and standardize user data.
    x, y, sample_weights = model._standardize_user_data(
        x, y, sample_weight=sample_weight, extract_tensors_from_dataset=True)
  
    batch_size = array_ops.shape(nest.flatten(x, expand_composites=True)[0])[0]
    outputs = training_eager.test_on_batch(
        model,
        x,
        y,
        sample_weights=sample_weights,
        output_loss_metrics=model._output_loss_metrics)
  
    if reset_metrics:
      model.reset_metrics()
  
    outputs['batch_size'] = math_ops.cast(batch_size, dtypes.int64)
    return outputs
  
  
  def predict_on_batch(model, x):
    """Returns predictions for a single batch of samples.
  
    Arguments:
        model: The model to predict with.
        x: Input data. It could be:
          - A Numpy array (or array-like), or a list of arrays
            (in case the model has multiple inputs).
          - A TensorFlow tensor, or a list of tensors
            (in case the model has multiple inputs).
          - A `tf.data` dataset.
  
    Returns:
        Numpy array(s) of predictions.
  
    Raises:
        ValueError: In case of mismatch between given number of inputs and
          expectations of the model.
    """
    # TODO(scottzhu): Standardization should happen in the data handlers,
    ## not on a per batch basis in the *_on_batch methods
    # Validate and standardize user data.
    inputs, _, _ = model._standardize_user_data(
        x, extract_tensors_from_dataset=True)
  
    # If `model._distribution_strategy` is True, then we are in a replica context
    # at this point.
    inputs = training_utils.cast_if_floating_dtype(inputs)
    if isinstance(inputs, collections.Sequence):
      # Unwrap lists with only one input, as we do when training on batch
      if len(inputs) == 1:
        inputs = inputs[0]
  
    with backend.eager_learning_phase_scope(0):
      return model(inputs)  # pylint: disable=not-callable
   # Copyright 2015 The TensorFlow Authors. All Rights Reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  # ==============================================================================
  """tensorboard_logging provides logging that is also written to the events file.
  
  Any messages logged via this module will be logged both via the platform logging
  mechanism and to the SummaryWriter set via `set_summary_writer`. This is useful
  for logging messages that you might want to be visible from inside TensorBoard
  or that should be permanently associated with the training session.
  
  You can use this just like the logging module:
  
  >>> tensorboard_logging.set_summary_writer(summary_writer)
  >>> tensorboard_logging.info("my %s", "message")
  >>> tensorboard_logging.log(tensorboard_logging.WARN, "something")
  """
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import time
  
  from tensorflow.core.util import event_pb2
  from tensorflow.python.platform import tf_logging as logging
  
  DEBUG = 'DEBUG'
  INFO = 'INFO'
  WARN = 'WARN'
  ERROR = 'ERROR'
  FATAL = 'FATAL'
  
  # Messages with levels below this verbosity will not be logged.
  _verbosity = WARN
  
  # A value meaning 'not set yet' so we can use None to mean 'user actively told
  # us they don't want a SummaryWriter'.
  _sentinel_summary_writer = object()
  
  # The SummaryWriter instance to use when logging, or None to not log, or
  # _sentinel_summary_writer to indicate that the user hasn't called
  # set_summary_writer yet.
  _summary_writer = _sentinel_summary_writer
  
  # Map from the tensorboard_logging logging enum values to the proto's enum
  # values.
  _LEVEL_PROTO_MAP = {
      DEBUG: event_pb2.LogMessage.DEBUGGING,
      INFO: event_pb2.LogMessage.INFO,
      WARN: event_pb2.LogMessage.WARN,
      ERROR: event_pb2.LogMessage.ERROR,
      FATAL: event_pb2.LogMessage.FATAL,
  }
  
  # Map from the tensorboard_logging module levels to the logging module levels.
  _PLATFORM_LOGGING_LEVEL_MAP = {
      DEBUG: logging.DEBUG,
      INFO: logging.INFO,
      WARN: logging.WARN,
      ERROR: logging.ERROR,
      FATAL: logging.FATAL
  }
  
  
  def get_verbosity():
    return _verbosity
  
  
  def set_verbosity(verbosity):
    _check_verbosity(verbosity)
    global _verbosity
    _verbosity = verbosity
  
  
  def _check_verbosity(verbosity):
    if verbosity not in _LEVEL_PROTO_MAP:
      raise ValueError('Level %s is not a valid tensorboard_logging level' %
                       verbosity)
  
  
  def set_summary_writer(summary_writer):
    """Sets the summary writer that events will be logged to.
  
    Calling any logging methods inside this module without calling this method
    will fail. If you don't want to log, call `set_summary_writer(None)`.
  
    Args:
      summary_writer: Either a SummaryWriter or None. None will cause messages not
      to be logged to any SummaryWriter, but they will still be passed to the
      platform logging module.
    """
    global _summary_writer
    _summary_writer = summary_writer
  
  
  def _clear_summary_writer():
    """Makes all subsequent log invocations error.
  
    This is only used for testing. If you want to disable TensorBoard logging,
    call `set_summary_writer(None)` instead.
    """
    global _summary_writer
    _summary_writer = _sentinel_summary_writer
  
  
  def log(level, message, *args):
    """Conditionally logs `message % args` at the level `level`.
  
    Note that tensorboard_logging verbosity and logging verbosity are separate;
    the message will always be passed through to the logging module regardless of
    whether it passes the tensorboard_logging verbosity check.
  
    Args:
      level: The verbosity level to use. Must be one of
        tensorboard_logging.{DEBUG, INFO, WARN, ERROR, FATAL}.
      message: The message template to use.
      *args: Arguments to interpolate to the message template, if any.
  
    Raises:
      ValueError: If `level` is not a valid logging level.
      RuntimeError: If the `SummaryWriter` to use has not been set.
    """
    if _summary_writer is _sentinel_summary_writer:
      raise RuntimeError('Must call set_summary_writer before doing any '
                         'logging from tensorboard_logging')
    _check_verbosity(level)
    proto_level = _LEVEL_PROTO_MAP[level]
    if proto_level >= _LEVEL_PROTO_MAP[_verbosity]:
      log_message = event_pb2.LogMessage(level=proto_level,
                                         message=message % args)
      event = event_pb2.Event(wall_time=time.time(), log_message=log_message)
  
      if _summary_writer:
        _summary_writer.add_event(event)
  
    logging.log(_PLATFORM_LOGGING_LEVEL_MAP[level], message, *args)
  
  
  def debug(message, *args):
    log(DEBUG, message, *args)
  
  
  def info(message, *args):
    log(INFO, message, *args)
  
  
  def warn(message, *args):
    log(WARN, message, *args)
  
  
  def error(message, *args):
    log(ERROR, message, *args)
  
  
  def fatal(message, *args):
    log(FATAL, message, *args)
  
apis: tensorflow.python.ops.array_ops.pad tensorflow.python.ops.variable_scope.variable_scope tensorflow.contrib.timeseries.python.timeseries.math_utils.variable_covariance_matrix tensorflow.python.framework.tensor_shape.dimension_value tensorflow.python.ops.init_ops.zeros_initializer tensorflow.python.ops.math_ops.reduce_mean tensorflow.python.ops.array_ops.concat tensorflow.contrib.timeseries.python.timeseries.state_space_models.state_space_model.StateSpaceModelConfiguration tensorflow.python.ops.linalg_ops.eye
apis: numpy.ones numpy.fft.fftshift tensorflow.python.ops.spectral_ops_test_util.fft_kernel_label_map tensorflow.python.platform.test.is_gpu_available numpy.fft.rfft2 numpy.random.rand tensorflow.python.ops.signal.fft_ops.fftshift tensorflow.python.ops.signal.fft_ops.ifftshift tensorflow.python.framework.dtypes.as_dtype numpy.random.uniform numpy.fft.fft2 numpy.zeros numpy.power numpy.prod tensorflow.python.framework.ops.convert_to_tensor numpy.fft.irfft2 tensorflow.core.protobuf.config_pb2.ConfigProto numpy.fft.ifft2 tensorflow.python.ops.math_ops.complex numpy.fft.ifftshift tensorflow.python.platform.test.main tensorflow.python.ops.math_ops.conj
apis: tensorflow.python.data.experimental.ops.optimization.assert_next tensorflow.python.platform.test.main tensorflow.python.data.ops.dataset_ops.Options tensorflow.python.data.experimental.ops.stats_aggregator.StatsAggregator tensorflow.python.data.ops.dataset_ops.Dataset.from_tensors
apis: tensorflow.python.util.deprecation.silence
apis: numpy.eye numpy.transpose numpy.linalg.inv tensorflow.python.ops.math_ops.matmul numpy.linalg.det numpy.random.RandomState tensorflow.python.platform.test.main numpy.arange tensorflow.python.ops.array_ops.matrix_band_part scipy.stats.multivariate_normal
apis: tensorflow.python.keras.initializers.ConstantV2 tensorflow.python.keras.initializers.lecun_normalV2 tensorflow.python.keras.backend.get_value tensorflow.python.keras.initializers.TruncatedNormalV2 tensorflow.python.keras.initializers.OrthogonalV2 tensorflow.python.keras.initializers.he_uniformV2 tensorflow.python.keras.initializers.GlorotUniformV2 tensorflow.python.ops.init_ops._compute_fans tensorflow.python.keras.Model tensorflow.python.keras.initializers.RandomUniformV2 tensorflow.python.keras.Input tensorflow.python.keras.initializers.OnesV2 tensorflow.python.tf2.enable tensorflow.python.keras.initializers.RandomNormalV2 tensorflow.python.ops.array_ops.ones tensorflow.python.keras.initializers.ZerosV2 tensorflow.python.keras.initializers.deserialize tensorflow.python.keras.layers.Dense tensorflow.python.keras.initializers.IdentityV2 tensorflow.python.keras.initializers.GlorotNormalV2 tensorflow.python.keras.initializers.he_normalV2 tensorflow.python.platform.test.main tensorflow.python.keras.initializers.lecun_uniformV2 numpy.sqrt tensorflow.python.keras.initializers.get
apis: numpy.ones tensorflow.contrib.rnn.python.kernel_tests.benchmarking.seconds_per_run tensorflow.python.ops.variable_scope.variable_scope tensorflow.python.ops.gradients_impl.gradients tensorflow.contrib.rnn.python.ops.lstm_ops.LSTMBlockFusedCell tensorflow.python.ops.array_ops.expand_dims tensorflow.python.ops.rnn.static_rnn tensorflow.python.framework.constant_op.constant tensorflow.contrib.rnn.python.kernel_tests.benchmarking.dict_product tensorflow.python.ops.init_ops.random_uniform_initializer tensorflow.python.framework.ops.Graph tensorflow.python.ops.gen_bitwise_ops.bitwise_and tensorflow.python.ops.variable_scope.get_variable tensorflow.python.ops.array_ops.placeholder tensorflow.python.ops.init_ops.zeros_initializer tensorflow.python.client.session.Session tensorflow.python.ops.init_ops.constant_initializer tensorflow.python.ops.rnn.dynamic_rnn tensorflow.contrib.rnn.python.ops.lstm_ops.LSTMBlockCell tensorflow.python.ops.gen_array_ops.bitcast tensorflow.python.ops.array_ops.zeros tensorflow.python.ops.variables.global_variables_initializer tensorflow.python.ops.variables.trainable_variables tensorflow.contrib.rnn.python.kernel_tests.benchmarking.device tensorflow.python.framework.ops.convert_to_tensor tensorflow.python.ops.variable_scope.get_variable_scope numpy.random.randn tensorflow.python.platform.test.main tensorflow.python.ops.array_ops.stack tensorflow.python.ops.rnn_cell.LSTMCell tensorflow.python.ops.rnn_cell.BasicLSTMCell numpy.array
apis: tensorflow.python.keras.engine.training_eager.test_on_batch tensorflow.python.keras.engine.training_utils.prepare_sample_weight_modes tensorflow.python.keras.distribute.distributed_training_utils.set_distributed_function tensorflow.python.keras.engine.training_eager.train_on_batch tensorflow.python.distribute.distribution_strategy_context.get_strategy tensorflow.python.util.nest.flatten tensorflow.python.ops.math_ops.cast tensorflow.python.keras.engine.training_utils.cast_if_floating_dtype tensorflow.python.keras.distribute.distributed_training_utils.unwrap_output_dict tensorflow.python.keras.backend.eager_learning_phase_scope tensorflow.python.eager.def_function.function tensorflow.python.keras.distribute.distributed_training_utils.get_distributed_function tensorflow.python.keras.engine.training_utils.ModelInputs tensorflow.python.framework.tensor_util.constant_value tensorflow.python.framework.tensor_util.is_tensor
apis: tensorflow.core.util.event_pb2.LogMessage tensorflow.python.platform.tf_logging.log