parameters: (optimizer, use_bias=False, use_callable_loss=True)
return_statement: return model_fn, dataset_fn, layer
docstring: Example of non-distribution-aware legacy code.
docstring_summary: Example of non-distribution-aware legacy code.
docstring_tokens: Example of non - distribution - aware legacy code .
function: def minimize_loss_example(optimizer, use_bias=False, use_callable_loss=True):
    """Example of non-distribution-aware legacy code."""
  
    def dataset_fn():
      dataset = dataset_ops.Dataset.from_tensors([[1.]]).repeat()
      # TODO(isaprykin): batch with drop_remainder causes shapes to be
      # fully defined for TPU.  Remove this when XLA supports dynamic shapes.
      return dataset.batch(1, drop_remainder=True)
  
    layer = core.Dense(1, use_bias=use_bias)
  
    def model_fn(x):
      """A very simple model written by the user."""
  
      def loss_fn():
        y = array_ops.reshape(layer(x), []) - constant_op.constant(1.)
        return y * y
  
      if isinstance(optimizer, optimizer_v2.OptimizerV2):
        return optimizer.minimize(loss_fn, lambda: layer.trainable_variables)
      elif use_callable_loss:
        return optimizer.minimize(loss_fn)
      else:
        return optimizer.minimize(loss_fn())
  
    return model_fn, dataset_fn, layer
function_tokens: def minimize_loss_example ( optimizer , use_bias = False , use_callable_loss = True ) : def dataset_fn ( ) : dataset = dataset_ops . Dataset . from_tensors ( [ [ 1. ] ] ) . repeat ( ) # TODO(isaprykin): batch with drop_remainder causes shapes to be # fully defined for TPU.  Remove this when XLA supports dynamic shapes. return dataset . batch ( 1 , drop_remainder = True ) layer = core . Dense ( 1 , use_bias = use_bias ) def model_fn ( x ) : """A very simple model written by the user.""" def loss_fn ( ) : y = array_ops . reshape ( layer ( x ) , [ ] ) - constant_op . constant ( 1. ) return y * y if isinstance ( optimizer , optimizer_v2 . OptimizerV2 ) : return optimizer . minimize ( loss_fn , lambda : layer . trainable_variables ) elif use_callable_loss : return optimizer . minimize ( loss_fn ) else : return optimizer . minimize ( loss_fn ( ) ) return model_fn , dataset_fn , layer