def representation_fn(*, flax_model: nn.Module, train_state: train_utils.TrainState, batch: bert_base_model.Batch, representation_layer: str, gather_to_host: bool=True) -> Tuple[(jnp.ndarray, jnp.ndarray, jnp.ndarray)]:
    'Feeds the inputs to the model and returns their representations.\n\n  Args:\n    flax_model: A Flax model.\n    train_state: TrainState, the state of training including the current\n      global_step, model_state, rng, and optimizer. The buffer of this argument\n      can be donated to the computation.\n    batch: A single batch of data from the dataset.\n    representation_layer: The name of the layer to use as the representation.\n    gather_to_host: Whether to gather results from all devices to the host,\n      rather than leaving them distributed.\n\n  Returns:\n    Representation learned by the model for the given inputs and the labels and\n    masks. If `gather_to_host` is True, these are collected from all hosts.\n  '
    variables = {'params': train_state.optimizer.target, **train_state.model_state}
    representation_layer_parts = representation_layer.split('/')
    filter_rep = (lambda mdl, _: (mdl.name == representation_layer_parts[(- 1)]))
    (_, model_state) = flax_model.apply(variables, batch, train=False, capture_intermediates=filter_rep, mutable=['intermediates'], transfer_mode=True, debug=False)
    if ('intermediates' not in model_state):
        raise ValueError(f'Layer with name "{representation_layer}" does not exist in your model.')
    representation = model_state['intermediates']
    for rep_layer in representation_layer_parts:
        if rep_layer:
            representation = representation[rep_layer]
    representation = representation['__call__'][0]
    if gather_to_host:
        representation = jax.lax.all_gather(representation, 'batch')
        batch = jax.lax.all_gather(batch, 'batch')
    return (representation, batch['label'], batch['batch_mask'])
Feeds the inputs to the model and returns their representations.

Args:
  flax_model: A Flax model.
  train_state: TrainState, the state of training including the current
    global_step, model_state, rng, and optimizer. The buffer of this argument
    can be donated to the computation.
  batch: A single batch of data from the dataset.
  representation_layer: The name of the layer to use as the representation.
  gather_to_host: Whether to gather results from all devices to the host,
    rather than leaving them distributed.

Returns:
  Representation learned by the model for the given inputs and the labels and
  masks. If `gather_to_host` is True, these are collected from all hosts.
def representation_fn(*, flax_model: nn.Module, train_state: train_utils.TrainState, batch: bert_base_model.Batch, representation_layer: str, gather_to_host: bool=True) -> Tuple[(jnp.ndarray, jnp.ndarray, jnp.ndarray)]:
    'Feeds the inputs to the model and returns their representations.\n\n  Args:\n    flax_model: A Flax model.\n    train_state: TrainState, the state of training including the current\n      global_step, model_state, rng, and optimizer. The buffer of this argument\n      can be donated to the computation.\n    batch: A single batch of data from the dataset.\n    representation_layer: The name of the layer to use as the representation.\n    gather_to_host: Whether to gather results from all devices to the host,\n      rather than leaving them distributed.\n\n  Returns:\n    Representation learned by the model for the given inputs and the labels and\n    masks. If `gather_to_host` is True, these are collected from all hosts.\n  '
    variables = {'params': train_state.optimizer.target, **train_state.model_state}
    representation_layer_parts = representation_layer.split('/')
    filter_rep = (lambda mdl, _: (mdl.name == representation_layer_parts[(- 1)]))
    (_, model_state) = flax_model.apply(variables, batch, train=False, capture_intermediates=filter_rep, mutable=['intermediates'], transfer_mode=True, debug=False)
    if ('intermediates' not in model_state):
        raise ValueError(f'Layer with name "{representation_layer}" does not exist in your model.')
    representation = model_state['intermediates']
    for rep_layer in representation_layer_parts:
        if rep_layer:
            representation = representation[rep_layer]
    representation = representation['__call__'][0]
    if gather_to_host:
        representation = jax.lax.all_gather(representation, 'batch')
        batch = jax.lax.all_gather(batch, 'batch')
    return (representation, batch['label'], batch['batch_mask'])