parameters: (model, strategy=None)
return_statement: return KerasTPUModel(cpu_model=cpu_model, strategy=strategy)
docstring: Copy `model` along with weights to the TPU.
  
    Returns a TPU model.
  
    Usage:
    ```
    a = Input(shape=(32,))
    b = Dense(32)(a)
    model = Model(inputs=a, outputs=b)
  
    # If `num_cores_per_host` is greater than one, batch parallelism will be used
    # to run on multiple TPU cores.
    strategy = keras_support.TPUDistributionStrategy(tpu_cluster_resolver)
    model = keras_support.tpu_model(model, strategy)
    model.compile(
        optimizer=tf.compat.v1.train.GradientDescentOptimizer(learning_rate=1.0),
        ...)
    ```
  
    Args:
      model: A `tf.keras.Model` instance.
      strategy: `TPUDistributionStrategy`.  The strategy to use for replicating
        model across multiple TPU cores.
  
    Returns:
      A new `KerasTPUModel` instance.
docstring_summary: Copy `model` along with weights to the TPU.
docstring_tokens: Copy model along with weights to the TPU .
function: def tpu_model(model, strategy=None):
    """Copy `model` along with weights to the TPU.
  
    Returns a TPU model.
  
    Usage:
    ```
    a = Input(shape=(32,))
    b = Dense(32)(a)
    model = Model(inputs=a, outputs=b)
  
    # If `num_cores_per_host` is greater than one, batch parallelism will be used
    # to run on multiple TPU cores.
    strategy = keras_support.TPUDistributionStrategy(tpu_cluster_resolver)
    model = keras_support.tpu_model(model, strategy)
    model.compile(
        optimizer=tf.compat.v1.train.GradientDescentOptimizer(learning_rate=1.0),
        ...)
    ```
  
    Args:
      model: A `tf.keras.Model` instance.
      strategy: `TPUDistributionStrategy`.  The strategy to use for replicating
        model across multiple TPU cores.
  
    Returns:
      A new `KerasTPUModel` instance.
    """
    _validate_shapes(model)
    # TODO(xiejw): Validate TPU model. TPUModel only?
    # TODO(xiejw): Validate replicas. Full or 1. Shall we allow subset?
    # TODO(xiejw): Adds reduction option.
  
    if strategy is None:
      strategy = TPUDistributionStrategy()
    else:
      if not isinstance(strategy, TPUDistributionStrategy):
        raise TypeError(
            '`strategy` must have type `tf.contrib.tpu.TPUDistributionStrategy`. '
            'Got: {}'.format(type(strategy)))
  
    # If the model has already been initialized, grab the optimizer configuration
    # and model weights before entering the TPU session.
    if model.optimizer:
      if (isinstance(model.optimizer, keras_optimizers.Optimizer) and not
          isinstance(model.optimizer, keras_optimizers.TFOptimizer)):
        optimizer_config = model.optimizer.get_config()
      else:
        optimizer_config = None
      model_weights = model.get_weights()
    else:
      model_weights = None
  
    setup_tpu_session(strategy._tpu_cluster_resolver)
  
    # Force initialization of the CPU model in the TPU session.
    cpu_model = models.clone_model(model)
    if model.optimizer:
      cpu_model.compile(
          _clone_optimizer(model.optimizer, optimizer_config),
          model.loss,
          metrics_module.clone_metrics(model._compile_metrics),
          model.loss_weights,
          model.sample_weight_mode,
          metrics_module.clone_metrics(model._compile_weighted_metrics),
      )
  
    if model_weights:
      cpu_model.set_weights(model_weights)
      cpu_model.reset_states()
  
    return KerasTPUModel(cpu_model=cpu_model, strategy=strategy)
function_tokens: def tpu_model ( model , strategy = None ) : _validate_shapes ( model ) # TODO(xiejw): Validate TPU model. TPUModel only? # TODO(xiejw): Validate replicas. Full or 1. Shall we allow subset? # TODO(xiejw): Adds reduction option. if strategy is None : strategy = TPUDistributionStrategy ( ) else : if not isinstance ( strategy , TPUDistributionStrategy ) : raise TypeError ( '`strategy` must have type `tf.contrib.tpu.TPUDistributionStrategy`. ' 'Got: {}' . format ( type ( strategy ) ) ) # If the model has already been initialized, grab the optimizer configuration # and model weights before entering the TPU session. if model . optimizer : if ( isinstance ( model . optimizer , keras_optimizers . Optimizer ) and not isinstance ( model . optimizer , keras_optimizers . TFOptimizer ) ) : optimizer_config = model . optimizer . get_config ( ) else : optimizer_config = None model_weights = model . get_weights ( ) else : model_weights = None setup_tpu_session ( strategy . _tpu_cluster_resolver ) # Force initialization of the CPU model in the TPU session. cpu_model = models . clone_model ( model ) if model . optimizer : cpu_model . compile ( _clone_optimizer ( model . optimizer , optimizer_config ) , model . loss , metrics_module . clone_metrics ( model . _compile_metrics ) , model . loss_weights , model . sample_weight_mode , metrics_module . clone_metrics ( model . _compile_weighted_metrics ) , ) if model_weights : cpu_model . set_weights ( model_weights ) cpu_model . reset_states ( ) return KerasTPUModel ( cpu_model = cpu_model , strategy = strategy )