CODEWITHSUNDEEP
tensorflow
Derk
Derk

Reputation: 1395

Tensorflow equivalent for this MATLAB code

I want to create a Tensor with as uppertriangular part the values from a vector. I have found in MATLAB this can be done with

a = [1 2 3 4 5 6 7 8 9 10];
b = triu(ones(5),1);
b = b'
b(b==1) = a
b = b'

My tensorflow implementation so far

b = tf.matrix_band_part(tf.ones([dim,dim]), 0, -1) # make upper triangular part 1
b = tf.transpose(b)
...
b = tf.transpose(b)

Who can help me?

Upvotes: 0

Views: 290

Answers (1)

Allen Lavoie
Allen Lavoie

Reputation: 5808

I haven't seen a fantastic way to do it, but it's certainly possible. Here's one way (expanding the last dimension of a tensor into a matrix; the preceding dimensions may be batch dimensions):

import tensorflow as tf

def matrix_with_upper_values(upper_values):
  # Check that the input is at least a vector
  upper_values = tf.convert_to_tensor(upper_values)
  upper_values.get_shape().with_rank_at_least(1)
  # Put the batch dimensions last
  upper_values = tf.transpose(
      upper_values,
      tf.concat(0, [[tf.rank(upper_values) - 1],
                    tf.range(tf.rank(upper_values) - 1)]))
  input_shape = tf.shape(upper_values)[0]
  # Compute the size of the matrix that would have this upper triangle
  matrix_size = (1 + tf.cast(tf.sqrt(tf.cast(input_shape * 8 + 1, tf.float32)),
                             tf.int32)) // 2
  # Check that the upper triangle size is valid
  check_size_op = tf.Assert(
      tf.equal(matrix_size ** 2, input_shape * 2 + matrix_size),
      ["Not a valid upper triangle size: ", input_shape])
  with tf.control_dependencies([check_size_op]):
    matrix_size = tf.identity(matrix_size)
  # Compute indices for the whole matrix and the upper diagonal
  index_matrix = tf.reshape(tf.range(matrix_size ** 2),
                            [matrix_size, matrix_size])
  diagonal_indicies = (matrix_size * tf.range(matrix_size)
                       + tf.range(matrix_size))
  upper_triangular_indices, _ = tf.unique(tf.reshape(
      tf.matrix_band_part(
          index_matrix, 0, -1)       # upper triangular part
      - tf.diag(diagonal_indicies),  # remove diagonal
      [-1]))
  batch_dimensions = tf.shape(upper_values)[1:]
  return_shape_transposed = tf.concat(0, [[matrix_size, matrix_size],
                                          batch_dimensions])
  # Fill everything else with zeros; later entries get priority
  # in dynamic_stitch
  result_transposed = tf.reshape(
      tf.dynamic_stitch(
          [index_matrix,
           upper_triangular_indices[1:]],  # discard 0
          [tf.zeros(return_shape_transposed, dtype=upper_values.dtype),
           upper_values]),
      return_shape_transposed)
  # Transpose the batch dimensions to be first again
  return tf.transpose(
      result_transposed,
      tf.concat(0, [tf.range(2, tf.rank(upper_values) + 1), [0, 1]]))

with tf.Session():
  print(matrix_with_upper_values([1]).eval())
  print(matrix_with_upper_values([2,7,1]).eval())
  print(matrix_with_upper_values([3,1,4,1,5,9]).eval())
  print(matrix_with_upper_values([]).eval())
  print(matrix_with_upper_values([[2,7,1],[4,3,5]]).eval())
  print(matrix_with_upper_values(tf.zeros([0, 3])).eval())

Prints:

[[0 1]
 [0 0]]
[[0 2 7]
 [0 0 1]
 [0 0 0]]
[[0 3 1 4]
 [0 0 1 5]
 [0 0 0 9]
 [0 0 0 0]]
[[ 0.]]
[[[0 2 7]
  [0 0 1]
  [0 0 0]]

 [[0 4 3]
  [0 0 5]
  [0 0 0]]]
[]

Upvotes: 1

Related Questions