"unknown rank error" in tensorflow.keras.layers.Layer with tensorflow.py_function output
I was able to reproduce this issue in TF 2.11 and TF Nightly 2.13.0-dev20230219 . Please find the gist [here](https://colab.research.google.com/gist/pjpratik/20cd82f24b7697da535ba222cd066b28/59726.ipynb). 

@SuryanarayanaY Could you please look into this issue. Thanks.  Hi @MegaCreater ,
The error might be due to python function using inside the Graph for which the shape is unknown.Hence getting the error:Unknown rank error.
The proposed workaround is you have add this code `inputs.set_shape([None,32,32,3])` after tf.nupmy_function like below.
@tf.function
```
def sample(inputs:tf.Tensor):                                                   # some sample function 

  def numpy_method(inputs:np.ndarray):
    # some complex method 
    return np.sqrt(inputs)# example code 

  inputs=tf.numpy_function(numpy_method,[inputs],inputs.dtype,name='outputs')
  inputs.set_shape([None,32,32,3])                                #New line of code
  inputs=tf.keras.layers.AvgPool2D()(inputs)                                   

  return inputs

sample(tf.random.uniform((1,32,32,3)))
```

Please refer to the attached gist [here](https://colab.research.google.com/gist/SuryanarayanaY/5c64d66b14b62ca6308674e4b80f6a13/59726.ipynb).

Hope this will solve your purpose. Thanks! @SuryanarayanaY thanks a lot. But this I already know. This solution I have already mentioned here - https://github.com/tensorflow/tensorflow/issues/37193#issuecomment-1434693519.  But this issue must be resolved. This issue is very old as stated here - #37193  but till now not solved. Adding `@tf.function` created issue only when with `tf.numpy_function` why cannot be redefine this `tf.numpy_function` function.  Can we update something here -> https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/ops/script_ops.py#L768 
somthing like this -> 

```python
# ...... 
def numpy_function(func, inp, Tout, outputs_shapes:tuple=None, stateful=True, name=None):

  ''' docs ------ 
  '''

  outputs = py_func_common(func, inp, Tout, stateful=stateful, name=name)# line 768
  if isinstance(outputs,tf.Tensor) and (outputs.shape==None):
    if outputs_shapes==None: 
      if isinstance(inp,(list,tuple)): outputs.set_shape(inp[0].shape)
      else: outputs.set_shape(inp.shape)
    else: outputs.set_shape(outputs_shapes)
  #elif isinstance(outputs,(list,tuple)): ..... 
  return outputs

# ..... 
``` Hi @MegaCreater ,

I have tested the proposed changes and seems working fine for this case.Please refer to attached tested [gist](https://colab.research.google.com/gist/SuryanarayanaY/f032311e9975716a69fc06d8ba0c06a0/59726_r3-code-gtest.ipynb). At least the changes works fine with and without @tf.function. I am not sure whether this covers all cases or not. Also it involves the structural changes in API as it adds new argument also.This may needs to be reviewed properly for short falls if any.

The proposed changes relies completely on input_shape which we are assuming that output_shape also being same as input_shape. What about the case when there is change in shapes that actually happens inside the numpy_method ?

Could we cover this to all use cases ?  Thanks!
 Yes we can. When the shapes of output is same then user wont have to pass `output_shape` variable (i.e `output_shape=None`, but when `output_shape` changes then `output_shape` has to be provided. Give me some time to work on it more. I will update you. 
But still with this update `rank` issue is not resolved. I am look for that also.   @MegaCreater ,

If `output_shape` shape has to be passed by the user explicitly when output_shape changes based on computation then the same thing can be done by `inputs.set_shape([None,32,32,3]) ` that proposed in earlier workaround right.

I would like to understand how it is different from earlier workaround. For me the earlier work around using `set_shape` might work for all cases.Is there any exception for this ?



 This issue is stale because it has been open for 7 days with no activity. It will be closed if no further activity occurs. Thank you. This issue was closed because it has been inactive for 7 days since being marked as stale. Please reopen if you'd like to work on this further. Are you satisfied with the resolution of your issue?
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=Yes&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/59726">Yes</a>
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=No&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/59726">No</a>
 I feel like this should be documented officially. I just wasted a few hours pulling my hairs :(
Issue Type -> Bug
Have you reproduced the bug with TF nightly? -> No
Tensorflow Version -> v2.11.0-rc2-17-gd5b57ca93e5,  2.11.0
Custom Code -> Yes
OS Platform and Distribution -> Google colab 

#### Standalone code to reproduce the issue
```python3
import numpy as np,tensorflow as tf 
@tf.function
def sample(inputs:tf.Tensor):                                                   # some sample function 

  def numpy_method(inputs:np.ndarray):
    # some complex method 
    return np.sqrt(inputs)# example code 

  inputs=tf.numpy_function(numpy_method,[inputs],inputs.dtype,name='outputs')   # make apply numpy method 
  inputs=tf.keras.layers.AvgPool2D()(inputs)                                    # ERROR !!!!!

  return inputs

sample(tf.random.uniform((1,32,32,3)))
```

#### Relevant log output

```shell

.....

ValueError: in user code:

    File "<ipython-input-189-8001acaa2b43>", line 10, in sample  *
        inputs=tf.keras.layers.AvgPool2D()(inputs)                                    # ERROR !!!!!
    File "/usr/local/lib/python3.8/dist-packages/keras/utils/traceback_utils.py", line 70, in error_handler  **
        raise e.with_traceback(filtered_tb) from None

    ValueError: Exception encountered when calling layer 'average_pooling2d' (type AveragePooling2D).
    
    Cannot take the length of shape with unknown rank.
    
    Call arguments received by layer 'average_pooling2d' (type AveragePooling2D):
      • inputs=tf.Tensor(shape=<unknown>, dtype=float32)
```

#### But works fine if we reverse the order (i.e. tf.numpy_function cannot properly cast back outputs to tensors)
```python3
import numpy as np,tensorflow as tf 
@tf.function
def sample(inputs:tf.Tensor):                                                   # some sample function 

  def numpy_method(inputs:np.ndarray):
    # some complex method 
    return np.sqrt(inputs)# example code 

  inputs=tf.keras.layers.AvgPool2D()(inputs)                                    # works fine 
  inputs=tf.numpy_function(numpy_method,[inputs],inputs.dtype,name='outputs')   # make apply numpy method 

  return inputs

sample(tf.random.uniform((1,32,32,3)))
```
#### Output 
```shell
<tf.Tensor: shape=(1, 16, 16, 3), dtype=float32, numpy=
array([[[[0.7783269 , 0.6260437 , 0.81261575] ......
```
This issue was also mention here -> #37193 ["Cannot take the length of shape with unknown rank" error](https://github.com/tensorflow/tensorflow/issues/37193#top) in but not resolved. 