def test_jacobian_fixed_shape(self):
   x = random_ops.random_uniform([2, 2])
   y = math_ops.matmul(x, x, transpose_a=True)
   jacobian_pfor = gradients.jacobian(y, x, use_pfor=True)
   jacobian_while = gradients.jacobian(y, x, use_pfor=False)
   answer = ops.convert_to_tensor([[
       gradient_ops.gradients(y[0][0], x)[0],
       gradient_ops.gradients(y[0][1], x)[0]
   ], [
       gradient_ops.gradients(y[1][0], x)[0],
       gradient_ops.gradients(y[1][1], x)[0]
   ]])
   self.run_and_assert_equal(answer, jacobian_pfor)
   self.run_and_assert_equal(answer, jacobian_while)

def create_lstm_hessian(batch_size, state_size, steps):
  _, output = lstm_model_fn(batch_size, state_size, steps)
  weights = variables.trainable_variables()
  pfor_jacobians = gradients.jacobian(output, weights, use_pfor=True)
  pfor_hessians = [
      gradients.jacobian(x, weights, use_pfor=True) for x in pfor_jacobians
  ]
  # TODO(agarwal): using two nested while_loop doesn't seem to work here.
  # Hence we use pfor_jacobians for computing while_hessians.
  while_jacobians = pfor_jacobians
  while_hessians = [
      gradients.jacobian(x, weights, use_pfor=False) for x in while_jacobians
  ]
  return pfor_hessians, while_hessians

 def test_jacobian_unknown_shape(self):
   with self.test_session() as sess:
     x = array_ops.placeholder(dtypes.float32, shape=[None, None])
     y = math_ops.matmul(x, x, transpose_a=True)
     jacobian_pfor = gradients.jacobian(y, x, use_pfor=True)
     jacobian_while = gradients.jacobian(y, x, use_pfor=False)
     answer = ops.convert_to_tensor([[
         gradient_ops.gradients(y[0][0], x)[0],
         gradient_ops.gradients(y[0][1], x)[0]
     ], [
         gradient_ops.gradients(y[1][0], x)[0],
         gradient_ops.gradients(y[1][1], x)[0]
     ]])
     ans, pfor_value, while_value = sess.run(
         [answer, jacobian_pfor, jacobian_while],
         feed_dict={x: [[1, 2], [3, 4]]})
     self.assertAllClose(ans, pfor_value)
     self.assertAllClose(ans, while_value)

  def test_jacobian_scan_shape(self):
    # Shape x: [3, 4]
    x = random_ops.random_uniform([3, 4])
    elems = random_ops.random_uniform([6])
    # Shape y: [6, 3, 4]
    y = functional_ops.scan(lambda a, e: a + e, elems, initializer=x)
    jacobian = gradients.jacobian(y, x)

    expected_shape = [6, 3, 4, 3, 4]
    self.assertAllEqual(expected_shape, jacobian.shape.as_list())

  def test_jacobian_while_loop_shape(self):
    # Shape x: [3, 4]
    x = random_ops.random_uniform([3, 4])
    _, y = tf_control_flow_ops.while_loop(lambda i, a: i > 5.,
                                          lambda i, a: (i + 1, a + i),
                                          (constant_op.constant(0.), x))
    # Shape y: [2, 3]
    y = y[:2, :3]
    jacobian = gradients.jacobian(y, x)

    expected_shape = [2, 3, 3, 4]
    self.assertAllEqual(expected_shape, jacobian.shape.as_list())

def create_fc_per_eg_jacobians(batch_size, activation_size, num_layers):
  model = FullyConnectedModel(activation_size=activation_size,
                              num_layers=num_layers)
  inp = random_ops.random_normal([batch_size, activation_size])
  output = model(inp)
  jacobians = gradients.jacobian(output, variables.trainable_variables())

  def loop_fn(i, use_pfor):
    inp_i = array_ops.expand_dims(array_ops.gather(inp, i), 0)
    output = array_ops.reshape(model(inp_i), [-1])
    return gradients.jacobian(
        output, variables.trainable_variables(), use_pfor=use_pfor)

  per_eg_jacobians_pfor = control_flow_ops.pfor(
      functools.partial(loop_fn, use_pfor=True),
      batch_size)
  per_eg_jacobians_while = control_flow_ops.for_loop(
      functools.partial(loop_fn, use_pfor=False),
      [dtypes.float32] * len(variables.trainable_variables()), batch_size)
  return jacobians, per_eg_jacobians_pfor, per_eg_jacobians_while

 def loop_fn(i, use_pfor):
   inp_i = array_ops.expand_dims(array_ops.gather(inp, i), 0)
   output = array_ops.reshape(model(inp_i), [-1])
   return gradients.jacobian(
       output, variables.trainable_variables(), use_pfor=use_pfor)

 def test_jacobian_parallel_iterations(self):
   x = constant_op.constant([[1., 2], [3, 4]])
   y = math_ops.matmul(x, x)
   self.assertAllClose(gradients.jacobian(y, x, parallel_iterations=2),
                       gradients.jacobian(y, x, parallel_iterations=3))