Python tensorflow.compat.v2.constant() Examples

def test_with_xla(self):
    @tf.function
    def fn():
      x = tf.constant([[3.0, 4.0], [30.0, 40.0]])
      y = 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):
          del i
          x, y = state
          return [x * alpha - beta, y * beta + x]
        out = for_loop(body, [x, y], [alpha, beta], 3)
      return tape.gradient(out[1], beta)

    grad = self.evaluate(tf.xla.experimental.compile(fn))[0]
    self.assertAllEqual(783, grad) 

def testFindsRootForFlatFunction(self):
    # Flat in the [-0.5, 0.5] range.
    objective_fn = lambda x: 0 if x == 0 else x * exp(-1 / x**2)

    left_bracket = [-10]
    right_bracket = [1]
    expected_num_iterations = [13]

    expected_num_iterations, result = self.evaluate([
        tf.constant(expected_num_iterations, dtype=tf.int32),
        root_search.brentq(objective_fn,
                           tf.constant(left_bracket, dtype=tf.float64),
                           tf.constant(right_bracket, dtype=tf.float64))
    ])

    _, value_at_roots, num_iterations, _ = result

    # Simply check that the objective function is close to the root for the
    # returned estimate. Do not check the estimate itself.
    # Unlike Brent's original algorithm (and the SciPy implementation), this
    # implementation stops the search as soon as a good enough root estimate is
    # found. As a result, the estimate may significantly differ from the one
    # returned by SciPy for functions which are extremely flat around the root.
    self.assertAllClose(value_at_roots, [0.])
    self.assertAllEqual(num_iterations, expected_num_iterations) 

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 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 testWithNoIteration(self):
    left_bracket = [-10, 1]
    right_bracket = [10, -1]

    first_guess = tf.constant(left_bracket, dtype=tf.float64)
    second_guess = tf.constant(right_bracket, dtype=tf.float64)

    # Skip iteration entirely.
    # Should return a Tensor built from the best guesses in input positions.
    guess, result = self.evaluate([
        tf.constant([-10, -1], dtype=tf.float64),
        root_search.brentq(
            polynomial5, first_guess, second_guess, max_iterations=0)
    ])

    self.assertAllEqual(result.estimated_root, guess) 

def test_piecewise_constant_integral_with_batch(self):
    """Tests PiecewiseConstantFunc with batching."""
    for dtype in [np.float32, np.float64]:
      x = np.array([[[0.0, 0.1, 2.0, 11.0], [0.0, 2.0, 3.0, 9.0]],
                    [[0.0, 1.0, 2.0, 3.0], [4.0, 5.0, 6.0, 7.0]]])
      jump_locations = np.array([[[0.1, 10.0], [1.5, 10.0]],
                                 [[1.0, 2.0], [5.0, 6.0]]])
      values = tf.constant([[[3, 4, 5], [3, 4, 5]],
                            [[3, 4, 5], [3, 4, 5]]], dtype=dtype)
      piecewise_func = piecewise.PiecewiseConstantFunc(jump_locations, values,
                                                       dtype=dtype)
      value = piecewise_func.integrate(x, x + 1.1)
      self.assertEqual(value.dtype.as_numpy_dtype, dtype)
      expected_value = np.array([[[4.3, 4.4, 4.4, 5.5],
                                  [3.3, 4.4, 4.4, 4.5]],
                                 [[3.4, 4.5, 5.5, 5.5],
                                  [3.4, 4.5, 5.5, 5.5]]])
      self.assertAllClose(value, expected_value, atol=1e-5, rtol=1e-5) 

def _test_batches_and_types(self, integrate_function, args):
    """Checks handling batches and dtypes."""
    dtypes = [np.float32, np.float64, np.complex64, np.complex128]
    a = [[0.0, 0.0], [0.0, 0.0]]
    b = [[np.pi / 2, np.pi], [1.5 * np.pi, 2 * np.pi]]
    a = [a, a]
    b = [b, b]
    k = tf.constant([[[[1.0]]], [[[2.0]]]])
    func = lambda x: tf.cast(k, dtype=x.dtype) * tf.sin(x)
    ans = [[[1.0, 2.0], [1.0, 0.0]], [[2.0, 4.0], [2.0, 0.0]]]

    results = []
    for dtype in dtypes:
      lower = tf.constant(a, dtype=dtype)
      upper = tf.constant(b, dtype=dtype)
      results.append(integrate_function(func, lower, upper, **args))

    results = self.evaluate(results)

    for i in range(len(results)):
      assert results[i].dtype == dtypes[i]
      assert np.allclose(results[i], ans, atol=1e-3) 

def test_piecewise_constant_value_with_batch_and_repetitions(self):
    """Tests PiecewiseConstantFunc with batching and repetitive values."""
    for dtype in [np.float32, np.float64]:
      x = tf.constant([[-4.1, 0.1, 1., 2., 10, 11.],
                       [1., 2., 3., 2., 5., 9.]], dtype=dtype)
      jump_locations = tf.constant([[0.1, 0.1, 1., 1., 10., 10.],
                                    [-1., 1.2, 2.2, 2.2, 2.2, 8.]], dtype=dtype)
      values = tf.constant([[3, 3, 4, 5, 5., 2, 6.],
                            [-1, -5, 2, 5, 5., 5., 1.]], dtype=dtype)
      piecewise_func = piecewise.PiecewiseConstantFunc(jump_locations, values,
                                                       dtype=dtype)
      # Also verifies left-continuity
      value = piecewise_func(x, left_continuous=True)
      self.assertEqual(value.dtype.as_numpy_dtype, dtype)
      expected_value = np.array([[3., 3., 4., 5., 5., 6.],
                                 [-5., 2., 5., 2., 5., 1.]])
      self.assertAllEqual(value, expected_value) 

def test_piecewise_constant_value_with_batch(self):
    """Tests PiecewiseConstantFunc with batching."""
    for dtype in [np.float32, np.float64]:
      x = np.array([[[0.0, 0.1, 2.0, 11.0], [0.0, 2.0, 3.0, 9.0]],
                    [[0.0, 1.0, 2.0, 3.0], [4.0, 5.0, 6.0, 7.0]]])
      jump_locations = np.array([[[0.1, 10.0], [1.5, 10.0]],
                                 [[1.0, 2.0], [5.0, 6.0]]])
      values = tf.constant([[[3, 4, 5], [3, 4, 5]],
                            [[3, 4, 5], [3, 4, 5]]], dtype=dtype)
      piecewise_func = piecewise.PiecewiseConstantFunc(jump_locations, values,
                                                       dtype=dtype)
      # Also verifies right-continuity
      value = piecewise_func(x, left_continuous=False)
      self.assertEqual(value.dtype.as_numpy_dtype, dtype)
      expected_value = np.array([[[3.0, 4.0, 4.0, 5.0],
                                  [3.0, 4.0, 4.0, 4.0]],
                                 [[3.0, 4.0, 5.0, 5.0],
                                  [3.0, 4.0, 5.0, 5.0]]])
      self.assertAllEqual(value, expected_value) 

def testFindsAllRootsUsingFloat16(self):
    left_bracket = [-2, 1]
    right_bracket = [2, -1]
    expected_num_iterations = [9, 4]

    expected_num_iterations, result = self.evaluate([
        tf.constant(expected_num_iterations, dtype=tf.int32),
        root_search.brentq(polynomial5,
                           tf.constant(left_bracket, dtype=tf.float16),
                           tf.constant(right_bracket, dtype=tf.float16))
    ])

    _, value_at_roots, num_iterations, _ = result

    # Simply check that the objective function is close to the root for the
    # returned estimates. Do not check the estimates themselves.
    # Using float16 may yield root estimates which differ from those returned
    # by the SciPy implementation.
    self.assertAllClose(value_at_roots, [0., 0.], atol=1e-3)
    self.assertAllEqual(num_iterations, expected_num_iterations) 

def test_incompatible_x1_x2_batch_shape(self):
    """Tests that `x1` and `x2` should have the same batch shape."""
    for dtype in [np.float32, np.float64]:
      x1 = np.array([[0., 0.1, 2., 11.],
                     [0., 0.1, 2., 11.]])
      x2 = np.array([[0., 0.1, 2., 11.],
                     [0., 0.1, 2., 11.], [0., 0.1, 2., 11.]])
      x3 = x2 + 1
      jump_locations = np.array([[0.1, 10]])
      values = tf.constant([[3, 4, 5]], dtype=dtype)
      piecewise_func = piecewise.PiecewiseConstantFunc(jump_locations, values,
                                                       dtype=dtype)
      with self.assertRaises(ValueError):
        piecewise_func.integrate(x1, x2)
      # `x2` and `x3` have the same batch shape but different from
      # the batch shape of `jump_locations`
      with self.assertRaises(ValueError):
        piecewise_func.integrate(x2, x3) 

def testPad(self):
    t = [[1, 2, 3], [4, 5, 6]]
    paddings = [[1, 1,], [2, 2]]
    self.assertAllEqual(
        array_ops.pad(t, paddings, 'constant'),
        [[0, 0, 0, 0, 0, 0, 0], [0, 0, 1, 2, 3, 0, 0], [0, 0, 4, 5, 6, 0, 0],
         [0, 0, 0, 0, 0, 0, 0]])

    self.assertAllEqual(
        array_ops.pad(t, paddings, 'reflect'),
        [[6, 5, 4, 5, 6, 5, 4], [3, 2, 1, 2, 3, 2, 1], [6, 5, 4, 5, 6, 5, 4],
         [3, 2, 1, 2, 3, 2, 1]])

    self.assertAllEqual(
        array_ops.pad(t, paddings, 'symmetric'),
        [[2, 1, 1, 2, 3, 3, 2], [2, 1, 1, 2, 3, 3, 2], [5, 4, 4, 5, 6, 6, 5],
         [5, 4, 4, 5, 6, 6, 5]]) 

def argsort(a, axis=-1, kind='quicksort', order=None):  # pylint: disable=missing-docstring
  # TODO(nareshmodi): make string tensors also work.
  if kind not in ('quicksort', 'stable'):
    raise ValueError("Only 'quicksort' and 'stable' arguments are supported.")
  if order is not None:
    raise ValueError("'order' argument to sort is not supported.")
  stable = (kind == 'stable')

  a = array_ops.array(a).data

  def _argsort(a, axis, stable):
    if axis is None:
      a = tf.reshape(a, [-1])
      axis = 0

    return tf.argsort(a, axis, stable=stable)

  tf_ans = tf.cond(
      tf.rank(a) == 0, lambda: tf.constant([0]),
      lambda: _argsort(a, axis, stable))

  return array_ops.array(tf_ans, dtype=np.intp) 

def testHomogeneousBackwards(self, scheme, accuracy_order):
    # Tests solving du/dt = At for a backward time step.
    # Compares with exact solution u(0) = exp(-At) u(t).
    time_step = 0.0001
    u = tf.constant([1, 2, -1, -2], dtype=tf.float64)
    matrix = tf.constant(
        [[1, -1, 0, 0], [3, 1, 2, 0], [0, -2, 1, 4], [0, 0, 3, 1]],
        dtype=tf.float64)

    tridiag_form = self._convert_to_tridiagonal_format(matrix)
    actual = self.evaluate(
        scheme(u, time_step, 0, lambda t: (tridiag_form, None)))

    expected = self.evaluate(
        tf.squeeze(
            tf.matmul(
                tf.linalg.expm(-matrix * time_step), tf.expand_dims(u, 1))))

    error_tolerance = 30 * time_step**(accuracy_order + 1)
    self.assertLess(np.max(np.abs(actual - expected)), error_tolerance) 

def _tf_gcd(x1, x2):
  def _gcd_cond_fn(x1, x2):
    return tf.reduce_any(x2 != 0)
  def _gcd_body_fn(x1, x2):
    # tf.math.mod will raise an error when any element of x2 is 0. To avoid
    # that, we change those zeros to ones. Their values don't matter because
    # they won't be used.
    x2_safe = tf.where(x2 != 0, x2, tf.constant(1, x2.dtype))
    x1, x2 = (tf.where(x2 != 0, x2, x1),
              tf.where(x2 != 0, tf.math.mod(x1, x2_safe),
                       tf.constant(0, x2.dtype)))
    return (tf.where(x1 < x2, x2, x1), tf.where(x1 < x2, x1, x2))
  if (not np.issubdtype(x1.dtype.as_numpy_dtype, np.integer) or
      not np.issubdtype(x2.dtype.as_numpy_dtype, np.integer)):
    raise ValueError("Arguments to gcd must be integers.")
  shape = tf.broadcast_static_shape(x1.shape, x2.shape)
  x1 = tf.broadcast_to(x1, shape)
  x2 = tf.broadcast_to(x2, shape)
  gcd, _ = tf.while_loop(_gcd_cond_fn, _gcd_body_fn,
                         (tf.math.abs(x1), tf.math.abs(x2)))
  return gcd 

def testInhomogeneous(self, scheme, accuracy_order):
    # Tests solving du/dt = At + b for a time step.
    # Compares with exact solution u(t) = exp(At) u(0) + (exp(At) - 1) A^(-1) b.
    time_step = 0.0001
    u = tf.constant([1, 2, -1, -2], dtype=tf.float64)
    matrix = tf.constant(
        [[1, -1, 0, 0], [3, 1, 2, 0], [0, -2, 1, 4], [0, 0, 3, 1]],
        dtype=tf.float64)
    b = tf.constant([1, -1, -2, 2], dtype=tf.float64)

    tridiag_form = self._convert_to_tridiagonal_format(matrix)
    actual = self.evaluate(scheme(u, 0, time_step, lambda t: (tridiag_form, b)))

    exponent = tf.linalg.expm(matrix * time_step)
    eye = tf.eye(4, 4, dtype=tf.float64)
    u = tf.expand_dims(u, 1)
    b = tf.expand_dims(b, 1)
    expected = (
        tf.matmul(exponent, u) +
        tf.matmul(exponent - eye, tf.matmul(tf.linalg.inv(matrix), b)))
    expected = self.evaluate(tf.squeeze(expected))

    error_tolerance = 30 * time_step**(accuracy_order + 1)
    self.assertLess(np.max(np.abs(actual - expected)), error_tolerance) 

def testInhomogeneousBackwards(self, scheme, accuracy_order):
    # Tests solving du/dt = At + b for a backward time step.
    # Compares with exact solution u(0) = exp(-At) u(t)
    # + (exp(-At) - 1) A^(-1) b.
    time_step = 0.0001
    u = tf.constant([1, 2, -1, -2], dtype=tf.float64)
    matrix = tf.constant(
        [[1, -1, 0, 0], [3, 1, 2, 0], [0, -2, 1, 4], [0, 0, 3, 1]],
        dtype=tf.float64)
    b = tf.constant([1, -1, -2, 2], dtype=tf.float64)

    tridiag_form = self._convert_to_tridiagonal_format(matrix)
    actual = self.evaluate(scheme(u, time_step, 0, lambda t: (tridiag_form, b)))

    exponent = tf.linalg.expm(-matrix * time_step)
    eye = tf.eye(4, 4, dtype=tf.float64)
    u = tf.expand_dims(u, 1)
    b = tf.expand_dims(b, 1)
    expected = (
        tf.matmul(exponent, u) +
        tf.matmul(exponent - eye, tf.matmul(tf.linalg.inv(matrix), b)))
    expected = self.evaluate(tf.squeeze(expected))

    error_tolerance = 30 * time_step**(accuracy_order + 1)
    self.assertLess(np.max(np.abs(actual - expected)), error_tolerance) 

def triu(m, k=0):  # pylint: disable=missing-docstring
  m = asarray(m).data
  m_shape = m.shape.as_list()

  if len(m_shape) < 2:
    raise ValueError('Argument to triu must have rank at least 2')

  if m_shape[-1] is None or m_shape[-2] is None:
    raise ValueError('Currently, the last two dimensions of the input array '
                     'need to be known.')

  z = tf.constant(0, m.dtype)

  mask = tri(*m_shape[-2:], k=k - 1, dtype=bool)
  return utils.tensor_to_ndarray(
      tf.where(tf.broadcast_to(mask, tf.shape(m)), z, m)) 

def tril(m, k=0):  # pylint: disable=missing-docstring
  m = asarray(m).data
  m_shape = m.shape.as_list()

  if len(m_shape) < 2:
    raise ValueError('Argument to tril must have rank at least 2')

  if m_shape[-1] is None or m_shape[-2] is None:
    raise ValueError('Currently, the last two dimensions of the input array '
                     'need to be known.')

  z = tf.constant(0, m.dtype)

  mask = tri(*m_shape[-2:], k=k, dtype=bool)
  return utils.tensor_to_ndarray(
      tf.where(tf.broadcast_to(mask, tf.shape(m)), m, z)) 

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 _seed2key(a):
  """Converts an RNG seed to an RNG key.

  Args:
    a: an RNG seed, a tensor of shape [2] and dtype `tf.int32`.

  Returns:
    an RNG key, an ndarray of shape [] and dtype `np.int64`.
  """

  def int32s_to_int64(a):
    """Converts an int32 tensor of shape [2] to an int64 tensor of shape []."""
    a = tf.bitwise.bitwise_or(
        tf.cast(a[0], tf.uint64),
        tf.bitwise.left_shift(
            tf.cast(a[1], tf.uint64), tf.constant(32, tf.uint64)))
    a = tf.cast(a, tf.int64)
    return a

  return tf_np.asarray(int32s_to_int64(a)) 

def _key2seed(a):
  """Converts an RNG key to an RNG seed.

  Args:
    a: an RNG key, an ndarray of shape [] and dtype `np.int64`.

  Returns:
    an RNG seed, a tensor of shape [2] and dtype `tf.int32`.
  """

  def int64_to_int32s(a):
    """Converts an int64 tensor of shape [] to an int32 tensor of shape [2]."""
    a = tf.cast(a, tf.uint64)
    fst = tf.cast(a, tf.uint32)
    snd = tf.cast(
        tf.bitwise.right_shift(a, tf.constant(32, tf.uint64)), tf.uint32)
    a = [fst, snd]
    a = tf.nest.map_structure(lambda x: tf.cast(x, tf.int32), a)
    a = tf.stack(a)
    return a

  return int64_to_int32s(a.data) 

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_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 test_piecewise_constant_value_no_batch(self):
    """Tests PiecewiseConstantFunc with no batching."""
    for dtype in [np.float32, np.float64]:
      x = np.array([0., 0.1, 2., 11.])
      jump_locations = np.array([0.1, 10], dtype=dtype)
      values = tf.constant([3, 4, 5], dtype=dtype)
      piecewise_func = piecewise.PiecewiseConstantFunc(jump_locations, values,
                                                       dtype=dtype)
      # Also verifies left-continuity
      value = piecewise_func(x)
      self.assertEqual(value.dtype.as_numpy_dtype, dtype)
      expected_value = np.array([3., 3., 4., 5.])
      self.assertAllEqual(value, expected_value) 

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

def test_logistic_regression(self):
    dim = 5
    n_objs = 10000
    np.random.seed(1)
    betas = np.random.randn(dim)  # The true beta
    intercept = np.random.randn()  # The true intercept
    features = np.random.randn(n_objs, dim)  # The feature matrix
    probs = 1 / (1 + np.exp(
        -np.matmul(features, np.expand_dims(betas, -1)) - intercept))
    labels = np.random.binomial(1, probs)  # The true labels
    regularization = 0.8
    feat = tf.constant(features, dtype=tf.float64)
    lab = tf.constant(labels, dtype=feat.dtype)

    def f_negative_log_likelihood(params):
      intercept, beta = params[0], params[1:]
      logit = tf.matmul(feat, tf.expand_dims(beta, -1)) + intercept
      log_likelihood = tf.reduce_sum(
          tf.nn.sigmoid_cross_entropy_with_logits(labels=lab, logits=logit))
      l2_penalty = regularization * tf.reduce_sum(beta**2)
      total_loss = log_likelihood + l2_penalty
      return total_loss
    start_point = np.ones(dim + 1)
    argmin = [
        -2.38636155, 1.61778325, -0.60694238, -0.51523609, -1.09832275,
        0.88892742
    ]

    self._check_algorithm(
        func=f_negative_log_likelihood,
        start_point=start_point,
        expected_argmin=argmin,
        gtol=1e-5) 

def testNoTimeDependence(self):
    """Test for the case where all terms (quadratic, linear, shift) are null."""
    grid = grids.uniform_grid(
        minimums=[-10, -20],
        maximums=[10, 20],
        sizes=[201, 301],
        dtype=tf.float32)
    ys = self.evaluate(grid[0])
    xs = self.evaluate(grid[1])

    time_step = 0.1
    final_t = 1
    variance = 1

    final_cond = np.outer(_gaussian(ys, variance), _gaussian(xs, variance))
    final_values = tf.expand_dims(tf.constant(final_cond, dtype=tf.float32),
                                  axis=0)
    bound_cond = [(_zero_boundary, _zero_boundary),
                  (_zero_boundary, _zero_boundary)]
    step_fn = douglas_adi_step(theta=0.5)
    result = fd_solvers.solve_backward(
        start_time=final_t,
        end_time=0,
        coord_grid=grid,
        values_grid=final_values,
        time_step=time_step,
        one_step_fn=step_fn,
        boundary_conditions=bound_cond,
        dtype=grid[0].dtype)
    expected = final_cond  # No time dependence.
    self._assertClose(expected, result) 

def test_diffs_differentiable(self):
    """Tests that the diffs op is differentiable."""
    x = tf.constant(2.0)
    xv = tf.stack([x, x * x, x * x * x], axis=0)

    # Produces [x, x^2 - x, x^3 - x^2]
    dxv = self.evaluate(math.diff(xv))
    np.testing.assert_array_equal(dxv, [2., 2., 4.])

    grad = self.evaluate(tf.gradients(math.diff(xv), x)[0])
    # Note that TF gradients adds up the components of the jacobian.
    # The sum of [1, 2x-1, 3x^2-2x] at x = 2 is 12.
    self.assertEqual(grad, 12.0) 

def test_invalid_value_event_shape(self):
    """Tests that `values` event shape is `jump_locations` event shape + 1."""
    for dtype in [np.float32, np.float64]:
      jump_locations = np.array([[0.1, 10], [2., 10]])
      values = tf.constant([[3, 4, 5, 6], [3, 4, 5, 7]], dtype=dtype)
      with self.assertRaises(ValueError):
        piecewise.PiecewiseConstantFunc(jump_locations, values, dtype=dtype)