TensorFlowで低レベルな操作を行いたい

前提・実現したいこと

可変のバッチサイズを取得し，それを用いて新たな形状のテンソルを作成後，各要素に代入処理を行いたいです．
ソースコードは簡単化のために書き換えています．
PyTorchで実装したソースコードを追記しました．
これをTensorFlowで実装したいです．
よろしくお願いします．

該当のソースコード

Python
1import tensorflow as tf
2
3def func(x):
4	# バッチサイズを取得（可変）
5	# nbatch = x.get_shape().as_list()[0]
6	nbatch = tf.shape(x)[0]
7	# 任意の形状のテンソル作成
8	# output = tf.zeros([nbatch, 2], dtype=x.dtype)
9	output = tf.Variable(initial_value=tf.zeros([nbatch, 2], dtype=x.dtype), trainable=False)
10	# 作成したテンソルの任意の要素に代入
11	# output[0, 0] = x[0, 0]
12	output = tf.assign(output[0, 0], x[0, 0])
13	output = tf.assign(output[0, 1], x[0, 1])
14	output = tf.assign(output[1, 0], x[0, 2])
15	output = tf.assign(output[1, 1], x[1, 0])
16	output = tf.assign(output[2, 0], x[1, 1])
17	output = tf.assign(output[2, 1], x[1, 2])
18	return output
19
20if __name__ == '__main__':
21	holder = tf.placeholder(tf.float32, [None, 3])
22	op = func(holder)
23	with tf.Session() as sess:
24		sess.run(tf.global_variables_initializer())
25		x = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
26		result = sess.run(op, feed_dict={holder: x})
27		print(result)

Python
1import torch
2
3def func(x):
4	# バッチサイズを取得（可変）
5	nbatch = x.size(0)
6	# 任意の形状のテンソル作成
7	output = x.new_empty((nbatch, 2))
8	# 作成したテンソルの任意の要素に代入
9	output[0, 0] = x[0, 0]
10	output[0, 1] = x[0, 1]
11	output[1, 0] = x[0, 2]
12	output[1, 1] = x[1, 0]
13	output[2, 0] = x[1, 1]
14	output[2, 1] = x[1, 2]
15	return output
16
17if __name__ == '__main__':
18	x = torch.Tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
19	result = func(x)
20	print(result)

発生している問題・エラーメッセージ

Traceback (most recent call last):
  File "run.py", line 23, in <module>
    op = func(holder)
  File "run.py", line 10, in func
    output = tf.Variable(tf.zeros([nbatch, 2]), trainable=False)
  File "C:\Anaconda3\Lib\site-packages\tensorflow_core\python\ops\variables.py", line 258, in __call__
    return cls._variable_v1_call(*args, **kwargs)
  File "C:\Anaconda3\Lib\site-packages\tensorflow_core\python\ops\variables.py", line 219, in _variable_v1_call
    shape=shape)
  File "C:\Anaconda3\Lib\site-packages\tensorflow_core\python\ops\variables.py", line 197, in <lambda>
    previous_getter = lambda **kwargs: default_variable_creator(None, **kwargs)
  File "C:\Anaconda3\Lib\site-packages\tensorflow_core\python\ops\variable_scope.py", line 2519, in default_variable_creator
    shape=shape)
  File "C:\Anaconda3\Lib\site-packages\tensorflow_core\python\ops\variables.py", line 262, in __call__
    return super(VariableMetaclass, cls).__call__(*args, **kwargs)
  File "C:\Anaconda3\Lib\site-packages\tensorflow_core\python\ops\variables.py", line 1688, in __init__
    shape=shape)
  File "C:\Anaconda3\Lib\site-packages\tensorflow_core\python\ops\variables.py", line 1853, in _init_from_args
    self._initial_value)
ValueError: initial_value must have a shape specified: Tensor("zeros:0", shape=(?, 2), dtype=float32)