# Path: third_party/bi_att_flow/my/tensorflow/general.py
# def add_wd(wd, scope=None):
#   scope = scope or tf.get_variable_scope().name
#   variables = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, scope=scope)
#   with tf.name_scope('weight_decay'):
#     for var in variables:
#       weight_decay = tf.mul(
#           tf.nn.l2_loss(var), wd, name='{}/wd'.format(var.op.name))
#       tf.add_to_collection('losses', weight_decay)
# 
# Path: third_party/bi_att_flow/my/tensorflow/general.py
# def exp_mask(val, mask, name=None):
#   """Give very negative number to unmasked elements in val.
# 
#     For example, [-3, -2, 10], [True, True, False] -> [-3, -2, -1e9].
#     Typically, this effectively masks in exponential space (e.g. softmax)
#     Args:
#         val: values to be masked
#         mask: masking boolean tensor, same shape as tensor
#         name: name for output tensor
# 
#     Returns:
#         Same shape as val, where some elements are very small (exponentially
#         zero)
#     """
#   if name is None:
#     name = 'exp_mask'
#   return tf.add(
#       val, (1 - tf.cast(mask, 'float')) * VERY_NEGATIVE_NUMBER, name=name)
# 
# Path: third_party/bi_att_flow/my/tensorflow/general.py
# def flatten(tensor, keep):
#   fixed_shape = tensor.get_shape().as_list()
#   start = len(fixed_shape) - keep
#   left = reduce(mul,
#                 [fixed_shape[i] or tf.shape(tensor)[i] for i in range(start)])
#   out_shape = [left] + [
#       fixed_shape[i] or tf.shape(tensor)[i]
#       for i in range(start, len(fixed_shape))
#   ]
#   flat = tf.reshape(tensor, out_shape)
#   return flat
# 
# Path: third_party/bi_att_flow/my/tensorflow/general.py
# def reconstruct(tensor, ref, keep):
#   ref_shape = ref.get_shape().as_list()
#   tensor_shape = tensor.get_shape().as_list()
#   ref_stop = len(ref_shape) - keep
#   tensor_start = len(tensor_shape) - keep
#   pre_shape = [ref_shape[i] or tf.shape(ref)[i] for i in range(ref_stop)]
#   keep_shape = [
#       tensor_shape[i] or tf.shape(tensor)[i]
#       for i in range(tensor_start, len(tensor_shape))
#   ]
#   # pre_shape = [tf.shape(ref)[i] for i in range(len(ref.get_shape().as_list()[:-keep]))]
#   # keep_shape = tensor.get_shape().as_list()[-keep:]
#   target_shape = pre_shape + keep_shape
#   out = tf.reshape(tensor, target_shape)
#   return out


import tensorflow as tf
from third_party.bi_att_flow.my.tensorflow.general import add_wd
from third_party.bi_att_flow.my.tensorflow.general import exp_mask
from third_party.bi_att_flow.my.tensorflow.general import flatten
from third_party.bi_att_flow.my.tensorflow.general import reconstruct
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import init_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import variable_scope as vs
from tensorflow.python.util import nest
        tf.cond(is_train, lambda: tf.nn.dropout(arg, input_keep_prob),
                lambda: arg) for arg in flat_args
    ]

  flat_out = _linear(
      flat_args, output_size, bias, bias_start=bias_start, scope=scope)
  out = reconstruct(flat_out, args[0], 1)
  if squeeze:
    out = tf.squeeze(out, [len(args[0].get_shape().as_list()) - 1])
  if wd:
    add_wd(wd)

  return out


def dropout(x, keep_prob, is_train, noise_shape=None, seed=None, name=None):
  with tf.name_scope(name or "dropout"):
    if keep_prob < 1.0:
      d = tf.nn.dropout(x, keep_prob, noise_shape=noise_shape, seed=seed)
      out = tf.cond(is_train, lambda: d, lambda: x)
      return out
    return x


def softmax(logits, mask=None, scope=None):
  with tf.name_scope(scope or "Softmax"):
    if mask is not None:
      logits = exp_mask(logits, mask)
    flat_logits = flatten(logits, 1)
    flat_out = tf.nn.softmax(flat_logits)

# Path: third_party/bi_att_flow/my/tensorflow/general.py
# def add_wd(wd, scope=None):
#   scope = scope or tf.get_variable_scope().name
#   variables = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, scope=scope)
#   with tf.name_scope('weight_decay'):
#     for var in variables:
#       weight_decay = tf.mul(
#           tf.nn.l2_loss(var), wd, name='{}/wd'.format(var.op.name))
#       tf.add_to_collection('losses', weight_decay)
# 
# Path: third_party/bi_att_flow/my/tensorflow/general.py
# def exp_mask(val, mask, name=None):
#   """Give very negative number to unmasked elements in val.
# 
#     For example, [-3, -2, 10], [True, True, False] -> [-3, -2, -1e9].
#     Typically, this effectively masks in exponential space (e.g. softmax)
#     Args:
#         val: values to be masked
#         mask: masking boolean tensor, same shape as tensor
#         name: name for output tensor
# 
#     Returns:
#         Same shape as val, where some elements are very small (exponentially
#         zero)
#     """
#   if name is None:
#     name = 'exp_mask'
#   return tf.add(
#       val, (1 - tf.cast(mask, 'float')) * VERY_NEGATIVE_NUMBER, name=name)
# 
# Path: third_party/bi_att_flow/my/tensorflow/general.py
# def flatten(tensor, keep):
#   fixed_shape = tensor.get_shape().as_list()
#   start = len(fixed_shape) - keep
#   left = reduce(mul,
#                 [fixed_shape[i] or tf.shape(tensor)[i] for i in range(start)])
#   out_shape = [left] + [
#       fixed_shape[i] or tf.shape(tensor)[i]
#       for i in range(start, len(fixed_shape))
#   ]
#   flat = tf.reshape(tensor, out_shape)
#   return flat
# 
# Path: third_party/bi_att_flow/my/tensorflow/general.py
# def reconstruct(tensor, ref, keep):
#   ref_shape = ref.get_shape().as_list()
#   tensor_shape = tensor.get_shape().as_list()
#   ref_stop = len(ref_shape) - keep
#   tensor_start = len(tensor_shape) - keep
#   pre_shape = [ref_shape[i] or tf.shape(ref)[i] for i in range(ref_stop)]
#   keep_shape = [
#       tensor_shape[i] or tf.shape(tensor)[i]
#       for i in range(tensor_start, len(tensor_shape))
#   ]
#   # pre_shape = [tf.shape(ref)[i] for i in range(len(ref.get_shape().as_list()[:-keep]))]
#   # keep_shape = tensor.get_shape().as_list()[-keep:]
#   target_shape = pre_shape + keep_shape
#   out = tf.reshape(tensor, target_shape)
#   return out


# Path: third_party/bi_att_flow/my/tensorflow/nn.py
import tensorflow as tf
from third_party.bi_att_flow.my.tensorflow.general import add_wd
from third_party.bi_att_flow.my.tensorflow.general import exp_mask
from third_party.bi_att_flow.my.tensorflow.general import flatten
from third_party.bi_att_flow.my.tensorflow.general import reconstruct
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import init_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import variable_scope as vs
from tensorflow.python.util import nest

        tf.cond(is_train, lambda: tf.nn.dropout(arg, input_keep_prob),
                lambda: arg) for arg in flat_args
    ]

  flat_out = _linear(
      flat_args, output_size, bias, bias_start=bias_start, scope=scope)
  out = reconstruct(flat_out, args[0], 1)
  if squeeze:
    out = tf.squeeze(out, [len(args[0].get_shape().as_list()) - 1])
  if wd:
    add_wd(wd)

  return out


def dropout(x, keep_prob, is_train, noise_shape=None, seed=None, name=None):
  with tf.name_scope(name or "dropout"):
    if keep_prob < 1.0:
      d = tf.nn.dropout(x, keep_prob, noise_shape=noise_shape, seed=seed)
      out = tf.cond(is_train, lambda: d, lambda: x)
      return out
    return x


def softmax(logits, mask=None, scope=None):
  with tf.name_scope(scope or "Softmax"):
    if mask is not None:
      logits = exp_mask(logits, mask)
    flat_logits = flatten(logits, 1)
    flat_out = tf.nn.softmax(flat_logits)

    out = reconstruct(flat_out, logits, 1)