Cartesian Product in Tensorflow

后端 未结 4 491
醉梦人生
醉梦人生 2020-12-02 01:56

Is there any easy way to do cartesian product in Tensorflow like itertools.product? I want to get combination of elements of two tensors (a and b),

4条回答
  •  臣服心动
    2020-12-02 02:21

    I'm going to assume here that both a and b are 1-D tensors.

    To get the cartesian product of the two, I would use a combination of tf.expand_dims and tf.tile:

    a = tf.constant([1,2,3]) 
    b = tf.constant([4,5,6,7]) 
    
    tile_a = tf.tile(tf.expand_dims(a, 1), [1, tf.shape(b)[0]])  
    tile_a = tf.expand_dims(tile_a, 2) 
    tile_b = tf.tile(tf.expand_dims(b, 0), [tf.shape(a)[0], 1]) 
    tile_b = tf.expand_dims(tile_b, 2) 
    
    cartesian_product = tf.concat([tile_a, tile_b], axis=2) 
    
    cart = tf.Session().run(cartesian_product) 
    
    print(cart.shape) 
    print(cart) 
    

    You end up with a len(a) * len(b) * 2 tensor where each combination of the elements of a and b is represented in the last dimension.

提交回复
热议问题