tensorflow_probability/python/mcmc/internal/util.py
def _choose_base_case(is_accepted,
                      accepted,
                      rejected,
                      name=None):
  """Helper to `choose` which expand_dims `is_accepted` and applies tf.where."""
  def _expand_is_accepted_like(x):
    """Helper to expand `is_accepted` like the shape of some input arg."""
    with tf.compat.v1.name_scope('expand_is_accepted_like'):
      expand_shape = tf.concat([
          tf.shape(input=is_accepted),
          tf.ones([tf.rank(x) - tf.rank(is_accepted)], dtype=tf.int32),
      ],
                               axis=0)
      multiples = tf.concat([
          tf.ones([tf.rank(is_accepted)], dtype=tf.int32),
          tf.shape(input=x)[tf.rank(is_accepted):],
      ],
                            axis=0)
      m = tf.tile(tf.reshape(is_accepted, expand_shape),
                  multiples)
      m.set_shape(m.shape.merge_with(x.shape))
      return m
  def _where(accepted, rejected):
    if accepted is rejected:
      return accepted
    accepted = tf.convert_to_tensor(value=accepted, name='accepted')
    rejected = tf.convert_to_tensor(value=rejected, name='rejected')
    r = tf.where(_expand_is_accepted_like(accepted), accepted, rejected)
    r.set_shape(r.shape.merge_with(accepted.shape.merge_with(rejected.shape)))
    return r

  with tf.compat.v1.name_scope(
      name, 'choose', values=[is_accepted, accepted, rejected]):
    if not is_list_like(accepted):
      return _where(accepted, rejected)
    return [(choose(is_accepted, a, r, name=name) if is_namedtuple_like(a)
             else _where(a, r))
            for a, r in zip(accepted, rejected)]
def _choose_base_case(is_accepted,
                      accepted,
                      rejected,
                      name=None):
  """Helper to `choose` which expand_dims `is_accepted` and applies tf.where."""
  def _expand_is_accepted_like(x):
    """Helper to expand `is_accepted` like the shape of some input arg."""
    with tf.compat.v1.name_scope('expand_is_accepted_like'):
      expand_shape = tf.concat([
          tf.shape(input=is_accepted),
          tf.ones([tf.rank(x) - tf.rank(is_accepted)], dtype=tf.int32),
      ],
                               axis=0)
      multiples = tf.concat([
          tf.ones([tf.rank(is_accepted)], dtype=tf.int32),
          tf.shape(input=x)[tf.rank(is_accepted):],
      ],
                            axis=0)
      m = tf.tile(tf.reshape(is_accepted, expand_shape),
                  multiples)
      m.set_shape(m.shape.merge_with(x.shape))
      return m
  def _where(accepted, rejected):
    if accepted is rejected:
      return accepted
    accepted = tf.convert_to_tensor(value=accepted, name='accepted')
    rejected = tf.convert_to_tensor(value=rejected, name='rejected')
    r = tf.where(_expand_is_accepted_like(accepted), accepted, rejected)
    r.set_shape(r.shape.merge_with(accepted.shape.merge_with(rejected.shape)))
    return r

  with tf.compat.v1.name_scope(
      name, 'choose', values=[is_accepted, accepted, rejected]):
    if not is_list_like(accepted):
      return _where(accepted, rejected)
    return [(choose(is_accepted, a, r, name=name) if is_namedtuple_like(a)
             else _where(a, r))
            for a, r in zip(accepted, rejected)]
def _choose_base_case ( is_accepted , accepted , rejected , name = None ) : def _expand_is_accepted_like ( x ) : """Helper to expand `is_accepted` like the shape of some input arg.""" with tf . compat . v1 . name_scope ( 'expand_is_accepted_like' ) : expand_shape = tf . concat ( [ tf . shape ( input = is_accepted ) , tf . ones ( [ tf . rank ( x ) - tf . rank ( is_accepted ) ] , dtype = tf . int32 ) , ] , axis = 0 ) multiples = tf . concat ( [ tf . ones ( [ tf . rank ( is_accepted ) ] , dtype = tf . int32 ) , tf . shape ( input = x ) [ tf . rank ( is_accepted ) : ] , ] , axis = 0 ) m = tf . tile ( tf . reshape ( is_accepted , expand_shape ) , multiples ) m . set_shape ( m . shape . merge_with ( x . shape ) ) return m def _where ( accepted , rejected ) : if accepted is rejected : return accepted accepted = tf . convert_to_tensor ( value = accepted , name = 'accepted' ) rejected = tf . convert_to_tensor ( value = rejected , name = 'rejected' ) r = tf . where ( _expand_is_accepted_like ( accepted ) , accepted , rejected ) r . set_shape ( r . shape . merge_with ( accepted . shape . merge_with ( rejected . shape ) ) ) return r with tf . compat . v1 . name_scope ( name , 'choose' , values = [ is_accepted , accepted , rejected ] ) : if not is_list_like ( accepted ) : return _where ( accepted , rejected ) return [ ( choose ( is_accepted , a , r , name = name ) if is_namedtuple_like ( a ) else _where ( a , r ) ) for a , r in zip ( accepted , rejected ) ]
Helper to `choose` which expand_dims `is_accepted` and applies tf.where.
Helper to choose which expand_dims is_accepted and applies tf . where .