tf.tensor_scatter_nd_update() doesn't work in custom layer/graph to update a 2D Tensor with indices

[Sawatdatta]

Please make sure that this is a bug. As per our GitHub Policy, we only address code/doc bugs, performance issues, feature requests and build/installation issues on GitHub. tag:bug_template

System information

• Have I written custom code (as opposed to using a stock example script provided in TensorFlow): Yes

• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Linux Ubuntu 18.04

• Mobile device (e.g. iPhone 8, Pixel 2, Samsung Galaxy) if the issue happens on mobile device: NA

• TensorFlow installed from (source or binary): binary

• TensorFlow version (use command below): 2.0

• Python version: 3.7

• Bazel version (if compiling from source):

• GCC/Compiler version (if compiling from source):

• CUDA/cuDNN version:10.0/7.3.0

                                                                **
  

• GPU model and memory: GeForce 940m

You can collect some of this information using our environment capture
script
You can also obtain the TensorFlow version with: 1. TF 1.0: python -c "import tensorflow as tf; print(tf.GIT_VERSION, tf.VERSION)" 2. TF 2.0: python -c "import tensorflow as tf; print(tf.version.GIT_VERSION, tf.version.VERSION)"

Describe the current behavior ValueError: Shape must be at least rank 1 but is rank 0 for 'TensorScatterUpdate' (op: 'TensorScatterUpdate') with input shapes: [64,64], [2], [].

Describe the expected behavior : it should update the tensor given the row, column position

Code to reproduce the issue
The following code is part of the custom layer and written in call function, the layer is not trainable hence I used Tensor instead of tf.Variable, size of sub_sub_mapl Tensor is 64,64, temp_val will be scalar, and m, n are a row & column position to be updated

             ```

temp_val = tf.convert_to_tensor(
tf.add(tf.math.reduce_max(mat[i, j:j + 2]), tf.math.reduce_max(mat[i, l:l + 2])) / 2)
indices = tf.constant([m, n], dtype=tf.int32)
tf.tensor_scatter_nd_update(self.sub_sub_mapl, indices, temp_val)

**Other info / logs**
Include any logs or source code that would be helpful to diagnose the problem. If including tracebacks, please include the full traceback. Large logs and files should be attached.

Traceback (most recent call last):
  File "F:/pom/pom_cust_layer_parallel_v5.py", line 93, in <module>
    tf.keras.layers.Dense(num_classes, activation='softmax')])
  File "F:\tf2venv\lib\site-packages\tensorflow_core\python\training\tracking\base.py", line 457, in _method_wrapper
    result = method(self, *args, **kwargs)
  File "F:\tf2venv\lib\site-packages\tensorflow_core\python\keras\engine\sequential.py", line 114, in __init__
    self.add(layer)
  File "F:\tf2venv\lib\site-packages\tensorflow_core\python\training\tracking\base.py", line 457, in _method_wrapper
    result = method(self, *args, **kwargs)
  File "F:\tf2venv\lib\site-packages\tensorflow_core\python\keras\engine\sequential.py", line 196, in add
    output_tensor = layer(self.outputs[0])
  File "F:\tf2venv\lib\site-packages\tensorflow_core\python\keras\engine\base_layer.py", line 842, in __call__
    outputs = call_fn(cast_inputs, *args, **kwargs)
  File "F:\tf2venv\lib\site-packages\tensorflow_core\python\autograph\impl\api.py", line 237, in wrapper
    raise e.ag_error_metadata.to_exception(e)
ValueError: in converted code:
    relative to F::

    pom/pom_cust_layer_parallel_v5.py:71 call  *
        self.MAPL = tf.map_fn(fn=self.get_avg_pool, elems=input)
    tf2venv\lib\site-packages\tensorflow_core\python\ops\map_fn.py:268 map_fn
        maximum_iterations=n)
    tf2venv\lib\site-packages\tensorflow_core\python\ops\control_flow_ops.py:2675 while_loop
        back_prop=back_prop)
    tf2venv\lib\site-packages\tensorflow_core\python\ops\while_v2.py:198 while_loop
        add_control_dependencies=add_control_dependencies)
    tf2venv\lib\site-packages\tensorflow_core\python\framework\func_graph.py:915 func_graph_from_py_func
        func_outputs = python_func(*func_args, **func_kwargs)
    tf2venv\lib\site-packages\tensorflow_core\python\ops\while_v2.py:176 wrapped_body
        outputs = body(*_pack_sequence_as(orig_loop_vars, args))
    tf2venv\lib\site-packages\tensorflow_core\python\ops\map_fn.py:257 compute
        packed_fn_values = fn(packed_values)
    pom/pom_cust_layer_parallel_v5.py:62 get_avg_pool
        self.sub_mapl = tf.map_fn(fn=self.get_avg_pool_channel, elems=mat_unstacked)
    tf2venv\lib\site-packages\tensorflow_core\python\ops\map_fn.py:268 map_fn
        maximum_iterations=n)
    tf2venv\lib\site-packages\tensorflow_core\python\ops\control_flow_ops.py:2675 while_loop
        back_prop=back_prop)
    tf2venv\lib\site-packages\tensorflow_core\python\ops\while_v2.py:198 while_loop
        add_control_dependencies=add_control_dependencies)
    tf2venv\lib\site-packages\tensorflow_core\python\framework\func_graph.py:915 func_graph_from_py_func
        func_outputs = python_func(*func_args, **func_kwargs)
    tf2venv\lib\site-packages\tensorflow_core\python\ops\while_v2.py:176 wrapped_body
        outputs = body(*_pack_sequence_as(orig_loop_vars, args))
    tf2venv\lib\site-packages\tensorflow_core\python\ops\map_fn.py:257 compute
        packed_fn_values = fn(packed_values)
    pom/pom_cust_layer_parallel_v5.py:50 get_avg_pool_channel
        tf.tensor_scatter_nd_update(self.sub_sub_mapl, indices, temp_val)
    tf2venv\lib\site-packages\tensorflow_core\python\ops\gen_array_ops.py:11087 tensor_scatter_update
        updates=updates, name=name)
    tf2venv\lib\site-packages\tensorflow_core\python\framework\op_def_library.py:793 _apply_op_helper
        op_def=op_def)
    tf2venv\lib\site-packages\tensorflow_core\python\framework\func_graph.py:548 create_op
        compute_device)
    tf2venv\lib\site-packages\tensorflow_core\python\framework\ops.py:3429 _create_op_internal
        op_def=op_def)
    tf2venv\lib\site-packages\tensorflow_core\python\framework\ops.py:1773 __init__
        control_input_ops)
    tf2venv\lib\site-packages\tensorflow_core\python\framework\ops.py:1613 _create_c_op
        raise ValueError(str(e))

[gadagashwini-zz]

@Sawatdatta, Please provide the complete standalone code to reproduce the reported issue. Thanks!

[gadagashwini-zz]

@Sawatdatta, Looks like code is incomplete. Provide us the full code snippet. It will indeed help us to move faster. Thanks!

[gadagashwini-zz]

Issue is replicating with Tf 2.0.
Please take a look at gist. Thanks!

[Sawatdatta]

Any updates on this issue?