Inconsistent type handling between backends: ops.sqrt

[lbortolotti]

Sorry for the spam :-)

I'm migrating my codebase to exclusively use keras_core.ops methods, including all all "base" operations, based on feedback received in keras-team/keras-core#919 and others.

I'm still having some trouble, however.

The following code works correctly with the jax backend, but with tensorflow does the following:

import os

os.environ['KERAS_BACKEND'] = 'tensorflow'
from keras_core import ops

print(ops.sqrt(ops.arange(10)))

Throws:

  File "keras_core_experiments\repro_2.py", line 6, in <module>
    print(ops.sqrt(ops.arange(10)))
  File "keras_core_experiments\venv\lib\site-packages\keras_core\src\ops\numpy.py", line 5274, in sqrt
    return backend.numpy.sqrt(x)
  File "keras_core_experiments\venv\lib\site-packages\keras_core\src\backend\tensorflow\numpy.py", line 644, in sqrt
    return tfnp.sqrt(x)
  File "keras_core_experiments\venv\lib\site-packages\tensorflow\python\ops\numpy_ops\np_math_ops.py", line 615, in sqrt
    return _scalar(math_ops.sqrt, x, True)
  File "keras_core_experiments\venv\lib\site-packages\tensorflow\python\ops\numpy_ops\np_math_ops.py", line 599, in _scalar
    x = x.astype(np_dtypes.default_float_type())
  File "keras_core_experiments\venv\lib\site-packages\tensorflow\python\framework\ops.py", line 424, in __getattr__
    raise AttributeError(
AttributeError: EagerTensor object has no attribute 'astype'. 
        If you are looking for numpy-related methods, please run the following:
        from tensorflow.python.ops.numpy_ops import np_config
        np_config.enable_numpy_behavior()
      . Did you mean: 'dtype'?

[james77777778]

I believe this is the limitation of tensorflow, as it cannot directly perform sqrt on int32 tensor.

Workaround:

from keras_core import ops

print(ops.arange(10).dtype)
# <dtype: 'int32'>

# method 1: cast before sqrt
print(ops.sqrt(ops.cast(ops.arange(10), "float32")))
# [0.        1.        1.4142135 1.7320508 2.        2.2360678 2.4494896
#  2.6457512 2.828427  3.       ], shape=(10,), dtype=float32)

# method 2: create tensor with supported dtype
print(ops.sqrt(ops.arange(10, dtype="float32")))
# [0.        1.        1.4142135 1.7320508 2.        2.2360678 2.4494896
#  2.6457512 2.828427  3.       ], shape=(10,), dtype=float32)

print(tf.sqrt(tf.range(10)))
# InvalidArgumentError

[lbortolotti]

That is the underlying issue - however at this point I'm a little unclear as to what the design goal of the keras_core.ops API is.

If it's meant totally decouple the user from the underlying framework, this seems like a bug that needs addressing (potentially performing the cast you suggest within keras_core.ops).

If it keras_core.ops is in fact expected to behave differently across backends, then it loses some value.

In any case, in the interim, thanks for the workaround!

[james77777778]

Some additional information:

backend	dtype of arange(dtype=None)	can perform sqrt?	dtype after sqrt
numpy	int64	yes	float64
jax	int32	yes	float32
tensorflow	int32	no	X
torch	float32	yes	float32

I think we should standardize the behavior of these two ops?

[fchollet]

Thanks for the report. I have fixed some of these issues (the title issue in particular), but generally it's likely there are still various dtype inconsistencies lurking. We should do extensive testing of dtype handling in backend ops.

[jackd]

I noticed a lot of these while working on a PR to resolve #18415 . In particular, the ops will work when run symbolically - and often time return an incorrectly typed symbolic tensor. e.g. in the sqrt case I believe you'll get back a tensor with the same type, even if it's an integer. Same goes for most trig functions, and even a lot of logical operations that are intended to be used on bools.

[fchollet]

Good call, we also need to test consistency between symbolic outputs and real outputs.

[sachinprasadhs]

@lbortolotti , This is working fine in the Keras 3.0.2, please find the working Gist attached

[github-actions]

This issue is stale because it has been open for 14 days with no activity. It will be closed if no further activity occurs. Thank you.

[github-actions]

This issue was closed because it has been inactive for 28 days. Please reopen if you'd like to work on this further.

[google-ml-butler]

Are you satisfied with the resolution of your issue?
Yes
No