docstr: Test chunking here to exercise accelerator memory usage.
docstr_tok: ["Test","chunking","here","to","exercise","accelerator","memory","usage","."]
code: def test_chunk_memory(self):
      """Test chunking here to exercise accelerator memory usage."""
      layer = tl.Serial(tl.Dense(1024 * 1024), tl.Dense(128))
      chunked = tl.Chunk(layer, 256)
      x = np.random.uniform(size=(16 * 1024, 16))
      chunked.init(shapes.signature(x))
      y = chunked(x)
      z = tl.Accelerate(chunked)(x)
      self.assertEqual(y.shape, (16 * 1024, 128))
      self.assertEqual(z.shape, (16 * 1024, 128))
  
file_path: trax/trax/layers/acceleration_test.py
imports: from absl.testing import absltest
  from jax import config
  import numpy
  from trax import fastmath
  from trax import layers
  from trax import shapes
label_desc: Use image node_id 4 for calling the AccelerationTest obj's underlying member method code with example usage: obj.test_chunk_memory() without return types