def Chunk(layer, chunk_size, pass_unchunkable=True):
  """Executes `layer` using batch chunks of size `chunk_size` to save memory."""
  if chunk_size < 1:
    return layer
  def reshape_to_chunks(x):
    chunk_batch = x.shape[0]
    size = chunk_size
    n_chunks = chunk_batch // size
    if chunk_batch % size != 0:
      if pass_unchunkable:
        n_chunks = 1
        size = chunk_batch
      else:
        raise ValueError(f'Chunk size {size} must divide batch '
                         f'size {chunk_batch}')
    return jnp.reshape(x, [n_chunks, size] + list(x.shape[1:]))
  reshape_to_chunks_layer = base.PureLayer(
      lambda xs: fastmath.nested_map(reshape_to_chunks, xs),
      n_in=layer.n_in, n_out=layer.n_in, name='ReshapeToChunks')
  def reshape_from_chunks(x):
    batch_size = x.shape[0] * x.shape[1]
    return jnp.reshape(x, [batch_size] + list(x.shape[2:]))
  reshape_from_chunks_layer = base.PureLayer(
      lambda xs: fastmath.nested_map(reshape_from_chunks, xs),
      n_in=layer.n_out, n_out=layer.n_out, name='ReshapeFromChunks')
  return Serial(
      reshape_to_chunks_layer,
      Scan(layer, axis=0, n_carry=0, remat=True),
      reshape_from_chunks_layer,
  )
Executes `layer` using batch chunks of size `chunk_size` to save memory.
trax/layers/combinators.py