def create_mnist_batch_jacobian(batch_size, data_format, training):
  images = random_ops.random_uniform([batch_size, 28, 28])
  model = Mnist(data_format)
  logits = model(images, training=training)

  pfor_jacobian = gradients.batch_jacobian(logits, images, use_pfor=True)
  while_jacobian = gradients.batch_jacobian(logits, images, use_pfor=False)
  return pfor_jacobian, while_jacobian

def create_lstm_batch_hessian(batch_size, state_size, steps):
  inp, output = lstm_model_fn(batch_size, state_size, steps)
  pfor_jacobian = gradients.batch_jacobian(output, inp, use_pfor=True)
  pfor_jacobian = array_ops.reshape(pfor_jacobian, [batch_size, -1])
  pfor_hessian = gradients.batch_jacobian(pfor_jacobian, inp, use_pfor=True)
  # TODO(agarwal): using two nested while_loop doesn't seem to work here.
  # Hence we use pfor_jacobian for computing while_hessian.
  while_jacobian = pfor_jacobian
  while_hessian = gradients.batch_jacobian(while_jacobian, inp, use_pfor=False)
  return pfor_hessian, while_hessian

 def test_batch_jacobian_fixed_shape(self):
   x = random_ops.random_uniform([2, 3, 5])
   y = x * x
   batch_jacobian_pfor = gradients.batch_jacobian(y, x, use_pfor=True)
   batch_jacobian_while = gradients.batch_jacobian(y, x, use_pfor=False)
   two_x = 2 * x
   answer = array_ops.stack(
       [array_ops.diag(two_x[0]),
        array_ops.diag(two_x[1])])
   self.run_and_assert_equal(answer, batch_jacobian_pfor)
   self.run_and_assert_equal(answer, batch_jacobian_while)

 def test_batch_jacobian_unknown_shape(self):
   with self.test_session() as sess:
     x = array_ops.placeholder(dtypes.float32)
     y = x * x
     batch_jacobian_pfor = gradients.batch_jacobian(y, x, use_pfor=True)
     batch_jacobian_while = gradients.batch_jacobian(y, x, use_pfor=False)
     two_x = 2 * x
     answer = array_ops.stack(
         [array_ops.diag(two_x[0]),
          array_ops.diag(two_x[1])])
     ans, pfor_value, while_value = sess.run(
         [answer, batch_jacobian_pfor, batch_jacobian_while],
         feed_dict={x: [[1, 2], [3, 4]]})
     self.assertAllClose(ans, pfor_value)
     self.assertAllClose(ans, while_value)

 def test_batch_jacobian_bad_unknown_shapes(self):
   with self.test_session() as sess:
     x = array_ops.placeholder(dtypes.float32)
     y = array_ops.concat([x, x], axis=0)
     jacobian = gradients.batch_jacobian(y, x)
     with self.assertRaisesRegexp(errors.InvalidArgumentError,
                                  "assertion failed"):
       sess.run(jacobian, feed_dict={x: [[1, 2], [3, 4]]})

def create_dynamic_lstm_batch_jacobian(batch_size, state_size, max_steps):
  inp, (_, final_state) = dynamic_lstm_model_fn(batch_size, state_size,
                                                max_steps)
  pfor_jacobian = gradients.batch_jacobian(final_state.c, inp, use_pfor=True)
  # Note that use_pfor=False does not work above given the current limitations
  # on implementation of while_loop. So we statically unroll the looping in the
  # jacobian computation.
  while_gradients = [
      gradient_ops.gradients(array_ops.gather(final_state.c, i, axis=1), inp)[0]
      for i in range(state_size)
  ]
  return pfor_jacobian, while_gradients

def create_fc_batch_jacobian(batch_size, activation_size, num_layers):
  inp, output = fully_connected_model_fn(batch_size, activation_size,
                                         num_layers)
  pfor_jacobian = gradients.batch_jacobian(output, inp, use_pfor=True)
  while_jacobian = gradients.batch_jacobian(output, inp, use_pfor=False)
  return pfor_jacobian, while_jacobian

 def test_batch_jacobian_bad_shapes(self):
   x = random_ops.random_uniform([2, 2])
   y = random_ops.random_uniform([3, 2])
   with self.assertRaisesRegexp(ValueError, "Need first dimension of output"):
     gradients.batch_jacobian(y, x, use_pfor=True)

def create_lstm_batch_jacobian(batch_size, state_size, steps):
  inp, output = lstm_model_fn(batch_size, state_size, steps)
  pfor_jacobian = gradients.batch_jacobian(output, inp, use_pfor=True)
  while_jacobian = gradients.batch_jacobian(output, inp, use_pfor=False)
  return pfor_jacobian, while_jacobian

 def test_batch_jacobian_parallel_iterations(self):
   x = constant_op.constant([[1., 2], [3, 4]])
   w = constant_op.constant([[1., 2, 3, 4], [5, 6, 7, 8]])
   y = math_ops.matmul(x, w)
   self.assertAllClose(gradients.batch_jacobian(y, x, parallel_iterations=2),
                       gradients.batch_jacobian(y, x, parallel_iterations=3))