Python tensorflow.compat.v2.float32() Examples

def _check_config(self, metric_cls, init_args):
    """Checks if Keras metrics can be saved and restored.

    Args:
      metric_cls: Keras metric class.
      init_args: (dict) Initializer keyword arguments.
    """
    init_args.update({
        'name': 'my_metric',
        'dtype': tf.float32,
    })
    metric_obj = metric_cls(**init_args)
    config = metric_obj.get_config()
    self.assertIsNotNone(config)

    restored_metric_obj = metric_cls.from_config(config)
    for init_name, init_value in six.iteritems(init_args):
      self.assertEqual(init_value, getattr(restored_metric_obj,
                                           '_' + init_name)) 

def bernoulli(key, mean=np.float32(0.5), shape=None):
  """Sample Bernoulli random values with given shape and mean.

  Args:
    key: the RNG key.
    mean: optional, an array_like broadcastable to `shape` for the mean of the
      random variables (default 0.5).
    shape: optional, a tuple of nonnegative integers representing the shape
      (default to `mean`'s shape).

  Returns:
    A random array with the specified shape and boolean dtype.
  """
  mean = tf_np.asarray(mean)
  if shape is None:
    shape = mean.shape
  return uniform(key, shape) < mean 

def true_divide(x1, x2):
  def _avoid_float64(x1, x2):
    if x1.dtype == x2.dtype and x1.dtype in (tf.int32, tf.int64):
      x1 = tf.cast(x1, dtype=tf.float32)
      x2 = tf.cast(x2, dtype=tf.float32)
    return x1, x2

  def f(x1, x2):
    if x1.dtype == tf.bool:
      assert x2.dtype == tf.bool
      float_ = dtypes.default_float_type()
      x1 = tf.cast(x1, float_)
      x2 = tf.cast(x2, float_)
    if not dtypes.is_allow_float64():
      # tf.math.truediv in Python3 produces float64 when both inputs are int32
      # or int64. We want to avoid that when is_allow_float64() is False.
      x1, x2 = _avoid_float64(x1, x2)
    return tf.math.truediv(x1, x2)
  return _bin_op(f, x1, x2) 

def test_multiple_state_vars(self):
    x = tf.constant([3.0, 4.0])
    y = tf.constant([5.0, 6.0])
    z = tf.constant([7.0, 8.0])
    alpha = tf.constant(2.0)
    beta = tf.constant(1.0)

    with tf.GradientTape(persistent=True) as tape:
      tape.watch([alpha, beta])
      def body(i, state):
        x, y, z = state
        k = tf.cast(i + 1, tf.float32)
        return [x * alpha - beta, y * k * alpha * beta, z * beta + x]
      out = for_loop(body, [x, y, z], [alpha, beta], 3)

    with self.subTest("independent_vars"):
      grad = tape.gradient(out[1], alpha)
      self.assertAllEqual(792, grad)
    with self.subTest("dependent_vars"):
      grad = tape.gradient(out[2], beta)
      self.assertAllEqual(63, grad) 

def _testBinOp(self, a, b, out, f, types=None):
    a = t2a(tf.convert_to_tensor(value=a, dtype=np.int32))
    b = t2a(tf.convert_to_tensor(value=b, dtype=np.int32))
    if not isinstance(out, arrays.ndarray):
      out = t2a(tf.convert_to_tensor(value=out, dtype=np.int32))
    if types is None:
      types = [[np.int32, np.int32, np.int32],
               [np.int64, np.int32, np.int64],
               [np.int32, np.int64, np.int64],
               [np.float32, np.int32, np.float64],
               [np.int32, np.float32, np.float64],
               [np.float32, np.float32, np.float32],
               [np.float64, np.float32, np.float64],
               [np.float32, np.float64, np.float64]]
    for a_type, b_type, out_type in types:
      o = f(a.astype(a_type), b.astype(b_type))
      self.assertIs(o.dtype.type, out_type)
      self.assertAllEqual(out.astype(out_type), o) 

def __init__(self,
                 shift,
                 validate_args=False,
                 name='shift'):
        """Instantiates the `Shift` bijector which computes `Y = g(X; shift) = X + shift`
        where `shift` is a numeric `Tensor`.
        Args:
          shift: Floating-point `Tensor`.
          validate_args: Python `bool` indicating whether arguments should be
            checked for correctness.
          name: Python `str` name given to ops managed by this object.
        """
        with tf.name_scope(name) as name:
            dtype = dtype_util.common_dtype([shift], dtype_hint=tf.float32)
            self._shift = tensor_util.convert_nonref_to_tensor(shift, dtype=dtype, name='shift')
            super(Shift, self).__init__(
              forward_min_event_ndims=0,
              is_constant_jacobian=True,
              dtype=dtype,
              validate_args=validate_args,
              name=name
            ) 

def _base_expansion_size(num, bases):
  """Computes the number of terms in the place value expansion.

  Let num = a0 + a1 b + a2 b^2 + ... ak b^k be the place value expansion of
  `num` in base b (ak <> 0). This function computes and returns `k+1` for each
  base `b` specified in `bases`.

  This can be inferred from the base `b` logarithm of `num` as follows:
    $$k = Floor(log_b (num)) + 1  = Floor( log(num) / log(b)) + 1$$

  Args:
    num: Scalar numpy array of dtype either `float32` or `float64`. The number
      to compute the base expansion size of.
    bases: Numpy array of the same dtype as num. The bases to compute the size
      against.

  Returns:
    Tensor of same dtype and shape as `bases` containing the size of num when
    written in that base.
  """
  return np.floor(np.log(num) / np.log(bases)) + 1


# First 1000 prime numbers. 

def testOutputIsPermutation(self):
    """Checks that stateless_random_shuffle outputs a permutation."""
    for dtype in (tf.int32, tf.int64, tf.float32, tf.float64):
      identity_permutation = tf.range(10, dtype=dtype)
      random_shuffle_seed_1 = tff_rnd.stateless_random_shuffle(
          identity_permutation, seed=tf.constant((1, 42), tf.int64))
      random_shuffle_seed_2 = tff_rnd.stateless_random_shuffle(
          identity_permutation, seed=tf.constant((2, 42), tf.int64))
      # Check that the shuffles are of the correct dtype
      for shuffle in (random_shuffle_seed_1, random_shuffle_seed_2):
        np.testing.assert_equal(shuffle.dtype, dtype.as_numpy_dtype)
      random_shuffle_seed_1 = self.evaluate(random_shuffle_seed_1)
      random_shuffle_seed_2 = self.evaluate(random_shuffle_seed_2)
      identity_permutation = self.evaluate(identity_permutation)
      # Check that the shuffles are different
      self.assertTrue(
          np.abs(random_shuffle_seed_1 - random_shuffle_seed_2).max())
      # Check that the shuffles are indeed permutations
      for shuffle in (random_shuffle_seed_1, random_shuffle_seed_2):
        self.assertAllEqual(set(shuffle), set(identity_permutation)) 

def testOutputIsStatelessSession(self):
    """Checks that stateless_random_shuffle is stateless across Sessions."""
    random_permutation_next_call = None
    for dtype in (tf.int32, tf.int64, tf.float32, tf.float64):
      random_permutation = tff_rnd.stateless_random_shuffle(
          tf.range(10, dtype=dtype), seed=tf.constant((100, 42), tf.int64))
      with tf.compat.v1.Session() as sess:
        random_permutation_first_call = sess.run(random_permutation)
      if random_permutation_next_call is not None:
        # Checks that the values are the same across different dtypes
        np.testing.assert_array_equal(random_permutation_first_call,
                                      random_permutation_next_call)
      with tf.compat.v1.Session() as sess:
        random_permutation_next_call = sess.run(random_permutation)
      np.testing.assert_array_equal(random_permutation_first_call,
                                    random_permutation_next_call) 

def testMultiDimensionalShape(self):
    """Check that stateless_random_shuffle works with multi-dim shapes."""
    for dtype in (tf.int32, tf.int64, tf.float32, tf.float64):
      input_permutation = tf.constant([[[1], [2], [3]], [[4], [5], [6]]],
                                      dtype=dtype)
      random_shuffle = tff_rnd.stateless_random_shuffle(
          input_permutation, seed=(1, 42))
      random_permutation_first_call = self.evaluate(random_shuffle)
      random_permutation_next_call = self.evaluate(random_shuffle)
      input_permutation = self.evaluate(input_permutation)
      # Check that the dtype is correct
      np.testing.assert_equal(random_permutation_first_call.dtype,
                              dtype.as_numpy_dtype)
      # Check that the shuffles are the same
      np.testing.assert_array_equal(random_permutation_first_call,
                                    random_permutation_next_call)
      # Check that the output shape is correct
      np.testing.assert_equal(random_permutation_first_call.shape,
                              input_permutation.shape) 

def test_forward_unconnected_gradient(self):
    t = tf.range(1, 3, dtype=tf.float32)  # Shape [2]
    zeros = tf.zeros([2], dtype=t.dtype)
    func = lambda t: tf.stack([zeros, zeros, zeros], axis=0)  # Shape [3, 2]
    expected_result = [[0.0, 0.0], [0.0, 0.0], [0.0, 0.0]]
    with self.subTest("EagerExecution"):
      fwd_grad = self.evaluate(tff.math.fwd_gradient(
          func, t, unconnected_gradients=tf.UnconnectedGradients.ZERO))
      self.assertEqual(fwd_grad.shape, (3, 2))
      np.testing.assert_allclose(fwd_grad, expected_result)
    with self.subTest("GraphExecution"):
      @tf.function
      def grad_computation():
        y = func(t)
        return tff.math.fwd_gradient(
            y, t, unconnected_gradients=tf.UnconnectedGradients.ZERO)
      fwd_grad = self.evaluate(grad_computation())
      self.assertEqual(fwd_grad.shape, (3, 2))
      np.testing.assert_allclose(fwd_grad, expected_result) 

def compute_logits(self, token_ids: tf.Tensor, training: bool) -> tf.Tensor:
        """
        Implements a language model, where each output is conditional on the current
        input and inputs processed so far.

        Args:
            token_ids: int32 tensor of shape [B, T], storing integer IDs of tokens.
            training: Flag indicating if we are currently training (used to toggle dropout)

        Returns:
            tf.float32 tensor of shape [B, T, V], storing the distribution over output symbols
            for each timestep for each batch element.
        """
        # TODO 5# 1) Embed tokens
        # TODO 5# 2) Run RNN on embedded tokens
        # TODO 5# 3) Project RNN outputs onto the vocabulary to obtain logits.
        return rnn_output_logits 

def test_batching(self):
    x = tf.constant([[3.0, 4.0], [30.0, 40.0]])
    y = tf.constant([[5.0, 6.0], [50.0, 60.0]])
    z = tf.constant([[7.0, 8.0], [70.0, 80.0]])
    alpha = tf.constant(2.0)
    beta = tf.constant(1.0)

    with tf.GradientTape(persistent=True) as tape:
      tape.watch([alpha, beta])
      def body(i, state):
        x, y, z = state
        k = tf.cast(i + 1, tf.float32)
        return [x * alpha - beta, y * k * alpha * beta, z * beta + x]
      out = for_loop(body, [x, y, z], [alpha, beta], 3)
    with self.subTest("independent_vars"):
      grad = tape.gradient(out[1], alpha)
      self.assertAllEqual(8712, grad)
    with self.subTest("dependent_vars"):
      grad = tape.gradient(out[2], beta)
      self.assertAllEqual(783, grad) 

def test_douglas_step_2d(self):
    u = np.arange(1, 17, dtype=np.float32).reshape(4, 4)
    d = np.arange(11, 27, dtype=np.float32).reshape(4, 4)
    dx = np.array([d, -3 * d, 2 * d])
    dy = np.array([2 * d, -6 * d, 4 * d])
    dxy = np.arange(-8, 8, dtype=np.float32).reshape(4, 4)
    bx = np.arange(2, 18, dtype=np.float32).reshape(4, 4)
    by = np.arange(5, 21, dtype=np.float32).reshape(4, 4)
    theta = 0.3

    def equation_params_fn(t):
      del t
      return ([[_tfconst(dy), _spread_mixed_term(_tfconst(dxy))],
               [None, _tfconst(dx)]],
              [_tfconst(by), _tfconst(bx)])

    scheme = douglas_adi_scheme(theta=theta)
    actual = self.evaluate(
        scheme(value_grid=tf.constant(u, dtype=tf.float32), t1=0, t2=1,
               equation_params_fn=equation_params_fn,
               n_dims=2))
    expected = self._simplified_douglas_step_2d(u, dx, dy, dxy, bx, by,
                                                0, 1, theta)
    self.assertLess(np.max(np.abs(expected - actual)), 0.01) 

def test_dynamic_shapes(self):
    """Can build op with dynamic shapes in graph mode."""
    if tf.executing_eagerly():
      return
    minimum = np.array([1.0, 1.0])
    scales = np.array([2.0, 3.0])

    @tff.math.make_val_and_grad_fn
    def quadratic(x):
      return tf.reduce_sum(input_tensor=scales * (x - minimum)**2)

    # Test with a vector of unknown dimension.
    start = tf.compat.v1.placeholder(tf.float32, shape=[None])
    op = tff.math.optimizer.conjugate_gradient_minimize(
        quadratic, initial_position=start, tolerance=1e-8)
    self.assertFalse(op.position.shape.is_fully_defined())

    with self.cached_session() as session:
      results = session.run(op, feed_dict={start: [0.6, 0.8]})
    self.assertTrue(results.converged)
    self.assertLessEqual(_norm(results.objective_gradient), 1e-8)
    self.assertArrayNear(results.position, minimum, 1e-5) 

def test_sample_paths_dtypes(self):
    """Sampled paths have the expected dtypes."""
    for dtype in [np.float32, np.float64]:
      drift_fn = lambda t, x: tf.sqrt(t) * tf.ones_like(x, dtype=t.dtype)
      vol_fn = lambda t, x: t * tf.ones([1, 1], dtype=t.dtype)
      process = GenericItoProcess(
          dim=1, drift_fn=drift_fn, volatility_fn=vol_fn, dtype=dtype)

      paths = self.evaluate(
          process.sample_paths(
              times=[0.1, 0.2],
              num_samples=10,
              initial_state=[0.1],
              time_step=0.01,
              seed=123))
      self.assertEqual(paths.dtype, dtype)

  # Several tests below are unit tests for GenericItoProcess.fd_solver_backward:
  # they mock out the pde solver and check only the conversion of SDE to PDE,
  # but not PDE solving. There are also integration tests further below. 

def test_douglas_step_3d(self):
    u = np.arange(0, 80, dtype=np.float32).reshape(4, 4, 5)
    d = np.arange(10, 90, dtype=np.float32).reshape(4, 4, 5)
    dx = np.array([d, -3 * d, 2 * d])
    dy = 2 * dx
    dz = 3 * dx
    dxy = np.arange(-20, 60, dtype=np.float32).reshape(4, 4, 5)
    dyz = 2 * dxy
    dxz = 3 * dxy
    bx = np.arange(20, 100, dtype=np.float32).reshape(4, 4, 5)
    by = np.arange(30, 110, dtype=np.float32).reshape(4, 4, 5)
    bz = np.arange(40, 120, dtype=np.float32).reshape(4, 4, 5)
    theta = 0.3

    def equation_params_fn(t):
      del t
      dyz_spread = _spread_mixed_term(_tfconst(dyz))
      dxz_spread = _spread_mixed_term(_tfconst(dxz))
      dxy_spread = _spread_mixed_term(_tfconst(dxy))
      return ([[_tfconst(dz), dyz_spread, dxz_spread],
               [dyz_spread, _tfconst(dy), dxy_spread],
               [dxz_spread, dxy_spread, _tfconst(dx)]],
              [_tfconst(bz), _tfconst(by), _tfconst(bx)])

    scheme = douglas_adi_scheme(theta=theta)
    actual = self.evaluate(
        scheme(value_grid=tf.constant(u, dtype=tf.float32), t1=0, t2=1,
               equation_params_fn=equation_params_fn, n_dims=3))
    expected = self._simplified_douglas_step_3d(u, dx, dy, dz, dxy, dyz, dxz,
                                                bx, by, bz, 0, 1, theta)
    self.assertLess(np.max(np.abs(expected - actual)), 0.01) 

def _convert_to_trimesh_feature(self, obj):
    if isinstance(obj, triangle_mesh.Trimesh):
      vertices = np.array(obj.vertices)
      faces = np.array(obj.faces, dtype=np.uint64)
    elif isinstance(obj, triangle_mesh.Scene):
      # Concatenate all the vertices and faces of the triangle meshes in the
      # scene.
      # TODO(b/156117488): Change to a different merging algorithm to avoid
      # duplicated vertices.
      vertices_list = [
          np.array(mesh.vertices) for mesh in obj.geometry.values()
      ]
      faces_list = np.array([
          np.array(mesh.faces, dtype=np.uint64)
          for mesh in obj.geometry.values()
      ])
      faces_offset = np.cumsum(
          [vertices.shape[0] for vertices in vertices_list], dtype=np.uint64)
      faces_list[1:] += faces_offset[:-1]
      vertices = np.concatenate(vertices_list, axis=0)
      faces = np.concatenate(faces_list, axis=0)
    else:
      raise ValueError('obj should be either a Trimesh or a Scene')
    return {
        'vertices': vertices.astype(np.float32),
        'faces': faces,
    } 

def _np_tridiagonal_solve(self, diags, rhs):
    return self.evaluate(
        tf.linalg.tridiagonal_solve(
            tf.constant(diags, dtype=tf.float32),
            tf.constant(rhs, dtype=tf.float32))) 

def testAstype(self):
    a = t2a(tf.convert_to_tensor(value=1.1, dtype=tf.float32)).astype(np.int32)
    self.assertIs(a.dtype.type, np.int32)
    self.assertAllEqual(1, a)
    a = t2a(tf.convert_to_tensor(value=[0.0, 1.1],
                                 dtype=tf.float32)).astype(np.bool_)
    self.assertIs(a.dtype.type, np.bool_)
    self.assertAllEqual([False, True], a) 

def testHeatEquation_WithMixedBoundaryConditions(self):
    """Test for mixed boundary conditions.

    Tests solving heat equation with the following boundary conditions:
    `u(x, t=1) = e * sin(x)`, `u_x(0, t) = e^t`, and
    `u(2 pi n + pi/2, t) = e^t`, where `n` is some integer.

    The exact solution `u(x, t=0) = e^t sin(x)`.
    """

    def final_cond_fn(x):
      return math.e * math.sin(x)

    def expected_result_fn(x):
      return tf.sin(x)

    @neumann
    def lower_boundary_fn(t, x):
      del x
      return -tf.exp(t)

    @dirichlet
    def upper_boundary_fn(t, x):
      del x
      return tf.exp(t)

    grid = grids.uniform_grid(minimums=[0], maximums=[10.5 * math.pi],
                              sizes=[1000], dtype=np.float32)
    self._testHeatEquation(
        grid,
        final_t=1,
        time_step=0.01,
        final_cond_fn=final_cond_fn,
        expected_result_fn=expected_result_fn,
        one_step_fn=crank_nicolson_step(),
        lower_boundary_fn=lower_boundary_fn,
        upper_boundary_fn=upper_boundary_fn,
        error_tolerance=1e-3) 

def testHeatEquation_WithNeumannBoundaryConditions(self):
    """Test for Neumann boundary conditions.

    Tests solving heat equation with the following boundary conditions:
    `u(x, t=1) = e * sin(x)`, `u_x(0, t) = e^t`, and
    `u_x(2 pi n + pi/2, t) = 0`, where `n` is some integer.

    The exact solution `u(x, t=0) = e^t sin(x)`.
    """

    def final_cond_fn(x):
      return math.e * math.sin(x)

    def expected_result_fn(x):
      return tf.sin(x)

    @neumann
    def lower_boundary_fn(t, x):
      del x
      return -tf.exp(t)

    @neumann
    def upper_boundary_fn(t, x):
      del t, x
      return 0

    grid = grids.uniform_grid(
        minimums=[0.0], maximums=[10.5 * math.pi], sizes=[1000],
        dtype=np.float32)
    self._testHeatEquation(
        grid,
        final_t=1,
        time_step=0.01,
        final_cond_fn=final_cond_fn,
        expected_result_fn=expected_result_fn,
        one_step_fn=crank_nicolson_step(),
        lower_boundary_fn=lower_boundary_fn,
        upper_boundary_fn=upper_boundary_fn,
        error_tolerance=1e-3) 

def testInnerSecondOrderCoeff(self):
    """Tests handling inner_second_order_coeff.

    As in previous test, take the diffusion equation `v_{t} - v_{xx} = 0` and
    substitute `v = exp(x) u`, but this time keep exponent under the derivative:
    `u_{t} - exp(-x)[exp(x)u]_{xx} = 0`. Expect the same solution as in
    previous test.
    """
    grid = grids.uniform_grid(
        minimums=[0], maximums=[1], sizes=[501], dtype=tf.float32)
    xs = grid[0]

    final_t = 0.1
    time_step = 0.001

    def second_order_coeff_fn(t, coord_grid):
      del t
      x = coord_grid[0]
      return [[-tf.exp(-x)]]

    def inner_second_order_coeff_fn(t, coord_grid):
      del t
      x = coord_grid[0]
      return [[tf.exp(x)]]

    initial = tf.exp(-xs) * _reference_pde_initial_cond(xs)
    expected = tf.exp(-xs) * _reference_pde_solution(xs, final_t)

    actual = fd_solvers.solve_forward(
        start_time=0,
        end_time=final_t,
        coord_grid=grid,
        values_grid=initial,
        time_step=time_step,
        second_order_coeff_fn=second_order_coeff_fn,
        inner_second_order_coeff_fn=inner_second_order_coeff_fn)[0]

    self.assertAllClose(expected, actual, atol=1e-3, rtol=1e-3) 

def test_douglas_step_with_batching(self):
    u = np.arange(0, 80, dtype=np.float32).reshape(4, 4, 5)
    d = np.arange(10, 90, dtype=np.float32).reshape(4, 4, 5)
    dx = np.array([d, -3 * d, 2 * d])
    dy = 2 * dx
    dxy = (np.arange(-20, 60, dtype=np.float32).reshape(4, 4, 5)) * 4
    theta = 0.3
    bx = np.arange(20, 100, dtype=np.float32).reshape(4, 4, 5)
    by = np.arange(30, 110, dtype=np.float32).reshape(4, 4, 5)

    def equation_params_fn(t):
      del t
      return ([[_tfconst(dy), _spread_mixed_term(_tfconst(dxy))],
               [None, _tfconst(dx)]],
              [_tfconst(by), _tfconst(bx)])

    scheme = douglas_adi_scheme(theta=theta)
    actual = self.evaluate(
        scheme(value_grid=tf.constant(u, dtype=tf.float32), t1=0, t2=1,
               equation_params_fn=equation_params_fn, n_dims=2))
    expected = np.zeros_like(u)
    for i in range(4):
      expected[i] = self._simplified_douglas_step_2d(u[i], dx[:, i], dy[:, i],
                                                     dxy[i], bx[i],
                                                     by[i], 0, 1, theta)

    self.assertLess(np.max(np.abs(expected - actual)), 0.01) 

def _tfconst(np_array):
  return tf.constant(np_array, dtype=tf.float32) 

def testAnisotropicDiffusion_ListOfTensors(self):
    def pack_second_order_coeff_fn(u_yy, u_xy, u_xx):
      return [tf.constant([u_yy, u_xy], dtype=tf.float32),
              tf.constant([u_xy, u_xx], dtype=tf.float32)]
    self._testDiffusionInDiagonalDirection(pack_second_order_coeff_fn) 

def testAnisotropicDiffusion_2DTensor(self):
    def pack_second_order_coeff_fn(u_yy, u_xy, u_xx):
      return tf.convert_to_tensor([[u_yy, u_xy], [u_xy, u_xx]],
                                  dtype=tf.float32)
    self._testDiffusionInDiagonalDirection(pack_second_order_coeff_fn)

  # pylint: disable=g-doc-args 

def testInnerSecondOrderCoeff(self):
    """Tests handling inner_second_order_coeff.

    As in previous test, take the diffusion equation
    `v_{t} - v_{xx} - v_{yy} = 0` and substitute `v = exp(x + 2y) u`, but this
    time keep exponent under the derivative:
    `u_{t} - exp(-x)[exp(x)u]_{xx} - exp(-2y)[exp(2y)u]_{yy} = 0`.
    Expect the same solution as in previous test.
    """
    grid = grids.uniform_grid(
        minimums=[0, 0], maximums=[1, 1], sizes=[201, 251], dtype=tf.float32)
    ys, xs = grid

    final_t = 0.1
    time_step = 0.002

    def second_order_coeff_fn(t, coord_grid):
      del t
      y, x = tf.meshgrid(*coord_grid, indexing='ij')
      return [[-tf.exp(-2 * y), None], [None, -tf.exp(-x)]]

    def inner_second_order_coeff_fn(t, coord_grid):
      del t
      y, x = tf.meshgrid(*coord_grid, indexing='ij')
      return [[tf.exp(2 * y), None], [None, tf.exp(x)]]

    exp = _dir_prod(tf.exp(-2 * ys), tf.exp(-xs))
    initial = exp * _reference_2d_pde_initial_cond(xs, ys)
    expected = exp * _reference_2d_pde_solution(xs, ys, final_t)

    actual = fd_solvers.solve_forward(
        start_time=0,
        end_time=final_t,
        coord_grid=grid,
        values_grid=initial,
        time_step=time_step,
        second_order_coeff_fn=second_order_coeff_fn,
        inner_second_order_coeff_fn=inner_second_order_coeff_fn)[0]

    self.assertAllClose(expected, actual, atol=1e-3, rtol=1e-3) 

def test_float32(self):
    minimum = np.array([1.0, 1.0], dtype=np.float32)
    scales = np.array([2.0, 3.0], dtype=np.float32)
    start = np.zeros_like(minimum)

    @tff.math.make_val_and_grad_fn
    def quadratic(x):
      return tf.reduce_sum(input_tensor=scales * (x - minimum)**2)

    result = tff.math.optimizer.conjugate_gradient_minimize(
        quadratic, initial_position=start)
    self.assertEqual(result.position.dtype, tf.float32)
    self.assertArrayNear(self.evaluate(result.position), minimum, 1e-5) 

def testFindsAllRootsUsingFloat32(self):
    self._testFindsAllRoots(
        objective_fn=polynomial5,
        left_bracket=[-4, 1],
        right_bracket=[3, -1],
        dtype=tf.float32,
        expected_roots=[-0.14823253010472962, -0.14823253013216148],
        expected_num_iterations=[13, 7])