code: def assign(var, new_val, assign_fn=assign_slice):
    
    if isinstance(var, Tensor):
      var = var.operation
    if not isinstance(var, Variable):
      raise ValueError("var must be a mtf.Variable or its output Tensor.")
    return Assign([var], [new_val], assign_fn=assign_fn)
docs: Assign a new value to a variable.
  
    Args:
      var: either a Variable operation or its output Tensor.
      new_val: a Tensor
      assign_fn: a function from
          (mtf.Variable, tf.Variable, tf.Tensor) -> tf.Operation
    Returns:
      an Operation
    Raises:
      ValueError: if var is not a Variable and var.operation is not a Variable