hand
This commit is contained in:
@@ -0,0 +1,278 @@
|
||||
# Copyright 2022 The JAX Authors.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# https://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
# Note: import <name> as <name> is required for names to be exported.
|
||||
# See PEP 484 & https://github.com/jax-ml/jax/issues/7570
|
||||
|
||||
import jax._src.core as _src_core
|
||||
|
||||
from jax._src.core import (
|
||||
AbstractValue as AbstractValue,
|
||||
Atom as Atom,
|
||||
ParamDict as ParamDict,
|
||||
ShapedArray as ShapedArray,
|
||||
Trace as Trace,
|
||||
Tracer as Tracer,
|
||||
Value as Value,
|
||||
ensure_compile_time_eval as ensure_compile_time_eval,
|
||||
eval_context as eval_context,
|
||||
eval_jaxpr as eval_jaxpr,
|
||||
max_dim as max_dim,
|
||||
min_dim as min_dim,
|
||||
)
|
||||
|
||||
_deprecations = {
|
||||
# Deprecated in v0.8.2; finalized in v0.10.0.
|
||||
# TODO(jakevdp) remove entries in v0.11.0.
|
||||
"get_aval": (
|
||||
"jax.core.get_aval was deprecated in JAX v0.8.2 and removed in JAX v0.10.0;"
|
||||
" use jax.typeof instead.",
|
||||
None,
|
||||
),
|
||||
"mapped_aval": (
|
||||
"jax.core.mapped_aval was deprecated in JAX v0.8.2 and removed in JAX"
|
||||
" v0.10.0. Use jax.extend.core.mapped_aval.",
|
||||
None,
|
||||
),
|
||||
"unmapped_aval": (
|
||||
"jax.core.unmapped_aval was deprecated in JAX v0.8.2 and removed in JAX"
|
||||
" v0.10.0. Use jax.extend.core.unmapped_aval.",
|
||||
None,
|
||||
),
|
||||
"set_current_trace": (
|
||||
"jax.core.set_current_trace was deprecated in JAX v0.8.2 and removed in"
|
||||
" JAX v0.10.0. Use jax.extend.core.set_current_trace.",
|
||||
None,
|
||||
),
|
||||
"take_current_trace": (
|
||||
"jax.core.take_current_trace was deprecated in JAX v0.8.2 and removed in"
|
||||
" JAX v0.10.0. Use jax.extend.core.take_current_trace.",
|
||||
None,
|
||||
),
|
||||
"traverse_jaxpr_params": (
|
||||
"jax.core.traverse_jaxpr_params was deprecated in JAX v0.8.2 and removed in"
|
||||
" JAX v0.10.0.",
|
||||
None,
|
||||
),
|
||||
"TraceTag": (
|
||||
"jax.core.TraceTag was deprecated in JAX v0.8.2 and removed in JAX v0.10.0."
|
||||
" Use jax.extend.core.TraceTag.",
|
||||
None,
|
||||
),
|
||||
"call_impl": (
|
||||
"jax.core.call_impl was deprecated in JAX v0.8.2 and removed in JAX"
|
||||
" v0.10.0. Use jax.extend.core.call_impl.",
|
||||
None,
|
||||
),
|
||||
"subjaxprs": (
|
||||
"jax.core.subjaxprs was deprecated in JAX v0.8.2 and removed in JAX"
|
||||
" v0.10.0. Use jax.extend.core.subjaxprs.",
|
||||
None,
|
||||
),
|
||||
"AbstractToken": (
|
||||
"jax.core.AbstractToken was deprecated in JAX v0.8.2 and removed in JAX"
|
||||
" v0.10.0. Use jax.extend.core.AbstractToken.",
|
||||
None,
|
||||
),
|
||||
# Deprecated in JAX v0.10.0; TODO(jakevdp) finalize in v0.11.0
|
||||
"CallPrimitive": (
|
||||
"jax.core.CallPrimitive is deprecated. Use jax.extend.core.CallPrimitive.",
|
||||
_src_core.CallPrimitive,
|
||||
),
|
||||
"DebugInfo": (
|
||||
"jax.core.DebugInfo is deprecated. Use jax.extend.core.DebugInfo.",
|
||||
_src_core.DebugInfo,
|
||||
),
|
||||
"DropVar": (
|
||||
"jax.core.DropVar is deprecated. Use jax.extend.core.DropVar.",
|
||||
_src_core.DropVar,
|
||||
),
|
||||
"Effect": (
|
||||
"jax.core.Effect is deprecated. Use jax.extend.core.Effect.",
|
||||
_src_core.Effect,
|
||||
),
|
||||
"Effects": (
|
||||
"jax.core.Effects is deprecated. Use jax.extend.core.Effects.",
|
||||
_src_core.Effects,
|
||||
),
|
||||
"InconclusiveDimensionOperation": (
|
||||
"jax.core.InconclusiveDimensionOperation is deprecated. Use jax.extend.core.InconclusiveDimensionOperation.",
|
||||
_src_core.InconclusiveDimensionOperation,
|
||||
),
|
||||
"JaxprTypeError": (
|
||||
"jax.core.JaxprTypeError is deprecated. Use jax.extend.core.JaxprTypeError.",
|
||||
_src_core.JaxprTypeError,
|
||||
),
|
||||
"check_jaxpr": (
|
||||
"jax.core.check_jaxpr is deprecated. Use jax.extend.core.check_jaxpr.",
|
||||
_src_core.check_jaxpr,
|
||||
),
|
||||
"concrete_or_error": (
|
||||
"jax.core.concrete_or_error is deprecated. Use jax.extend.core.concrete_or_error.",
|
||||
_src_core.concrete_or_error,
|
||||
),
|
||||
"find_top_trace": (
|
||||
"jax.core.find_top_trace is deprecated. Use jax.extend.core.find_top_trace.",
|
||||
_src_core.find_top_trace,
|
||||
),
|
||||
"gensym": (
|
||||
"jax.core.gensym is deprecated. Use jax.extend.core.gensym.",
|
||||
_src_core.gensym,
|
||||
),
|
||||
"get_opaque_trace_state": (
|
||||
"jax.core.get_opaque_trace_state is deprecated. Use jax.extend.core.get_opaque_trace_state.",
|
||||
_src_core.get_opaque_trace_state,
|
||||
),
|
||||
"jaxprs_in_params": (
|
||||
"jax.core.jaxprs_in_params is deprecated. Use jax.extend.core.jaxprs_in_params.",
|
||||
_src_core.jaxprs_in_params,
|
||||
),
|
||||
"new_jaxpr_eqn": (
|
||||
"jax.core.new_jaxpr_eqn is deprecated. Use jax.extend.core.new_jaxpr_eqn.",
|
||||
_src_core.new_jaxpr_eqn,
|
||||
),
|
||||
"no_effects": (
|
||||
"jax.core.no_effects is deprecated. Use jax.extend.core.no_effects.",
|
||||
_src_core.no_effects,
|
||||
),
|
||||
"nonempty_axis_env_DO_NOT_USE": (
|
||||
"jax.core.nonempty_axis_env_DO_NOT_USE is deprecated.",
|
||||
_src_core.nonempty_axis_env,
|
||||
),
|
||||
"primal_dtype_to_tangent_dtype": (
|
||||
"jax.core.primal_dtype_to_tangent_dtype is deprecated. Use jax.extend.core.primal_dtype_to_tangent_dtype.",
|
||||
_src_core.primal_dtype_to_tangent_dtype,
|
||||
),
|
||||
"unsafe_am_i_under_a_jit_DO_NOT_USE": (
|
||||
"jax.core.unsafe_am_i_under_a_jit_DO_NOT_USE is deprecated.",
|
||||
_src_core.unsafe_am_i_under_a_jit,
|
||||
),
|
||||
"unsafe_am_i_under_a_vmap_DO_NOT_USE": (
|
||||
"jax.core.unsafe_am_i_under_a_vmap_DO_NOT_USE is deprecated.",
|
||||
_src_core.unsafe_am_i_under_a_vmap,
|
||||
),
|
||||
"unsafe_get_axis_names_DO_NOT_USE": (
|
||||
"jax.core.unsafe_get_axis_names_DO_NOT_USE is deprecated.",
|
||||
_src_core.unsafe_get_axis_names,
|
||||
),
|
||||
"valid_jaxtype": (
|
||||
"jax.core.valid_jaxtype is deprecated. Use jax.extend.core.valid_jaxtype.",
|
||||
_src_core.valid_jaxtype,
|
||||
),
|
||||
"JaxprPpContext": (
|
||||
"jax.core.JaxprPpContext is deprecated.",
|
||||
_src_core.JaxprPpContext,
|
||||
),
|
||||
"JaxprPpSettings": (
|
||||
"jax.core.JaxprPpSettings is deprecated.",
|
||||
_src_core.JaxprPpSettings,
|
||||
),
|
||||
"OutputType": (
|
||||
"jax.core.OutputType is deprecated.",
|
||||
_src_core.OutputType,
|
||||
),
|
||||
"abstract_token": (
|
||||
"jax.core.abstract_token is deprecated.",
|
||||
_src_core.abstract_token,
|
||||
),
|
||||
"aval_mapping_handlers": (
|
||||
"jax.core.aval_mapping_handlers is deprecated.",
|
||||
_src_core.aval_mapping_handlers,
|
||||
),
|
||||
"call": (
|
||||
"jax.core.call is deprecated.",
|
||||
_src_core.call,
|
||||
),
|
||||
"concretization_function_error": (
|
||||
"jax.core.concretization_function_error is deprecated.",
|
||||
_src_core.concretization_function_error,
|
||||
),
|
||||
"custom_typechecks": (
|
||||
"jax.core.custom_typechecks is deprecated.",
|
||||
_src_core.custom_typechecks,
|
||||
),
|
||||
"is_concrete": (
|
||||
"jax.core.is_concrete is deprecated.",
|
||||
_src_core.is_concrete,
|
||||
),
|
||||
"is_constant_dim": (
|
||||
"jax.core.is_constant_dim is deprecated.",
|
||||
_src_core.is_constant_dim,
|
||||
),
|
||||
"is_constant_shape": (
|
||||
"jax.core.is_constant_shape is deprecated.",
|
||||
_src_core.is_constant_shape,
|
||||
),
|
||||
"literalable_types": (
|
||||
"jax.core.literalable_types is deprecated.",
|
||||
_src_core.literalable_types,
|
||||
),
|
||||
"no_axis_name": (
|
||||
"jax.core.no_axis_name is deprecated.",
|
||||
_src_core.no_axis_name,
|
||||
),
|
||||
"pytype_aval_mappings": (
|
||||
"jax.core.pytype_aval_mappings is deprecated.",
|
||||
_src_core.pytype_aval_mappings,
|
||||
),
|
||||
"trace_ctx": (
|
||||
"jax.core.trace_ctx is deprecated.",
|
||||
_src_core.trace_ctx,
|
||||
),
|
||||
}
|
||||
|
||||
import typing as _typing
|
||||
if _typing.TYPE_CHECKING:
|
||||
CallPrimitive = _src_core.CallPrimitive
|
||||
DebugInfo = _src_core.DebugInfo
|
||||
DropVar = _src_core.DropVar
|
||||
Effect = _src_core.Effect
|
||||
Effects = _src_core.Effects
|
||||
InconclusiveDimensionOperation = _src_core.InconclusiveDimensionOperation
|
||||
JaxprPpContext = _src_core.JaxprPpContext
|
||||
JaxprPpSettings = _src_core.JaxprPpSettings
|
||||
JaxprTypeError = _src_core.JaxprTypeError
|
||||
OutputType = _src_core.OutputType
|
||||
abstract_token = _src_core.abstract_token
|
||||
aval_mapping_handlers = _src_core.aval_mapping_handlers
|
||||
call = _src_core.call
|
||||
check_jaxpr = _src_core.check_jaxpr
|
||||
concrete_or_error = _src_core.concrete_or_error
|
||||
concretization_function_error = _src_core.concretization_function_error
|
||||
custom_typechecks = _src_core.custom_typechecks
|
||||
find_top_trace = _src_core.find_top_trace
|
||||
gensym = _src_core.gensym
|
||||
get_opaque_trace_state = _src_core.get_opaque_trace_state
|
||||
is_concrete = _src_core.is_concrete
|
||||
is_constant_dim = _src_core.is_constant_dim
|
||||
is_constant_shape = _src_core.is_constant_shape
|
||||
jaxprs_in_params = _src_core.jaxprs_in_params
|
||||
literalable_types = _src_core.literalable_types
|
||||
new_jaxpr_eqn = _src_core.new_jaxpr_eqn
|
||||
no_axis_name = _src_core.no_axis_name
|
||||
no_effects = _src_core.no_effects
|
||||
nonempty_axis_env_DO_NOT_USE = _src_core.nonempty_axis_env
|
||||
primal_dtype_to_tangent_dtype = _src_core.primal_dtype_to_tangent_dtype
|
||||
pytype_aval_mappings = _src_core.pytype_aval_mappings
|
||||
trace_ctx = _src_core.trace_ctx
|
||||
unsafe_am_i_under_a_jit_DO_NOT_USE = _src_core.unsafe_am_i_under_a_jit
|
||||
unsafe_am_i_under_a_vmap_DO_NOT_USE = _src_core.unsafe_am_i_under_a_vmap
|
||||
unsafe_get_axis_names_DO_NOT_USE = _src_core.unsafe_get_axis_names
|
||||
valid_jaxtype = _src_core.valid_jaxtype
|
||||
else:
|
||||
from jax._src.deprecations import deprecation_getattr as _deprecation_getattr
|
||||
__getattr__ = _deprecation_getattr(__name__, _deprecations)
|
||||
del _deprecation_getattr
|
||||
del _typing
|
||||
del _src_core
|
||||
Reference in New Issue
Block a user