def test_chunk_grad_memory(self):
    """Test chunking gradient here to exercise accelerator memory usage."""
    layer = tl.Serial(tl.Dense(1024*1024), tl.Dense(24))
    chunked = tl.Chunk(layer, 256)

    @fastmath.jit
    def mock_training_step(x, weights, state, rng):
      def compute_mock_loss(weights):
        logits, new_state = chunked.pure_fn(x, weights, state, rng)
        loss = fastmath.numpy.mean(logits)
        return loss, (new_state, logits)
      gradients, (new_state, logits) = fastmath.grad(
          compute_mock_loss, has_aux=True)(weights)
      new_weights = fastmath.nested_map_multiarg(
          lambda w, g: w - 1e-4 * g, weights, gradients)
      return new_weights, new_state, logits

    x = np.random.uniform(size=(32*1024, 16))
    chunked.init(shapes.signature(x))
    weights, _, logits = mock_training_step(
        x, chunked.weights, chunked.state, fastmath.random.get_prng(0))
    self.assertEqual(logits.shape, (32*1024, 24))
    self.assertEqual(weights[1][0][0][0].shape, (16, 1024*1024))
Test chunking gradient here to exercise accelerator memory usage.
trax/layers/acceleration_test.py