code: def _test_finding_sparse_solution(self, batch_shape=None):
      # Test that Proximal Hessian descent prefers sparse solutions when
      # l1_regularizer is large enough.
      #
      # Define
      #
      #     Loss(x) := (x[0] - a[0])**2 + epsilon * sum(
      #                    (x[i] - a[i])**2 for i in range(1, n))
      #
      # where `a` is a constant and epsilon is small.  Set
      # l2_regularizer = 0 and set l1_regularizer such that
      #
      #     epsilon << l1_regularizer << 1.
      #
      # L1 regularization should cause the computed optimum to have zeros in all
      # but the 0th coordinate: optimal_x ~= [a[0], 0, ..., 0].
      n = 10
      epsilon = 1e-6
      if batch_shape is None:
          batch_shape = []
      # Set a[0] explicitly to make sure it's not very close to zero
      a0 = 6.0
      a_ = np.concatenate(
          [
              np.full(batch_shape + [1], a0),
              np.random.random(size=batch_shape + [n - 1]),
          ],
          axis=-1,
      )
      a = self._adjust_dtype_and_shape_hints(a_)
  
      def _grad_and_hessian_unregularized_loss_fn(x):
          diff = x - a
          grad = 2.0 * tf.concat(
              [diff[..., :1], epsilon * diff[..., 1:]], axis=-1
          )
  
          hessian_outer = tf.SparseTensor(
              indices=[
                  b + (i, i)
                  for i in range(n)
                  for b in np.ndindex(*batch_shape)
              ],
              values=tf.ones(
                  shape=[np.prod(batch_shape, dtype=int) * n],
                  dtype=self.dtype,
              ),
              dense_shape=batch_shape + [n, n],
          )
  
          hessian_middle_per_batch = 2 * tf.concat(
              [[1.0], epsilon * tf.ones([n - 1], dtype=self.dtype)],
              axis=0,
          )
          hessian_middle = (
              tf.zeros(batch_shape + [n], dtype=self.dtype)
              + hessian_middle_per_batch
          )
          return grad, hessian_outer, hessian_middle
  
      w, is_converged, num_iter = proximal_hessian_sparse.minimize(
          _grad_and_hessian_unregularized_loss_fn,
          x_start=tf.zeros(batch_shape + [n], dtype=self.dtype),
          l1_regularizer=1e-2,
          l2_regularizer=None,
          maximum_iterations=10,
          maximum_full_sweeps_per_iteration=10,
          tolerance=1e-5,
          learning_rate=1.0,
      )
  
      w_, is_converged_, _ = self.evaluate([w, is_converged, num_iter])
  
      expected_w = tf.concat(
          [a[..., :1], tf.zeros(batch_shape + [n - 1], self.dtype)],
          axis=-1,
      )
  
      # Using atol=0 ensures that w must be exactly zero in all coordinates
      # where expected_w is exactly zero.
      self.assertAllEqual(is_converged_, True)
      self.assertAllClose(w_, expected_w, atol=0.0, rtol=1e-3)
  
imports: import numpy
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.optimizer import proximal_hessian_sparse
label_desc: Use image node_id 3 for calling the _ProximalHessianTest obj's underlying member method code with example usage: obj._test_finding_sparse_solution(batch_shape) and returns: grad,  hessian_outer,  hessian_middle