docstr: Remove all the model variable for memory saving.
  
  During JAX training, since the training function are stateless, we have
  to pass in and get the model weights over and over, during which the
  copy of the weights that attached to the KerasVariable are still and
  occupying extra memory. We remove those variable to save memory (for
  better memory utilization) at the beginning of the epoch, and reattach
  the value back to variables at the end of the epoch, via
  `jax_state_sync()`.
docstr_tok: ["Remove","all","the","model","variable","for","memory","saving",".","During","JAX","training",",","since","the","training","function","are","stateless",",","we","have","to","pass","in","and","get","the","model","weights","over","and","over",",","during","which","the","copy","of","the","weights","that","attached","to","the","KerasVariable","are","still","and","occupying","extra","memory",".","We","remove","those","variable","to","save","memory","(","for","better","memory","utilization",")","at","the","beginning","of","the","epoch",",","and","reattach","the","value","back","to","variables","at","the","end","of","the","epoch",",","via","`","jax_state_sync","(",")","`","."]
code: def _purge_model_variables(
      self,
      trainable_variables=True,
      non_trainable_variables=True,
      optimizer_variables=True,
      metric_variables=True,
  ):
      """Remove all the model variable for memory saving.
  
      During JAX training, since the training function are stateless, we have
      to pass in and get the model weights over and over, during which the
      copy of the weights that attached to the KerasVariable are still and
      occupying extra memory. We remove those variable to save memory (for
      better memory utilization) at the beginning of the epoch, and reattach
      the value back to variables at the end of the epoch, via
      `jax_state_sync()`.
      """
      if trainable_variables:
          for v in self.trainable_variables:
              v._value = None
      if non_trainable_variables:
          for v in self.non_trainable_variables:
              v._value = None
      if optimizer_variables:
          for v in self.optimizer.variables:
              v._value = None
      if metric_variables:
          for v in self.metrics_variables:
              v._value = None
  
imports: import collections
  import itertools
  from functools import partial
  import jax
  import numpy
  import tree
  from keras import backend
  from keras import callbacks
  from keras import ops
  from keras import optimizers
  from keras.backend import distribution_lib
  from keras.distribution import distribution_lib
  from keras.trainers import trainer
  from keras.trainers.data_adapters import data_adapter_utils
  from keras.trainers.epoch_iterator import EpochIterator
  from keras.utils import traceback_utils
label_desc: Use image node_id 18 for calling the JAXTrainer obj's underlying member method code with example usage: obj._purge_model_variables(trainable_variables, non_trainable_variables, optimizer_variables, metric_variables) without return types