mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-18 12:56:07 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #14790 from patrick-kidger:patch-3

Browse Source 
PiperOrigin-RevId: 514437379
This commit is contained in:
jax authors 2023-03-06 09:46:59 -08:00
commit 6aed604789
1 changed files with 7 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
jax/experimental/jax2tf/__init__.py
View File

@ -12,6 +12,10 @@
# See the License for the specific language governing permissions and
# limitations under the License.
from jax.experimental.jax2tf.jax2tf import (convert, dtype_of_val,
split_to_logical_devices, PolyShape)
from jax.experimental.jax2tf.call_tf import call_tf
from jax.experimental.jax2tf.jax2tf import (
convert as convert,
dtype_of_val as dtype_of_val,
split_to_logical_devices as split_to_logical_devices,
PolyShape as PolyShape
)
from jax.experimental.jax2tf.call_tf import call_tf as call_tf