tensorflow: @tf.function triggers "AttributeError: 'Tensor' object has no attribute 'ndim'"

TensorFlow version: v2.4.0-49-g85c8b2a817f 2.4.1

Minimal Example

import tensorflow as tf

@tf.function
def f(a):
    return a.ndim

x = tf.zeros((3, 4))
f(x)

Describe the current behavior

AttributeError: 'Tensor' object has no attribute 'ndim'

Describe the expected behavior

<tf.Tensor: shape=(), dtype=int32, numpy=2>

This actually works

import tensorflow as tf

@tf.function
def f(a):
    return len(a.shape)

x = tf.zeros((3, 4))
f(x)

This problem occured in https://github.com/jonasrauber/eagerpy/issues/36

About this issue

Original URL
State: closed
Created 3 years ago
Comments: 19 (14 by maintainers)

Commits related to this issue

correct implementation of ndim in order to work in compile mode with tensorflow (see issue https://github.com/tensorflow/tensorflow/issues/48612 and https://github.com/jonasrauber/eagerpy/issues/36) — committed to eserie/eagerpy by eserie 3 years ago
correct implementation of ndim ... (#40) * correct implementation of ndim in order to work in compile mode with tensorflow (see issue https://github.com/tensorflow/tensorflow/issues/48612 and https:/... — committed to jonasrauber/eagerpy by eserie 3 years ago

Most upvoted comments

Yes. If you want, you could still raise an AttributeError, but only selectively if ndim is actually unknown.

It doesn’t make sense to delegate this to downstream code / libraries. Imho, they should have to care as little as possible about whether the code is eventually wrapped by @tf.function or not.

jonasrauber on Apr 20, 2021

In the end it boils down to this: the API should be consistent, so that the same code works with and without @tf.function.

jonasrauber on Apr 19, 2021

Yes, that’s what I think should be changed 😉 Or is there any reason this has to stay like it is?

jonasrauber on Apr 19, 2021

This is an inconsistency in the Tensor class. I see ndim was added to EagerTensor (forwarding to .shape.ndims), but not to the base Tensor class). That should be a straightforward fix.

The recommended method right now is .shape.rank and that should work consistently, but all these do need to be cleaned up.

mdanatg on May 19, 2021

You are in graph mode I think you need something like:

import tensorflow as tf

@tf.function
def f(a):
    tf.print(tf.rank(a))
    return tf.rank(a)

x = tf.zeros((3, 4))
f(x)

bhack on Apr 19, 2021