mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-19 05:16:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

Cast "axis" arg of tf.concat to tf.int32.

Browse Source 
PiperOrigin-RevId: 537478940
This commit is contained in:
Yu Emma Wang 2023-06-02 21:11:54 -07:00 committed by jax authors
parent 5639e194be
commit 983e1c0fd1
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
jax/experimental/jax2tf/jax2tf.py
View File

@ -1923,7 +1923,7 @@ tf_impl_with_avals[lax.clamp_p] = _clamp
def _concatenate(*operands, dimension):
return tf.concat(operands, axis=dimension)
return tf.concat(operands, axis=tf.cast(dimension, tf.int32))
tf_impl[lax.concatenate_p] = _concatenate