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))
Test chunking here to exercise accelerator memory usage.
trax/layers/acceleration_test.py