from tensorflow.python.framework import ops

    for input_tensor in inputs:
      op = state_ops.assign_add(var, input_tensor, use_locking=True)
      update_ops.append(op)
    with ops.control_dependencies(update_ops):
      return gen_state_ops._destroy_temporary_variable(var,
                                                       var_name=var_name,
                                                       name=name)


@ops.RegisterShape("BatchMatMul")
def _BatchMatMulShape(op):
  """Shape function for BatchMatMul op."""
  a_shape = op.inputs[0].get_shape()
  adj_a = op.get_attr("adj_x")
  b_shape = op.inputs[1].get_shape()
  adj_b = op.get_attr("adj_y")
  if not a_shape.is_fully_defined() or not b_shape.is_fully_defined():
    return [tensor_shape.unknown_shape()]
