Bfloat16 on GPU for  sigmoid/swish activations, dropout and LSTM layers
@yufang67 As per the guide for [Mixed precision](https://www.tensorflow.org/guide/mixed_precision#setting_the_dtype_policy), the bfloat16 should be used to run on TPUs. Using `mixed_float16` is suggested when running on GPUs, and CPUs.

I did try with `mixed_float16` and the model works fine. Please find the gist [here](https://colab.research.google.com/gist/pjpratik/86f09732242b0595d6b3b505e230200e/59743.ipynb).

Can you please elaborate your feature request with use case. Thanks! Thanks @pjpratik . 
Yes, we are trying to train an ASR model on GPU with bfloat16. We would like to check the benefit of using bfloat16 like speedup and accuracy comparing to the mixed_float16 which we usually use. The activations and layers above are used in the model. when training on GPU, other ops and layers are on GPU, but they are on CPU. it has an huge impact on speed.  As a result, they are also expected to supported on GPU.

similar threads:
https://github.com/tensorflow/tensorflow/issues/59050#issuecomment-1428862660
https://github.com/tensorflow/tensorflow/issues/59728

Thanks Hi @yufang67 Thanks for the information.

@SuryanarayanaY Could you please look at this feature request? Thanks. Hi @yufang67 ,

`bfloat16` data type specifically designed for TPUs only and on GPUs this wont work and hence throwing user error as `ValueError:Value passed to parameter 'input' has DataType bfloat16 not in list of allowed values: float16, float32, float64` which is intended behaviour.

Please use `mixed_bfloat16` for TPUs environment only. For GPUs you should use `mixed_float16` .

Thanks!
 @SuryanarayanaY I dont think mixed_bfloat16 is for TPUs only (A100 support bfloat16). As more and more layers support bfloat16 on GPU, i would like to test the gain for the whole model. And the above activations and layers are blockers. thats why i open this feature request. 

if i understand you correctly, do you means that there is no plan to support bfloat16 for LSTM layer  ?

Thanks This issue has been automatically marked as stale because it has no recent activity. It will be closed if no further activity occurs. Thank you.
 @SuryanarayanaY Is that possible to handle sigmoid function as done in https://github.com/tensorflow/tensorflow/issues/59050#issuecomment-1428862660 for relu ? thanks > `bfloat16` data type specifically designed for TPUs only and on GPUs this wont work and hence throwing user error as `ValueError:Value passed to parameter 'input' has DataType bfloat16 not in list of allowed values: float16, float32, float64` which is intended behaviour.

I find this is surprising. PyTorch has enabled bfloat16 support on compatible GPU (e.g. A100) quite some time ago.  Does this mean TF prioritize the optimization implementation on TPUs over GPUs? Hi @tengwenxuan ,

Let me correct myself. On GPUs `bfloat16` may be workable if used with XLA i.e j`it_compile=True` either with `model.compile()` or with `tf.function()` for the Ops that supported on XLA. Support to add 'bfloat16' on GPUs also started recently and in current nightly version few ops might already support `bfloat16` with GPU. Right now I don't have any exclusive list of Ops for same.

@yufang67 ,

Could you please compile your model with XLA and use nightly version and let us know the output.

Thanks!



 @SuryanarayanaY Thanks for response. Yes, i have already tried XLA and it runs without error, but the training runs extremely slowly because of the variable length of data.  > Hi @tengwenxuan ,
> 
> Let me correct myself. On GPUs `bfloat16` may be workable if used with XLA i.e j`it_compile=True` either with `model.compile()` or with `tf.function()` for the Ops that supported on XLA. Support to add 'bfloat16' on GPUs also started recently and in current nightly version few ops might already support `bfloat16` with GPU. Right now I don't have any exclusive list of Ops for same.
> 
> @yufang67 ,
> 
> Could you please compile your model with XLA and use nightly version and let us know the output.
> 
> Thanks!

Thanks. however, XLA is not suitable here due to varying tensor sizes, I think we need native support of bfloat16 in order to have real performance gain.  @trevor-m do you know if it's feasible to add GPU bf16 support to such ops? @reedwm We still plan to support cwise ops like this through MLIR gen which is not yet ready. Until then we can easily add bf16 registrations. I originally thought there was way too many, but I checked and it doesn't look that bad if we were to add gpu registrations for the cwise ops which have bfloat16 cpu implementations already. I can take a look at that soon.


 Thanks! Adding gpu registrations to cwise ops that have bf16 cpu implementations SGTM. Thank you for the follow-up. re: "support this thru MLIR gen", I'm curious, is this what was explained in [this blog post](https://medium.com/tensorflow/mlir-a-new-intermediate-representation-and-compiler-framework-beba999ed18d)? If so, will it still suffer from the performance issue of changing tensor size like in current XLA/JIT implementation? Do you expect the support thru MLIR will provide better performance similar as native kernel support?
 I opened a PR for the cwise op support.

Looks like we will also need bfloat16 support for CudnnRNN as well for this example since it is used by keras - not sure how to tell keras not to use this. This might require casting to float if bf16 is not supported by cudnnRnn.

> Thank you for the follow-up. re: "support this thru MLIR gen", I'm curious, is this what was explained in [this blog post](https://medium.com/tensorflow/mlir-a-new-intermediate-representation-and-compiler-framework-beba999ed18d)? If so, will it still suffer from the performance issue of changing tensor size like in current XLA/JIT implementation? Do you expect the support thru MLIR will provide better performance similar as native kernel support?

I think this isn't the same as what's described by that blog post. A large number of TF kernel today use MLIR generation by default. It can be either JIT or AOT (ahead of time). I think the perf of mlir generated kernels is expected to be better, but not 100% sure. @kushanam might be able to provide more details. @reedwm Looks like CudnnRNN doesn't support BF16 yet. Should I add a bf16 registration for the op and just cast to float always? I guess the alternative would be to exclude CudnnRNN from the auto mixed precision allowlist for bf16 only - I don't see any place where there is currently a way to have a different list for bf16 vs fp16. @yufang67,
The related PR(https://github.com/tensorflow/tensorflow/pull/59972) which was raised has been merged and alsoI tried to execute the mentioned code on tensorflow v2.15 (GPU) by importing the keras from the keras_core and the code was executed without any issue/error. Kindly find the gist of it [here](https://colab.research.google.com/gist/tilakrayal/fdc27fa1701f846a63b859b9dab4132c/untitled1575.ipynb). Thank you!

 Thanks for the update. This issue can be closed. @yufang67,
Glad the issue is resolved for you, Could you please feel free to move this to closed status. Thank you!
<details><summary>Click to expand!</summary> 
 
 ### Issue Type

Feature Request

### Have you reproduced the bug with TF nightly?

Yes

### Source

source

### Tensorflow Version

TensorFlow version 2.13.0-dev20230215

### Custom Code

Yes

### OS Platform and Distribution

ubuntu 20.04

### Mobile device

_No response_

### Python version

3.8

### Bazel version

_No response_

### GCC/Compiler version

_No response_

### CUDA/cuDNN version

11.8/8.6

### GPU model and memory

single A100 80G

### Current Behaviour?

```shell
current behaviour:
when using mixed_bfloat16 policy, sigmoid type activation and dropout layer run on CPU, and LSTM layer does not support bfloat16 input.

expected behavior:
sigmoid/swish activation, dropout and LSTM layers run on GPU with mixed_bfloat16 policy
```


### Standalone code to reproduce the issue

```shell
import tensorflow as tf
from tensorflow.keras.models import Sequential
from tensorflow.keras import mixed_precision
from tensorflow.keras.layers import LSTM, Embedding

tf.debugging.set_log_device_placement(True)
#from tensorflow.python.framework.ops import disable_eager_execution
#disable_eager_execution()


policy = mixed_precision.Policy('mixed_bfloat16')
# float32
print(policy.name, policy.variable_dtype, policy.compute_dtype)

mixed_precision.set_global_policy(policy)

input_shape = (4, 28, 28, 3)
x = tf.random.normal((4, 28, 28, 3))
#layer = tf.keras.layers.Conv2D(2, 3, activation='relu', input_shape=(28, 28, 3))
layer1 = tf.keras.layers.Activation(activation=tf.keras.activations.swish)
layer2 = tf.keras.layers.Activation(activation=tf.keras.activations.sigmoid)
#layer3 = tf.keras.layers.Activation(activation=tf.keras.activations.relu)
layer4 = tf.keras.layers.Dropout(0.1, name="dropout")

x1 = layer1(x)
x2 = layer2(x1)
y = layer4(x2, training=True)
print(' ====== output ======== ', y.dtype, layer1.dtype, layer2.dtype, layer4.dtype)

vocab_size = 50
input_shape = (4, vocab_size, 1)
x = tf.random.normal(input_shape)

def customer_model():
    model = Sequential()
    model.add(LSTM(128))
    return model

model = customer_model()
model.build(input_shape=input_shape)
model.summary()
y = model(x)
```


### Relevant log output

```shell
2023-02-20 15:15:13.745147: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
To enable the following instructions: AVX2 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
2023-02-20 15:15:14.363645: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT
mixed_bfloat16 float32 bfloat16
2023-02-20 15:15:15.799744: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1635] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 78915 MB memory:  -> device: 0, name: NVIDIA A100 80GB PCIe, pci bus id: 0001:00:00.0, compute capability: 8.0
input: (_Arg): /job:localhost/replica:0/task:0/device:CPU:0
2023-02-20 15:15:15.809731: I tensorflow/core/common_runtime/placer.cc:114] input: (_Arg): /job:localhost/replica:0/task:0/device:CPU:0
_EagerConst: (_EagerConst): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:15.809763: I tensorflow/core/common_runtime/placer.cc:114] _EagerConst: (_EagerConst): /job:localhost/replica:0/task:0/device:GPU:0
output_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:15.809775: I tensorflow/core/common_runtime/placer.cc:114] output_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:15.813051: I tensorflow/core/common_runtime/eager/execute.cc:1514] Executing op _EagerConst in device /job:localhost/replica:0/task:0/device:GPU:0
input: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.219908: I tensorflow/core/common_runtime/placer.cc:114] input: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
_EagerConst: (_EagerConst): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.219944: I tensorflow/core/common_runtime/placer.cc:114] _EagerConst: (_EagerConst): /job:localhost/replica:0/task:0/device:GPU:0
output_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.219952: I tensorflow/core/common_runtime/placer.cc:114] output_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.221005: I tensorflow/core/common_runtime/eager/execute.cc:1514] Executing op _EagerConst in device /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.221972: I tensorflow/core/common_runtime/eager/execute.cc:1514] Executing op _EagerConst in device /job:localhost/replica:0/task:0/device:GPU:0
shape: (_DeviceArg): /job:localhost/replica:0/task:0/device:CPU:0
2023-02-20 15:15:16.222322: I tensorflow/core/common_runtime/placer.cc:114] shape: (_DeviceArg): /job:localhost/replica:0/task:0/device:CPU:0
RandomStandardNormal: (RandomStandardNormal): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.222337: I tensorflow/core/common_runtime/placer.cc:114] RandomStandardNormal: (RandomStandardNormal): /job:localhost/replica:0/task:0/device:GPU:0
output_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.222345: I tensorflow/core/common_runtime/placer.cc:114] output_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.222865: I tensorflow/core/common_runtime/eager/execute.cc:1514] Executing op RandomStandardNormal in device /job:localhost/replica:0/task:0/device:GPU:0
x: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.223723: I tensorflow/core/common_runtime/placer.cc:114] x: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
y: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.223736: I tensorflow/core/common_runtime/placer.cc:114] y: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
Mul: (Mul): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.223750: I tensorflow/core/common_runtime/placer.cc:114] Mul: (Mul): /job:localhost/replica:0/task:0/device:GPU:0
z_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.223761: I tensorflow/core/common_runtime/placer.cc:114] z_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.224202: I tensorflow/core/common_runtime/eager/execute.cc:1514] Executing op Mul in device /job:localhost/replica:0/task:0/device:GPU:0
x: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.224907: I tensorflow/core/common_runtime/placer.cc:114] x: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
y: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.224922: I tensorflow/core/common_runtime/placer.cc:114] y: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
AddV2: (AddV2): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.224936: I tensorflow/core/common_runtime/placer.cc:114] AddV2: (AddV2): /job:localhost/replica:0/task:0/device:GPU:0
z_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.224944: I tensorflow/core/common_runtime/placer.cc:114] z_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.225386: I tensorflow/core/common_runtime/eager/execute.cc:1514] Executing op AddV2 in device /job:localhost/replica:0/task:0/device:GPU:0
x: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.229609: I tensorflow/core/common_runtime/placer.cc:114] x: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
Cast: (Cast): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.229623: I tensorflow/core/common_runtime/placer.cc:114] Cast: (Cast): /job:localhost/replica:0/task:0/device:GPU:0
y_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.229630: I tensorflow/core/common_runtime/placer.cc:114] y_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.230004: I tensorflow/core/common_runtime/eager/execute.cc:1514] Executing op Cast in device /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.230386: I tensorflow/core/common_runtime/eager/execute.cc:1514] Executing op _EagerConst in device /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.230496: I tensorflow/core/common_runtime/eager/execute.cc:1514] Executing op Cast in device /job:localhost/replica:0/task:0/device:GPU:0
x: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.230820: I tensorflow/core/common_runtime/placer.cc:114] x: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
y: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.230831: I tensorflow/core/common_runtime/placer.cc:114] y: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
Mul: (Mul): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.230843: I tensorflow/core/common_runtime/placer.cc:114] Mul: (Mul): /job:localhost/replica:0/task:0/device:GPU:0
z_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.230850: I tensorflow/core/common_runtime/placer.cc:114] z_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.231215: I tensorflow/core/common_runtime/eager/execute.cc:1514] Executing op Mul in device /job:localhost/replica:0/task:0/device:GPU:0
x: (_Arg): /job:localhost/replica:0/task:0/device:CPU:0
2023-02-20 15:15:16.231575: I tensorflow/core/common_runtime/placer.cc:114] x: (_Arg): /job:localhost/replica:0/task:0/device:CPU:0
Sigmoid: (Sigmoid): /job:localhost/replica:0/task:0/device:CPU:0
2023-02-20 15:15:16.231596: I tensorflow/core/common_runtime/placer.cc:114] Sigmoid: (Sigmoid): /job:localhost/replica:0/task:0/device:CPU:0
y_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:CPU:0
2023-02-20 15:15:16.231608: I tensorflow/core/common_runtime/placer.cc:114] y_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:CPU:0
2023-02-20 15:15:16.232198: I tensorflow/core/common_runtime/eager/execute.cc:1514] Executing op Sigmoid in device /job:localhost/replica:0/task:0/device:CPU:0
2023-02-20 15:15:16.233223: I tensorflow/core/common_runtime/eager/execute.cc:1514] Executing op Mul in device /job:localhost/replica:0/task:0/device:GPU:0
input: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.233727: I tensorflow/core/common_runtime/placer.cc:114] input: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
Identity: (Identity): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.233747: I tensorflow/core/common_runtime/placer.cc:114] Identity: (Identity): /job:localhost/replica:0/task:0/device:GPU:0
output_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.233759: I tensorflow/core/common_runtime/placer.cc:114] output_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.234163: I tensorflow/core/common_runtime/eager/execute.cc:1514] Executing op Identity in device /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.234923: I tensorflow/core/common_runtime/eager/execute.cc:1514] Executing op Sigmoid in device /job:localhost/replica:0/task:0/device:CPU:0
x: (_Arg): /job:localhost/replica:0/task:0/device:CPU:0
2023-02-20 15:15:16.235519: I tensorflow/core/common_runtime/placer.cc:114] x: (_Arg): /job:localhost/replica:0/task:0/device:CPU:0
Cast: (Cast): /job:localhost/replica:0/task:0/device:CPU:0
2023-02-20 15:15:16.235536: I tensorflow/core/common_runtime/placer.cc:114] Cast: (Cast): /job:localhost/replica:0/task:0/device:CPU:0
y_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:CPU:0
2023-02-20 15:15:16.235546: I tensorflow/core/common_runtime/placer.cc:114] y_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:CPU:0
2023-02-20 15:15:16.235889: I tensorflow/core/common_runtime/eager/execute.cc:1514] Executing op Cast in device /job:localhost/replica:0/task:0/device:CPU:0
input: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.236126: I tensorflow/core/common_runtime/placer.cc:114] input: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
_EagerConst: (_EagerConst): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.236140: I tensorflow/core/common_runtime/placer.cc:114] _EagerConst: (_EagerConst): /job:localhost/replica:0/task:0/device:GPU:0
output_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.236147: I tensorflow/core/common_runtime/placer.cc:114] output_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.236694: I tensorflow/core/common_runtime/eager/execute.cc:1514] Executing op _EagerConst in device /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.237002: I tensorflow/core/common_runtime/eager/execute.cc:1514] Executing op Mul in device /job:localhost/replica:0/task:0/device:GPU:0
input: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.237338: I tensorflow/core/common_runtime/placer.cc:114] input: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
Shape: (Shape): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.237358: I tensorflow/core/common_runtime/placer.cc:114] Shape: (Shape): /job:localhost/replica:0/task:0/device:GPU:0
output_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.237369: I tensorflow/core/common_runtime/placer.cc:114] output_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.237779: I tensorflow/core/common_runtime/eager/execute.cc:1514] Executing op Shape in device /job:localhost/replica:0/task:0/device:GPU:0
shape: (_DeviceArg): /job:localhost/replica:0/task:0/device:CPU:0
2023-02-20 15:15:16.238193: I tensorflow/core/common_runtime/placer.cc:114] shape: (_DeviceArg): /job:localhost/replica:0/task:0/device:CPU:0
RandomUniform: (RandomUniform): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.238212: I tensorflow/core/common_runtime/placer.cc:114] RandomUniform: (RandomUniform): /job:localhost/replica:0/task:0/device:GPU:0
output_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.238223: I tensorflow/core/common_runtime/placer.cc:114] output_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.238690: I tensorflow/core/common_runtime/eager/execute.cc:1514] Executing op RandomUniform in device /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.238925: I tensorflow/core/common_runtime/eager/execute.cc:1514] Executing op Cast in device /job:localhost/replica:0/task:0/device:CPU:0
2023-02-20 15:15:16.239130: I tensorflow/core/common_runtime/eager/execute.cc:1514] Executing op _EagerConst in device /job:localhost/replica:0/task:0/device:GPU:0
x: (_Arg): /job:localhost/replica:0/task:0/device:CPU:0
2023-02-20 15:15:16.239387: I tensorflow/core/common_runtime/placer.cc:114] x: (_Arg): /job:localhost/replica:0/task:0/device:CPU:0
y: (_Arg): /job:localhost/replica:0/task:0/device:CPU:0
2023-02-20 15:15:16.239404: I tensorflow/core/common_runtime/placer.cc:114] y: (_Arg): /job:localhost/replica:0/task:0/device:CPU:0
GreaterEqual: (GreaterEqual): /job:localhost/replica:0/task:0/device:CPU:0
2023-02-20 15:15:16.239434: I tensorflow/core/common_runtime/placer.cc:114] GreaterEqual: (GreaterEqual): /job:localhost/replica:0/task:0/device:CPU:0
z_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:CPU:0
2023-02-20 15:15:16.239442: I tensorflow/core/common_runtime/placer.cc:114] z_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:CPU:0
2023-02-20 15:15:16.240101: I tensorflow/core/common_runtime/eager/execute.cc:1514] Executing op GreaterEqual in device /job:localhost/replica:0/task:0/device:CPU:0
x: (_Arg): /job:localhost/replica:0/task:0/device:CPU:0
2023-02-20 15:15:16.240458: I tensorflow/core/common_runtime/placer.cc:114] x: (_Arg): /job:localhost/replica:0/task:0/device:CPU:0
Cast: (Cast): /job:localhost/replica:0/task:0/device:CPU:0
2023-02-20 15:15:16.240477: I tensorflow/core/common_runtime/placer.cc:114] Cast: (Cast): /job:localhost/replica:0/task:0/device:CPU:0
y_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:CPU:0
2023-02-20 15:15:16.240489: I tensorflow/core/common_runtime/placer.cc:114] y_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:CPU:0
2023-02-20 15:15:16.240879: I tensorflow/core/common_runtime/eager/execute.cc:1514] Executing op Cast in device /job:localhost/replica:0/task:0/device:CPU:0
2023-02-20 15:15:16.241141: I tensorflow/core/common_runtime/eager/execute.cc:1514] Executing op _EagerConst in device /job:localhost/replica:0/task:0/device:GPU:0
condition: (_Arg): /job:localhost/replica:0/task:0/device:CPU:0
2023-02-20 15:15:16.241482: I tensorflow/core/common_runtime/placer.cc:114] condition: (_Arg): /job:localhost/replica:0/task:0/device:CPU:0
t: (_Arg): /job:localhost/replica:0/task:0/device:CPU:0
2023-02-20 15:15:16.241500: I tensorflow/core/common_runtime/placer.cc:114] t: (_Arg): /job:localhost/replica:0/task:0/device:CPU:0
e: (_Arg): /job:localhost/replica:0/task:0/device:CPU:0
2023-02-20 15:15:16.241508: I tensorflow/core/common_runtime/placer.cc:114] e: (_Arg): /job:localhost/replica:0/task:0/device:CPU:0
SelectV2: (SelectV2): /job:localhost/replica:0/task:0/device:CPU:0
2023-02-20 15:15:16.241516: I tensorflow/core/common_runtime/placer.cc:114] SelectV2: (SelectV2): /job:localhost/replica:0/task:0/device:CPU:0
output_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:CPU:0
2023-02-20 15:15:16.241527: I tensorflow/core/common_runtime/placer.cc:114] output_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:CPU:0
2023-02-20 15:15:16.242187: I tensorflow/core/common_runtime/eager/execute.cc:1514] Executing op SelectV2 in device /job:localhost/replica:0/task:0/device:CPU:0
 ====== output ========  <dtype: 'bfloat16'> float32 float32 float32
2023-02-20 15:15:16.242688: I tensorflow/core/common_runtime/eager/execute.cc:1514] Executing op _EagerConst in device /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.242816: I tensorflow/core/common_runtime/eager/execute.cc:1514] Executing op RandomStandardNormal in device /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.242920: I tensorflow/core/common_runtime/eager/execute.cc:1514] Executing op Mul in device /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.243029: I tensorflow/core/common_runtime/eager/execute.cc:1514] Executing op AddV2 in device /job:localhost/replica:0/task:0/device:GPU:0
input: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.244373: I tensorflow/core/common_runtime/placer.cc:114] input: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
_EagerConst: (_EagerConst): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.244402: I tensorflow/core/common_runtime/placer.cc:114] _EagerConst: (_EagerConst): /job:localhost/replica:0/task:0/device:GPU:0
output_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.244422: I tensorflow/core/common_runtime/placer.cc:114] output_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.245327: I tensorflow/core/common_runtime/eager/execute.cc:1514] Executing op _EagerConst in device /job:localhost/replica:0/task:0/device:GPU:0
resource_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.245736: I tensorflow/core/common_runtime/placer.cc:114] resource_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:GPU:0
VarHandleOp: (VarHandleOp): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.245756: I tensorflow/core/common_runtime/placer.cc:114] VarHandleOp: (VarHandleOp): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.246125: I tensorflow/core/common_runtime/eager/execute.cc:1514] Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:0
resource: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.253589: I tensorflow/core/common_runtime/placer.cc:114] resource: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
value: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.253619: I tensorflow/core/common_runtime/placer.cc:114] value: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
AssignVariableOp: (AssignVariableOp): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.253635: I tensorflow/core/common_runtime/placer.cc:114] AssignVariableOp: (AssignVariableOp): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.254140: I tensorflow/core/common_runtime/eager/execute.cc:1514] Executing op AssignVariableOp in device /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.259081: I tensorflow/core/common_runtime/eager/execute.cc:1514] Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.259243: I tensorflow/core/common_runtime/eager/execute.cc:1514] Executing op AssignVariableOp in device /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.259523: I tensorflow/core/common_runtime/eager/execute.cc:1514] Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.259624: I tensorflow/core/common_runtime/eager/execute.cc:1514] Executing op AssignVariableOp in device /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.266083: I tensorflow/core/common_runtime/eager/execute.cc:1514] Executing op _EagerConst in device /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.266318: I tensorflow/core/common_runtime/eager/execute.cc:1514] Executing op _EagerConst in device /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.266583: I tensorflow/core/common_runtime/eager/execute.cc:1514] Executing op _EagerConst in device /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.266830: I tensorflow/core/common_runtime/eager/execute.cc:1514] Executing op _EagerConst in device /job:localhost/replica:0/task:0/device:GPU:0
seed: (_Arg): /job:localhost/replica:0/task:0/device:CPU:0
2023-02-20 15:15:16.267103: I tensorflow/core/common_runtime/placer.cc:114] seed: (_Arg): /job:localhost/replica:0/task:0/device:CPU:0
StatelessRandomGetKeyCounter: (StatelessRandomGetKeyCounter): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.267120: I tensorflow/core/common_runtime/placer.cc:114] StatelessRandomGetKeyCounter: (StatelessRandomGetKeyCounter): /job:localhost/replica:0/task:0/device:GPU:0
key_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.267134: I tensorflow/core/common_runtime/placer.cc:114] key_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:GPU:0
counter_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.267143: I tensorflow/core/common_runtime/placer.cc:114] counter_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.267699: I tensorflow/core/common_runtime/eager/execute.cc:1514] Executing op StatelessRandomGetKeyCounter in device /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.286569: I tensorflow/core/common_runtime/eager/execute.cc:1514] Executing op _EagerConst in device /job:localhost/replica:0/task:0/device:GPU:0
shape: (_DeviceArg): /job:localhost/replica:0/task:0/device:CPU:0
2023-02-20 15:15:16.286927: I tensorflow/core/common_runtime/placer.cc:114] shape: (_DeviceArg): /job:localhost/replica:0/task:0/device:CPU:0
key: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.286943: I tensorflow/core/common_runtime/placer.cc:114] key: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
counter: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.286957: I tensorflow/core/common_runtime/placer.cc:114] counter: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
alg: (_DeviceArg): /job:localhost/replica:0/task:0/device:CPU:0
2023-02-20 15:15:16.286966: I tensorflow/core/common_runtime/placer.cc:114] alg: (_DeviceArg): /job:localhost/replica:0/task:0/device:CPU:0
StatelessRandomUniformV2: (StatelessRandomUniformV2): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.286977: I tensorflow/core/common_runtime/placer.cc:114] StatelessRandomUniformV2: (StatelessRandomUniformV2): /job:localhost/replica:0/task:0/device:GPU:0
output_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.286994: I tensorflow/core/common_runtime/placer.cc:114] output_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.287706: I tensorflow/core/common_runtime/eager/execute.cc:1514] Executing op StatelessRandomUniformV2 in device /job:localhost/replica:0/task:0/device:GPU:0
x: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.291713: I tensorflow/core/common_runtime/placer.cc:114] x: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
y: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.291736: I tensorflow/core/common_runtime/placer.cc:114] y: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
Sub: (Sub): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.291752: I tensorflow/core/common_runtime/placer.cc:114] Sub: (Sub): /job:localhost/replica:0/task:0/device:GPU:0
z_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.291766: I tensorflow/core/common_runtime/placer.cc:114] z_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.292313: I tensorflow/core/common_runtime/eager/execute.cc:1514] Executing op Sub in device /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.300187: I tensorflow/core/common_runtime/eager/execute.cc:1514] Executing op Mul in device /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.300343: I tensorflow/core/common_runtime/eager/execute.cc:1514] Executing op AddV2 in device /job:localhost/replica:0/task:0/device:GPU:0
resource_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.300704: I tensorflow/core/common_runtime/placer.cc:114] resource_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:GPU:0
VarHandleOp: (VarHandleOp): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.300720: I tensorflow/core/common_runtime/placer.cc:114] VarHandleOp: (VarHandleOp): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.301056: I tensorflow/core/common_runtime/eager/execute.cc:1514] Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:0
resource: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.301395: I tensorflow/core/common_runtime/placer.cc:114] resource: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
value: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.301410: I tensorflow/core/common_runtime/placer.cc:114] value: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
AssignVariableOp: (AssignVariableOp): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.301424: I tensorflow/core/common_runtime/placer.cc:114] AssignVariableOp: (AssignVariableOp): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.301772: I tensorflow/core/common_runtime/eager/execute.cc:1514] Executing op AssignVariableOp in device /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.302407: I tensorflow/core/common_runtime/eager/execute.cc:1514] Executing op _EagerConst in device /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.302584: I tensorflow/core/common_runtime/eager/execute.cc:1514] Executing op _EagerConst in device /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.302674: I tensorflow/core/common_runtime/eager/execute.cc:1514] Executing op StatelessRandomGetKeyCounter in device /job:localhost/replica:0/task:0/device:GPU:0
shape: (_DeviceArg): /job:localhost/replica:0/task:0/device:CPU:0
2023-02-20 15:15:16.303396: I tensorflow/core/common_runtime/placer.cc:114] shape: (_DeviceArg): /job:localhost/replica:0/task:0/device:CPU:0
key: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.303409: I tensorflow/core/common_runtime/placer.cc:114] key: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
counter: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.303444: I tensorflow/core/common_runtime/placer.cc:114] counter: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
alg: (_DeviceArg): /job:localhost/replica:0/task:0/device:CPU:0
2023-02-20 15:15:16.303454: I tensorflow/core/common_runtime/placer.cc:114] alg: (_DeviceArg): /job:localhost/replica:0/task:0/device:CPU:0
StatelessRandomNormalV2: (StatelessRandomNormalV2): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.303464: I tensorflow/core/common_runtime/placer.cc:114] StatelessRandomNormalV2: (StatelessRandomNormalV2): /job:localhost/replica:0/task:0/device:GPU:0
output_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.303474: I tensorflow/core/common_runtime/placer.cc:114] output_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.304084: I tensorflow/core/common_runtime/eager/execute.cc:1514] Executing op StatelessRandomNormalV2 in device /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.304383: I tensorflow/core/common_runtime/eager/execute.cc:1514] Executing op Mul in device /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.304498: I tensorflow/core/common_runtime/eager/execute.cc:1514] Executing op AddV2 in device /job:localhost/replica:0/task:0/device:GPU:0
input: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.304807: I tensorflow/core/common_runtime/placer.cc:114] input: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
Qr: (Qr): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.304826: I tensorflow/core/common_runtime/placer.cc:114] Qr: (Qr): /job:localhost/replica:0/task:0/device:GPU:0
q_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.304838: I tensorflow/core/common_runtime/placer.cc:114] q_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:GPU:0
r_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.304845: I tensorflow/core/common_runtime/placer.cc:114] r_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.305406: I tensorflow/core/common_runtime/eager/execute.cc:1514] Executing op Qr in device /job:localhost/replica:0/task:0/device:GPU:0
input: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.321530: I tensorflow/core/common_runtime/placer.cc:114] input: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
DiagPart: (DiagPart): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.321554: I tensorflow/core/common_runtime/placer.cc:114] DiagPart: (DiagPart): /job:localhost/replica:0/task:0/device:GPU:0
diagonal_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.321568: I tensorflow/core/common_runtime/placer.cc:114] diagonal_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.322028: I tensorflow/core/common_runtime/eager/execute.cc:1514] Executing op DiagPart in device /job:localhost/replica:0/task:0/device:GPU:0
x: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.322425: I tensorflow/core/common_runtime/placer.cc:114] x: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
Sign: (Sign): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.322448: I tensorflow/core/common_runtime/placer.cc:114] Sign: (Sign): /job:localhost/replica:0/task:0/device:GPU:0
y_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.322462: I tensorflow/core/common_runtime/placer.cc:114] y_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.322849: I tensorflow/core/common_runtime/eager/execute.cc:1514] Executing op Sign in device /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.333823: I tensorflow/core/common_runtime/eager/execute.cc:1514] Executing op Mul in device /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.334332: I tensorflow/core/common_runtime/eager/execute.cc:1514] Executing op _EagerConst in device /job:localhost/replica:0/task:0/device:GPU:0
x: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.334700: I tensorflow/core/common_runtime/placer.cc:114] x: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
perm: (_DeviceArg): /job:localhost/replica:0/task:0/device:CPU:0
2023-02-20 15:15:16.334718: I tensorflow/core/common_runtime/placer.cc:114] perm: (_DeviceArg): /job:localhost/replica:0/task:0/device:CPU:0
Transpose: (Transpose): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.334731: I tensorflow/core/common_runtime/placer.cc:114] Transpose: (Transpose): /job:localhost/replica:0/task:0/device:GPU:0
y_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.334742: I tensorflow/core/common_runtime/placer.cc:114] y_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.335267: I tensorflow/core/common_runtime/eager/execute.cc:1514] Executing op Transpose in device /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.342633: I tensorflow/core/common_runtime/eager/execute.cc:1514] Executing op _EagerConst in device /job:localhost/replica:0/task:0/device:GPU:0
tensor: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.343070: I tensorflow/core/common_runtime/placer.cc:114] tensor: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
shape: (_DeviceArg): /job:localhost/replica:0/task:0/device:CPU:0
2023-02-20 15:15:16.343097: I tensorflow/core/common_runtime/placer.cc:114] shape: (_DeviceArg): /job:localhost/replica:0/task:0/device:CPU:0
Reshape: (Reshape): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.343117: I tensorflow/core/common_runtime/placer.cc:114] Reshape: (Reshape): /job:localhost/replica:0/task:0/device:GPU:0
output_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.343132: I tensorflow/core/common_runtime/placer.cc:114] output_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.343815: I tensorflow/core/common_runtime/eager/execute.cc:1514] Executing op Reshape in device /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.347316: I tensorflow/core/common_runtime/eager/execute.cc:1514] Executing op Mul in device /job:localhost/replica:0/task:0/device:GPU:0
resource_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.348034: I tensorflow/core/common_runtime/placer.cc:114] resource_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:GPU:0
VarHandleOp: (VarHandleOp): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.348058: I tensorflow/core/common_runtime/placer.cc:114] VarHandleOp: (VarHandleOp): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.348550: I tensorflow/core/common_runtime/eager/execute.cc:1514] Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:0
resource: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.348896: I tensorflow/core/common_runtime/placer.cc:114] resource: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
value: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.348912: I tensorflow/core/common_runtime/placer.cc:114] value: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
AssignVariableOp: (AssignVariableOp): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.348922: I tensorflow/core/common_runtime/placer.cc:114] AssignVariableOp: (AssignVariableOp): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.349273: I tensorflow/core/common_runtime/eager/execute.cc:1514] Executing op AssignVariableOp in device /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.349657: I tensorflow/core/common_runtime/eager/execute.cc:1514] Executing op _EagerConst in device /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.349863: I tensorflow/core/common_runtime/eager/execute.cc:1514] Executing op _EagerConst in device /job:localhost/replica:0/task:0/device:GPU:0
dims: (_DeviceArg): /job:localhost/replica:0/task:0/device:CPU:0
2023-02-20 15:15:16.350162: I tensorflow/core/common_runtime/placer.cc:114] dims: (_DeviceArg): /job:localhost/replica:0/task:0/device:CPU:0
value: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.350183: I tensorflow/core/common_runtime/placer.cc:114] value: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
Fill: (Fill): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.350203: I tensorflow/core/common_runtime/placer.cc:114] Fill: (Fill): /job:localhost/replica:0/task:0/device:GPU:0
output_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.350213: I tensorflow/core/common_runtime/placer.cc:114] output_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.350732: I tensorflow/core/common_runtime/eager/execute.cc:1514] Executing op Fill in device /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.354235: I tensorflow/core/common_runtime/eager/execute.cc:1514] Executing op _EagerConst in device /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.354513: I tensorflow/core/common_runtime/eager/execute.cc:1514] Executing op _EagerConst in device /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.354638: I tensorflow/core/common_runtime/eager/execute.cc:1514] Executing op Fill in device /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.354780: I tensorflow/core/common_runtime/eager/execute.cc:1514] Executing op _EagerConst in device /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.354859: I tensorflow/core/common_runtime/eager/execute.cc:1514] Executing op Fill in device /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.354995: I tensorflow/core/common_runtime/eager/execute.cc:1514] Executing op _EagerConst in device /job:localhost/replica:0/task:0/device:GPU:0
values_0: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.355260: I tensorflow/core/common_runtime/placer.cc:114] values_0: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
values_1: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.355270: I tensorflow/core/common_runtime/placer.cc:114] values_1: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
values_2: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.355276: I tensorflow/core/common_runtime/placer.cc:114] values_2: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
axis: (_DeviceArg): /job:localhost/replica:0/task:0/device:CPU:0
2023-02-20 15:15:16.355285: I tensorflow/core/common_runtime/placer.cc:114] axis: (_DeviceArg): /job:localhost/replica:0/task:0/device:CPU:0
ConcatV2: (ConcatV2): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.355299: I tensorflow/core/common_runtime/placer.cc:114] ConcatV2: (ConcatV2): /job:localhost/replica:0/task:0/device:GPU:0
output_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.355305: I tensorflow/core/common_runtime/placer.cc:114] output_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.355806: I tensorflow/core/common_runtime/eager/execute.cc:1514] Executing op ConcatV2 in device /job:localhost/replica:0/task:0/device:GPU:0
resource_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.359677: I tensorflow/core/common_runtime/placer.cc:114] resource_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:GPU:0
VarHandleOp: (VarHandleOp): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.359707: I tensorflow/core/common_runtime/placer.cc:114] VarHandleOp: (VarHandleOp): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.360132: I tensorflow/core/common_runtime/eager/execute.cc:1514] Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:0
resource: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.360527: I tensorflow/core/common_runtime/placer.cc:114] resource: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
value: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.360544: I tensorflow/core/common_runtime/placer.cc:114] value: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
AssignVariableOp: (AssignVariableOp): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.360554: I tensorflow/core/common_runtime/placer.cc:114] AssignVariableOp: (AssignVariableOp): /job:localhost/replica:0/task:0/device:GPU:0
2023-02-20 15:15:16.360967: I tensorflow/core/common_runtime/eager/execute.cc:1514] Executing op AssignVariableOp in device /job:localhost/replica:0/task:0/device:GPU:0
Traceback (most recent call last):
  File "/home/azureuser/env_tf_nightly2/lib/python3.8/site-packages/keras/engine/training.py", line 510, in build
    self.call(x, **kwargs)
  File "/home/azureuser/env_tf_nightly2/lib/python3.8/site-packages/keras/engine/sequential.py", line 427, in call
    outputs = layer(inputs, **kwargs)
  File "/home/azureuser/env_tf_nightly2/lib/python3.8/site-packages/keras/layers/rnn/base_rnn.py", line 556, in __call__
    return super().__call__(inputs, **kwargs)
  File "/home/azureuser/env_tf_nightly2/lib/python3.8/site-packages/keras/utils/traceback_utils.py", line 70, in error_handler
    raise e.with_traceback(filtered_tb) from None
  File "/home/azureuser/env_tf_nightly2/lib/python3.8/site-packages/tensorflow/python/framework/op_def_library.py", line 56, in _SatisfiesTypeConstraint
    raise TypeError(
TypeError: Exception encountered when calling layer 'lstm' (type LSTM).

Value passed to parameter 'input' has DataType bfloat16 not in list of allowed values: float16, float32, float64

Call arguments received by layer 'lstm' (type LSTM):
  • inputs=tf.Tensor(shape=(4, 50, 1), dtype=bfloat16)
  • mask=None
  • training=None
  • initial_state=None

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "test_bfloat16.py", line 72, in <module>
    model.build(input_shape=input_shape)
  File "/home/azureuser/env_tf_nightly2/lib/python3.8/site-packages/keras/engine/sequential.py", line 383, in build
    super().build(input_shape)
  File "/home/azureuser/env_tf_nightly2/lib/python3.8/site-packages/keras/engine/training.py", line 512, in build
    raise ValueError(
ValueError: You cannot build your model by calling `build` if your layers do not support float type inputs. Instead, in order to instantiate and build your model, call your model on real tensor data (of the correct dtype).

The actual error from `call` is: Exception encountered when calling layer 'lstm' (type LSTM).

Value passed to parameter 'input' has DataType bfloat16 not in list of allowed values: float16, float32, float64

Call arguments received by layer 'lstm' (type LSTM):
  • inputs=tf.Tensor(shape=(4, 50, 1), dtype=bfloat16)
  • mask=None
  • training=None
  • initial_state=None.
```
</details>