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

remove core.{aval_method,aval_property}

Browse Source 
PiperOrigin-RevId: 510043837
This commit is contained in:
Roy Frostig 2023-02-15 22:21:27 -08:00 committed by jax authors
parent d8514d0ec6
commit 26045c49e7
2 changed files with 8 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
jax/core.py
View File

@ -114,8 +114,6 @@ from jax._src.core import (
as_named_shape as as_named_shape,
attrgetter as attrgetter,
aval_mapping_handlers as aval_mapping_handlers,
aval_method as aval_method,
aval_property as aval_property,
axis_frame as axis_frame,
call as call,
call_bind_with_continuation as call_bind_with_continuation,

13
tests/custom_object_test.py
View File

@ -18,20 +18,23 @@ import unittest
import numpy as np
import jax
from jax._src import test_util as jtu
import jax.numpy as jnp
from jax import core, jit, lax, make_jaxpr
from jax import jit, lax, make_jaxpr
from jax.config import config
from jax.interpreters import mlir
from jax.interpreters import xla
from jax._src import core
from jax._src import device_array
from jax._src import dispatch
from jax._src import dtypes
from jax.interpreters import mlir
from jax.interpreters import xla
from jax._src import test_util as jtu
from jax._src.lib.mlir import ir
from jax._src.lib import xla_bridge, xla_client
xc = xla_client
xb = xla_bridge
from jax.config import config
config.parse_flags_with_absl()
# TODO(jakevdp): use a setup/teardown method to populate and unpopulate all the