code: def test_polynomial_is_almost_a_polynomial(self):
      poly_bnn = kernels.PolynomialBNN(degree=3)
      init_params = poly_bnn.init(
          jax.random.PRNGKey(0), jnp.ones((10, 1))
      )
  
      # compute power series
      func = lambda x: poly_bnn.apply(init_params, x)[0]
      params = [func(0.0)]
      for _ in range(4):
          func = jax.grad(func)
          params.append(func(0.0))
  
      # Last 4th degree coefficient should be around 0.
      self.assertAlmostEqual(params[-1], 0.0)
  
      # Check that the random initialization is approximately a polynomial by
      # evaluating far away from the expansion.
      x = 17.0
      self.assertAlmostEqual(
          poly_bnn.apply(init_params, x)[0],
          params[0]
          + x * params[1]
          + x**2 * params[2] / 2
          + x**3 * params[3] / 6,
          places=3,
      )
  
imports: from absl.testing import parameterized
  import jax
  import jax.numpy
  import numpy
  from tensorflow_probability.python.experimental.autobnn import kernels
  from tensorflow_probability.python.experimental.autobnn import util
  from tensorflow_probability.substrates.jax.distributions import lognormal
  from absl.testing import absltest
label_desc: Use image node_id 5 for calling the KernelsTest obj's underlying member method code with example usage: obj.test_polynomial_is_almost_a_polynomial() without return types