jax の API リファレンス

jax (jax-ml/jax) の公開 API 400 件 —— クラス 60、関数 303、メソッド 37。実際のソースを静的解析して抽出した正確なシグネチャを掲載しています。

リポジトリ: jax-ml/jax

種別	件数
クラス	60
関数	303
メソッド	37

API 一覧

funcbuild_wheel.copy_file(src_file:str, dst_dir:str) -> None

Copy a file to the destination directory.

funcbuild_wheel.prepare_srcs(deps:list[str], srcs_dir:str) -> None

Filter the sources and copy them to the destination directory.

func

jax._src.ad_checkpoint.remat(fun:Callable, *prevent_cse:bool=True, *policy:Callable[..., bool] | None=None, *static_argnums:int | tuple[int, ...]=()) -> Callable

Alias of :func:`jax.checkpoint`.

funcjax._src.ad_checkpoint.save_any_names_but_these(*names_not_to_save)

Save only named values, i.e.

funcjax._src.ad_checkpoint.save_from_both_policies(policy_1, policy_2)

Logical OR of the given policies.

classjax._src.api.NotSpecified

Sentinel for use in jax.jit

funcjax._src.api.clear_caches()

Clear all compilation and staging caches.

funcjax._src.api.device_get(x:Any)

Transfer ``x`` to host.

func

jax._src.api.device_put(x, device:None | xc.Device | Sharding | P | Format | Any=None, *src:None | xc.Device | Sharding | P | Format | Any=None, *donate:bool | Any=False, *may_alias:bool | None | Any=None)

Transfers ``x`` to ``device``.

funcjax._src.api.eval_shape(fun:Callable, *args, **kwargs)

Compute the shape/dtype of ``fun`` without any FLOPs.

func

jax._src.api.hessian(fun:Callable, argnums:int | Sequence[int]=0, has_aux:bool=False, holomorphic:bool=False) -> Callable

Hessian of ``fun`` as a dense array.

func

jax._src.api.jacobian(fun:Callable, argnums:int | Sequence[int]=0, has_aux:bool=False, holomorphic:bool=False, allow_int:bool=False) -> Callable

Alias of :func:`jax.jacrev`.

funcjax._src.api.jvp(fun:Callable, primals, tangents, has_aux:bool=False) -> tuple[Any, ...]

Computes a (forward-mode) Jacobian-vector product of ``fun``.

funcjax._src.api.linear_transpose(fun:Callable, *reduce_axes=(), *primals) -> Callable

Transpose a function that is promised to be linear.

funcjax._src.api.tuptree_flags(prefix, treedef, name:str, full_name:str) -> list[bool]

Expand a flags prefix into per-leaf flags for `treedef`.

func

jax._src.api.vmap(fun:F, in_axes:int | None | Sequence[Any]=0, out_axes:Any=0, axis_name:AxisName | None=None, axis_size:int | None=None, spmd_axis_name:AxisName | tuple[AxisName, ...] | None=None, sum_match:bool=False) -> F

Vectorizing map.

funcjax._src.api_util.rebase_donate_argnums(donate_argnums, static_argnums) -> tuple[int, ...]

Shifts donate to account for static.

classjax._src.array.Shard

A single data shard of an Array.

funcjax._src.array.as_slice_indices(arr:Any, idx:Index) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]

Returns start_indices, limit_indices, removed_dims

methodjax._src.basearray.Array.addressable_data(index:int) -> Array

Return an array of the addressable data at a particular index.

methodjax._src.basearray.Array.addressable_shards() -> Sequence[Shard]

List of addressable shards.

methodjax._src.basearray.Array.committed() -> bool

Whether the array is committed or not.

methodjax._src.basearray.Array.device() -> Device | Sharding

Array API-compatible device attribute.

methodjax._src.basearray.Array.dtype() -> np.dtype

The data type (:class:`numpy.dtype`) of the array.

methodjax._src.basearray.Array.global_shards() -> Sequence[Shard]

List of global shards.

methodjax._src.basearray.Array.is_fully_addressable() -> bool

Is this Array fully addressable?

methodjax._src.basearray.Array.is_fully_replicated() -> bool

Is this Array fully replicated?

methodjax._src.basearray.Array.ndim() -> int

The number of dimensions in the array.

methodjax._src.basearray.Array.shape() -> tuple[int, ...]

The shape of the array.

methodjax._src.basearray.Array.sharding() -> Sharding

The sharding for the array.

methodjax._src.basearray.Array.size() -> int

The total number of elements in the array.

func

jax._src.blocked_sampler.sample_block(sampler_fn:SampleFn, keys:NdKeyList, block_size:Shape, tile_size:Shape, *args, **kwargs) -> Array

Draws random samples for a single block.

funcjax._src.cache_key.add_flag_prefixes(flag_prefixes:list[str]) -> None

Add flag prefixes to include in the cache key.

funcjax._src.cache_key.clear_flag_prefixes() -> None

Clear flag prefixes added by add_flag_prefixes().

funcjax._src.cache_key.custom_hook() -> str

Custom hook for any addition to the cache key.

func

jax._src.cache_key.get(module:ir.Module, devices:np.ndarray, compile_options:xla_client.CompileOptions, backend:xla_client.Client, compression_algorithm:str='zstandard', ignore_custom_partitioning:bool=False) -> str

Creates a hashed string to use as a key to the compilation cache.

funcjax._src.cache_key.get_flag_prefixes() -> list[str]

Return flag prefixes added by add_flag_prefixes().

funcjax._src.checkify.check(pred:Bool, msg:str, *debug:bool=False, *fmt_args, **fmt_kwargs) -> None

Check a predicate, add an error with msg if predicate is False.

funcjax._src.checkify.check_error(error:Error) -> None

Raise an Exception if ``error`` represents a failure.

funcjax._src.checkify.debug_check(pred:Bool, msg:str, *fmt_args, **fmt_kwargs) -> None

Check a predicate when running under checkify, otherwise is a no-op.

funcjax._src.checkify.div_error_check(error, enabled_errors, x, y)

Checks for division by zero and NaN.

funcjax._src.cloud_tpu_init.cloud_tpu_init() -> None

Automatically sets Cloud TPU topology and other env vars.

funcjax._src.cloud_tpu_init.is_libtpu_at_least(version_str:str) -> bool

Returns True if not running on Cloud TPU.

classjax._src.clusters.cluster.ClusterEnv

Interface for defining a cluster environment.

method

jax._src.clusters.cluster.ClusterEnv.get_coordinator_address(timeout_secs:int | None, override_coordinator_port:str | None) -> str

Returns address and port used by JAX to bootstrap.

methodjax._src.clusters.cluster.ClusterEnv.get_local_process_id() -> int | None

Get index of current process inside a host.

classjax._src.compilation_cache.VerificationCache

A cache that wraps another cache and verifies its contents.

funcjax._src.compilation_cache.get_file_cache(path:str) -> tuple[CacheInterface, str] | None

Returns the file cache and the path to the cache.

funcjax._src.compilation_cache.initialize_cache(path) -> None

This API is deprecated; use set_cache_dir instead.

funcjax._src.compilation_cache.is_executable_in_cache(backend, cache_key:str) -> bool

Checks if the executable is in the cache.

funcjax._src.compilation_cache.is_initialized() -> bool

Deprecated.

funcjax._src.compilation_cache.reset_cache() -> None

Get back to pristine, uninitialized state.

funcjax._src.compilation_cache.set_cache_dir(path) -> None

Sets the persistent compilation cache directory.

classjax._src.compiler.CompilerEffortLevel

Effort level enumeration for XLA.

funcjax._src.compiler.use_detailed_logging(module:ir.Module) -> bool

Returns 'true' if detailed logging should be enabled for 'module'.

classjax._src.config.ValueHolder

A holder for a configuration value.

funcjax._src.config.bool_env(varname:str, default:bool) -> bool

Read an environment variable and interpret it as a boolean.

funcjax._src.config.explicit_device_put_scope() -> Generator[None]

Indicates that the current context is an explicit device_put*() call.

funcjax._src.config.int_env(varname:str, default:int) -> int

Read an environment variable and interpret it as an integer.

classjax._src.core.Ref

Mutable array reference.

methodjax._src.core.Trace.stage_value(val)

Lifts a value into a trace.

funcjax._src.core.canonicalize_shape(shape:Shape, context:str='') -> tuple[Any, ...]

Canonicalizes and checks for errors in a user-provided shape value.

funcjax._src.core.canonicalize_slice(s:slice, axis_size:DimSize) -> tuple[DimSize, DimSize, DimSize]

Computes the start index, step, and size of the slice `x[s]`.

funcjax._src.core.check_jaxpr(jaxpr:Jaxpr)

Checks well-formedness of a jaxpr.

funcjax._src.core.concrete_or_error(force:Any, val:Any, context='')

Like force(val), but gives the context in the error message.

funcjax._src.core.definitely_equal_shape(s1:Shape, s2:Shape) -> bool

Check that two shapes are guaranteed to be element-wise equal.

funcjax._src.core.dilate_dim(d:DimSize, dilation:DimSize) -> DimSize

max(0, 1 + dilation * (d - 1)).

funcjax._src.core.dim_value_dtype()

The dtype to be used for dimension values.

funcjax._src.core.dimension_as_value(d:DimSize)

Turns a dimension size into a JAX array.

funcjax._src.core.divide_shape_sizes(s1:Shape, s2:Shape) -> DimSize

Returns an integer "i" s.t., i * size(s2) == size(s1).

funcjax._src.core.evaluate_shape(shape:Shape, dim_vars:Sequence[str], *dim_values:Array) -> Sequence[Array]

Evaluates a shape possibly containing non-constants.

funcjax._src.core.free_ref(ref:Ref)

Invalidate a given reference.

funcjax._src.core.freeze(ref:Ref) -> Array

Invalidate a given reference and return its final value.

funcjax._src.core.get_sharding(sharding, shape)

Modifies and checks the sharding.

funcjax._src.core.is_symbolic_dim(v:Any) -> bool

Checks if a value is a symbolic dimension used for shape polymorphism.

funcjax._src.core.max_dim(d1:DimSize, d2:DimSize) -> DimSize

Like max(d1, d2) but for both constant and symbolic dimensions.

funcjax._src.core.maybe_find_leaked_tracers(trace:Trace) -> list[Tracer]

Find the leaked tracers holding a reference to the Trace

funcjax._src.core.mem_space_to_kind(mem_space:Any) -> str

Converts a memory space to its corresponding XLA memory kind string.

funcjax._src.core.min_dim(d1:DimSize, d2:DimSize) -> DimSize

Like min(d1, d2) but for both constant and symbolic dimensions.

funcjax._src.core.subjaxprs(jaxpr:Jaxpr) -> Iterator[Jaxpr]

Generator for all subjaxprs found in the params of jaxpr.eqns.

funcjax._src.core.typecompat(aval_ref:AbstractValue, aval:AbstractValue) -> bool

Determine whether `aval` conforms to `aval_ref`.

funcjax._src.core.typematch(t1:AbstractValue, t2:AbstractValue, no_dtype_check:bool=False) -> bool

Determine whether `t1` and `t2` are equivalent.

funcjax._src.core.typeof(x:Any) -> Any

Return the JAX type (i.e.

funcjax._src.cudnn.fusion.cudnn_fusion(f)

Makes a function become a cuDNN kernel.

classjax._src.custom_batching.custom_vmap

Customize the vmap behavior of a JAX-transformable function.

method

jax._src.custom_batching.custom_vmap.def_vmap(vmap_rule:Callable[..., tuple[Any, Any]]) -> Callable[..., tuple[Any, Any]]

Define the vmap rule for this custom_vmap function.

classjax._src.custom_partitioning_sharding_rule.ArrayMapping

Describes the factors for an operand or result.

classjax._src.custom_partitioning_sharding_rule.CompoundFactor

Describes the factors for a compound factor.

classjax._src.custom_partitioning_sharding_rule.SdyShardingRule

Represents a Shardy sharding rule.

classjax._src.debugger.cli_debugger.CliDebugger

A text-based debugger.

methodjax._src.debugger.cli_debugger.CliDebugger.default(line)

Evaluates an expression.

methodjax._src.debugger.cli_debugger.CliDebugger.do_down(arg)

d(own) Move down a stack frame.

methodjax._src.debugger.cli_debugger.CliDebugger.do_quit(_)

q(uit) (exit) Quit the debugger.

methodjax._src.debugger.cli_debugger.CliDebugger.do_up(arg)

u(p) Move up a stack frame.

classjax._src.debugger.colab_debugger.CodeViewer

A mutable DOM element that displays code as HTML.

methodjax._src.debugger.colab_debugger.CodeViewer.update_code(code_, highlights, *linenostart:int=1)

Updates the code viewer to use new code.

classjax._src.debugger.colab_debugger.ColabDebugger

A JAX debugger for a Colab environment.

classjax._src.debugger.colab_debugger.DebuggerView

Main view for the Colab debugger.

classjax._src.debugger.colab_debugger.FramePreview

Displays information about a stack frame.

classjax._src.debugger.colab_lib.DynamicDOMElement

A DOM element that can be mutated.

classjax._src.debugger.colab_lib.DynamicDiv

A `div` that can be edited.

classjax._src.debugger.colab_lib.StaticDOMElement

An immutable DOM element.

classjax._src.debugger.core.DebuggerFrame

Encapsulates Python frame information.

classjax._src.debugger.web_debugger.WebDebugger

A web-based debugger.

funcjax._src.debugging.visualize_array_sharding(arr, **kwargs)

Visualizes an array's sharding.

funcjax._src.deprecations.accelerate_getattr_deprecation(module:ModuleType, *names:str) -> None

Accelerate the deprecation of a module-level attribute.

funcjax._src.deprecations.is_accelerated_attribute(module:ModuleType, name:str) -> bool

Returns true if given name is accelerated.

func

jax._src.deprecations.warn(deprecation_id:str, message:str, stacklevel:int, *error_class:type[Exception]=ValueError) -> None

Warns about a deprecation, or errors if the deprecation is accelerated.

funcjax._src.dispatch.jaxpr_has_primitive(jaxpr:core.Jaxpr, prim_name:str) -> bool

Whether there is a primitive given by user anywhere inside a Jaxpr.

funcjax._src.distributed.is_initialized() -> bool

Check if the JAX distributed system is initialized.

funcjax._src.distributed.shutdown()

Shuts down the distributed system.

funcjax._src.dlpack.is_supported_dtype(dtype:DTypeLike) -> bool

Check if dtype is supported by jax.dlpack.

classjax._src.dtypes.ExtendedDType

Abstract Base Class for extended dtypes

classjax._src.dtypes.TypePromotionError

Raised when JAX type promotion fails.

funcjax._src.dtypes.check_and_canonicalize_user_dtype(dtype, fun_name=None, *allow_non_jax_dtypes:bool=False) -> DType

Checks validity of a user-provided dtype, and returns its canonical form.

funcjax._src.dtypes.coerce_to_array(x:Any, dtype:DTypeLike | None=None) -> np.ndarray

Coerces a scalar or NumPy array to an np.array.

funcjax._src.dtypes.dtype(x:Any) -> DType

Return the dtype object for a value or type.

classjax._src.dtypes.extended

Scalar class for extended dtypes.

funcjax._src.dtypes.issubdtype(a:DTypeLike | ExtendedDType | None, b:DTypeLike | ExtendedDType | None) -> bool

Returns True if first argument is a typecode lower/equal in type hierarchy.

funcjax._src.dtypes.itemsize_bits(dtype:DTypeLike) -> int

Number of bits per element for the dtype.

funcjax._src.dtypes.jax_dtype(obj:DTypeLike | None, *align:bool=False, *copy:bool=False) -> DType

Cast an object to a dtype, respecting JAX dtype defaults.

classjax._src.dtypes.prng_key

Scalar class for PRNG Key dtypes.

funcjax._src.dtypes.promote_types(a:DTypeLike, b:DTypeLike) -> DType

Returns the type to which a binary operation should cast its arguments.

funcjax._src.dtypes.register_weak_scalar_type(typ:type)

funcjax._src.dtypes.scalar_type_of(x:Any) -> type

Return the scalar type associated with a JAX value.

funcjax._src.dtypes.scalar_type_to_dtype(typ:type, value:Any=None) -> DType

Return the numpy dtype for the given scalar type.

funcjax._src.dtypes.supports_inf(dtype:DTypeLike) -> bool

Return true if the dtype supports infinity, else return False.

funcjax._src.dtypes.to_floating_dtype(dtype:DTypeLike) -> DType

Promotes a dtype to a non-complex floating dtype.

funcjax._src.dtypes.to_inexact_dtype(dtype:DTypeLike) -> DType

Promotes a dtype into an inexact dtype, if it is not already one.

funcjax._src.dtypes.to_numeric_dtype(dtype:DTypeLike) -> DType

Promotes a dtype into an numeric dtype, if it is not already one.

classjax._src.effects.Effect

A generic side-effect.

funcjax._src.error_check.raise_if_error() -> None

Raise an exception if the internal error state is set.

funcjax._src.error_check.set_error_if(pred:Array, msg:str) -> None

Set the internal error state if any element of `pred` is `True`.

classjax._src.errors.JAXIndexError

JAX-specific :class:`IndexError`

classjax._src.errors.JAXTypeError

JAX-specific :class:`TypeError`

classjax._src.errors.KeyReuseError

This error occurs when a PRNG key is reused in an unsafe manner.

classjax._src.export._export.DisabledSafetyCheck

A safety check that should be skipped on (de)serialization.

classjax._src.export._export.Exported

A JAX function lowered to StableHLO.

methodjax._src.export._export.Exported.has_vjp() -> bool

Returns if this Exported supports VJP.

methodjax._src.export._export.Exported.mlir_module(serialized:bool=True) -> Any

A string or Module representation of the ``mlir_module_serialized``.

methodjax._src.export._export.Exported.serialize(vjp_order:int=0) -> bytearray

Serializes an Exported.

methodjax._src.export._export.Exported.vjp() -> Exported

Gets the exported VJP.

funcjax._src.export._export.default_export_platform() -> str

Retrieves the default export platform.

funcjax._src.export._export.deserialize(blob:bytearray) -> Exported

Deserializes an Exported.

funcjax._src.export.serialization.deserialize(ser:bytearray) -> _export.Exported

Deserializes an Exported.

funcjax._src.export.serialization.serialize(exp:_export.Exported, vjp_order:int=0) -> bytearray

Serializes an Exported.

methodjax._src.export.serialization_generated.Effect.GetRootAsEffect(buf, offset=0)

This method is deprecated.

methodjax._src.export.serialization_generated.Exported.GetRootAsExported(buf, offset=0)

This method is deprecated.

methodjax._src.export.serialization_generated.Sharding.GetRootAsSharding(buf, offset=0)

This method is deprecated.

classjax._src.export.shape_poly.BoundsPrecision

Specifies desired precision for the bounds calculation.

classjax._src.export.shape_poly.PolyShape

Tuple of polymorphic dimension specifications.

classjax._src.export.shape_poly.SymbolicScope

Identifies a scope for symbolic expressions.

funcjax._src.export.shape_poly.cmp_sequence(s1, s2, elem_cmp) -> int

Compares two sequences using `elem_cmp`.

funcjax._src.export.shape_poly.is_symbolic_dim(p:DimSize) -> TypeGuard[_DimExpr]

Checks if a dimension is symbolic.

func

jax._src.export.shape_poly.shape_assertion(assert_what:typing.Array, *error_message:str, *error_message_inputs:typing.Array) -> None

Adds a shape assertion in the code.

classjax._src.ffi.TypeRegistration

A dictionary type for registering FFI types.

funcjax._src.ffi.register_ffi_target(name:str, fn:Any, platform:str='cpu', api_version:int=1, **kwargs:Any) -> None

Registers a foreign function target.

funcjax._src.ffi.register_ffi_target_as_batch_partitionable(name:str) -> None

Registers an FFI target as batch partitionable.

funcjax._src.ffi.register_ffi_type(name:str, type_registration:TypeRegistration, platform:str='cpu') -> None

Registers a custom type for a FFI target.

funcjax._src.ffi.register_ffi_type_id(name:str, obj:Any, platform:str='cpu') -> None

Registers a custom type ID for a FFI target.

funcjax._src.flatten_util.ravel_pytree(pytree:Any) -> tuple[Array, Callable[[Array], Any]]

Ravel (flatten) a pytree of arrays down to a 1D array.

classjax._src.flattree.FlatTree

FlatTree is a Python OOP version of this functor.

funcjax._src.hardware_utils.get_shm_size_in_mb()

Get /dev/shm size in MB.

funcjax._src.hardware_utils.num_available_amd_gpus(stop_at:int | None=None) -> int

Count AMD GPUs available via KFD kernel driver.

classjax._src.hypothesis_test_util.HypothesisShardedTestCase

Runs Hypothesis tests in a sharded manner.

classjax._src.hypothesis_test_util.HypothesisShardedTestLoader

A TestLoader that bypasses method-level sharding.

funcjax._src.hypothesis_test_util.setup_hypothesis(max_examples=30) -> None

Sets up the hypothesis profiles.

classjax._src.image.scale.ResizeMethod

Image resize method.

func

jax._src.image.scale.resize(image, shape:core.Shape, method:str | ResizeMethod, antialias:bool=True, precision=lax.Precision.HIGHEST)

Image resize.

classjax._src.indexing.Slice

A slice with a start index and a size.

classjax._src.interpreters.mlir.LoweringRuleContext

Per-rule context information for MLIR lowering.

classjax._src.interpreters.mlir.ModuleContext

Module-wide context information for MLIR lowering.

funcjax._src.interpreters.mlir.aval_to_ir_type(ctx:ModuleContext, aval:core.AbstractValue) -> ir.Type

Converts a JAX aval to a single MLIR IR type.

funcjax._src.interpreters.mlir.aval_to_ir_types(ctx:ModuleContext, aval:core.AbstractValue) -> tuple[ir.Type, ...]

Converts a JAX aval to one or more MLIR IR types.

funcjax._src.interpreters.mlir.compare_hlo(x, y, direction:str, comparison_type:str | None=None)

Creates CompareOp.

funcjax._src.interpreters.mlir.convert_hlo(ctx:LoweringRuleContext, x, aval_in, aval_out)

Variant of convert that has HLO semantics.

funcjax._src.interpreters.mlir.delegate_lowering(ctx, lowering_fun, *args, **ctx_override_kwargs)

Side-effects on `ctx`

funcjax._src.interpreters.mlir.dump_module_to_file(module:ir.Module, stage_name:str) -> str | None

Dumps the `module` IR to a file.

funcjax._src.interpreters.mlir.eval_dynamic_shape_as_tensor(ctx:LoweringRuleContext, shape:core.Shape) -> ir.Value

Evaluates the dynamic shapes as one 1d int32 tensor.

funcjax._src.interpreters.mlir.eval_dynamic_shape_as_vals(ctx:LoweringRuleContext, shape:core.Shape) -> tuple[ir.Value, ...]

Evaluates the dynamic shapes as int32 values.

funcjax._src.interpreters.mlir.ir_attribute(val:Any) -> ir.Attribute

Convert a Python value to an MLIR attribute.

funcjax._src.interpreters.mlir.lower_fun(fun:Callable, multiple_results:bool=True) -> Callable

Converts a traceable JAX function `fun` into a lowering rule.

funcjax._src.interpreters.mlir.make_ir_context() -> ir.Context

Creates an MLIR context suitable for JAX IR.

funcjax._src.interpreters.mlir.refine_polymorphic_shapes(module:ir.Module) -> ir.Module

Refines the polymorphic shapes inside a module.

func

jax._src.interpreters.mlir.register_lowering(prim:core.Primitive, rule:LoweringRule, platform:str | None=None, inline:bool=True, cacheable:bool=True) -> None

Registers a lowering rule for a primitive.

funcjax._src.interpreters.mlir.sanitize_name(name:str) -> str

Ensure a name is usable as module or function name.

funcjax._src.interpreters.mlir.sharded_aval(aval:core.AbstractValue, sharding:JSharding | None) -> core.AbstractValue

Returns the new aval sharded based on sharding proto.

funcjax._src.interpreters.partial_eval.convert_constvars_jaxpr(jaxpr:Jaxpr) -> Jaxpr

Detaches the consts, exposing the constant inputs as leading invars.

func

jax._src.interpreters.partial_eval.dce_jaxpr(jaxpr:Jaxpr, used_outputs:bool | Sequence[bool], instantiate:bool | Sequence[bool]=False) -> tuple[Jaxpr, list[bool]]

Runs dead-code elementation on a given jaxpr.

funcjax._src.interpreters.partial_eval.separate_consts(jaxpr:Jaxpr) -> tuple[Jaxpr, list[Any]]

Detaches the consts and returns them explicitly.

classjax._src.interpreters.pxla.AllArgsInfo

Avals and debug_info for all arguments prior to DCE.

funcjax._src.interpreters.pxla.clear_in_memory_compilation_cache() -> None

Clears the in-memory compilation cache.

funcjax._src.jaxpr_util.eqns_using_var(jaxpr:core.Jaxpr, invar:core.Var) -> Iterator[core.JaxprEqn]

Find the leaf equations using a variable

funcjax._src.jaxpr_util.jaxpr_to_html(jaxpr:core.Jaxpr) -> str

Renders a Jaxpr as HTML with interactive tracebacks and search.

funcjax._src.jaxpr_util.maybe_dump_jaxpr_to_file(fun_name:str, jaxpr:core.Jaxpr) -> str | None

Maybe dumps the `jaxpr` to a file.

classjax._src.lax.control_flow.conditionals.CondSum

A cond-shaped sum, represented as a tagged product.

funcjax._src.lax.control_flow.loops.cumlogsumexp(operand:Array, axis:int=0, reverse:bool=False) -> Array

Computes a cumulative logsumexp along `axis`.

funcjax._src.lax.control_flow.loops.cummax(operand:Array, axis:int=0, reverse:bool=False) -> Array

Computes a cumulative maximum along `axis`.

funcjax._src.lax.control_flow.loops.cummin(operand:Array, axis:int=0, reverse:bool=False) -> Array

Computes a cumulative minimum along `axis`.

funcjax._src.lax.control_flow.loops.cumprod(operand:Array, axis:int=0, reverse:bool=False) -> Array

Computes a cumulative product along `axis`.

funcjax._src.lax.control_flow.loops.cumsum(operand:Array, axis:int=0, reverse:bool=False) -> Array

Computes a cumulative sum along `axis`.

funcjax._src.lax.control_flow.loops.map(f, xs, *batch_size:int | None=None)

Map a function over leading array axes.

funcjax._src.lax.eval_jaxpr.create_call_primitive(name:str) -> core.Primitive

Creates a JAX Primitive with standard call rules registered.

classjax._src.lax.fft.FftType

Describes which FFT operation to perform.

classjax._src.lax.lax.DotAlgorithm

Specify the algorithm used for computing dot products.

classjax._src.lax.lax.DotAlgorithmPreset

An enum of known algorithms for computing dot products.

classjax._src.lax.lax.Precision

Precision enum for lax matrix multiply related functions.

classjax._src.lax.lax.RandomAlgorithm

Describes which PRNG algorithm to use for rng_bit_generator.

classjax._src.lax.lax.Tolerance

Specify the tolerances used for computing unary functions.

funcjax._src.lax.lax.abs(x:ArrayLike) -> Array

Elementwise absolute value: :math:`|x|`.

funcjax._src.lax.lax.acos(x:ArrayLike) -> Array

Elementwise arc cosine: :math:`\mathrm{acos}(x)`.

funcjax._src.lax.lax.acosh(x:ArrayLike) -> Array

Elementwise inverse hyperbolic cosine: :math:`\mathrm{acosh}(x)`.

funcjax._src.lax.lax.add(x:ArrayLike, y:ArrayLike) -> Array

Elementwise addition: :math:`x + y`.

funcjax._src.lax.lax.after_all(*operands)

Merges one or more XLA token values.

funcjax._src.lax.lax.argmax(operand:ArrayLike, axis:int, index_dtype:DTypeLike) -> Array

Computes the index of the maximum element along ``axis``.

funcjax._src.lax.lax.argmin(operand:ArrayLike, axis:int, index_dtype:DTypeLike) -> Array

Computes the index of the minimum element along ``axis``.

funcjax._src.lax.lax.asarray(x:ArrayLike) -> Array

Lightweight conversion of ArrayLike input to Array output.

funcjax._src.lax.lax.asin(x:ArrayLike) -> Array

Elementwise arc sine: :math:`\mathrm{asin}(x)`.

funcjax._src.lax.lax.asinh(x:ArrayLike) -> Array

Elementwise inverse hyperbolic sine: :math:`\mathrm{asinh}(x)`.

funcjax._src.lax.lax.atan(x:ArrayLike) -> Array

Elementwise arc tangent: :math:`\mathrm{atan}(x)`.

funcjax._src.lax.lax.atan2(x:ArrayLike, y:ArrayLike) -> Array

Elementwise two-term arc tangent: :math:`\mathrm{atan}({x \over y})`.

funcjax._src.lax.lax.atanh(x:ArrayLike) -> Array

Elementwise inverse hyperbolic tangent: :math:`\mathrm{atanh}(x)`.

funcjax._src.lax.lax.batch_matmul(lhs:Array, rhs:Array, precision:PrecisionLike=None) -> Array

Batch matrix multiplication.

funcjax._src.lax.lax.bitcast_convert_type(operand:ArrayLike, new_dtype:DTypeLike) -> Array

Elementwise bitcast.

funcjax._src.lax.lax.bitwise_and(x:ArrayLike, y:ArrayLike) -> Array

Elementwise AND: :math:`x \wedge y`.

funcjax._src.lax.lax.bitwise_not(x:ArrayLike) -> Array

Elementwise NOT: :math:`\neg x`.

funcjax._src.lax.lax.bitwise_or(x:ArrayLike, y:ArrayLike) -> Array

Elementwise OR: :math:`x \vee y`.

funcjax._src.lax.lax.bitwise_xor(x:ArrayLike, y:ArrayLike) -> Array

Elementwise exclusive OR: :math:`x \oplus y`.

funcjax._src.lax.lax.broadcast(operand:ArrayLike, sizes:Sequence[int], *out_sharding=None) -> Array

Broadcasts an array, adding new leading dimensions only.

func

jax._src.lax.lax.broadcast_in_dim(operand:ArrayLike, shape:Shape, broadcast_dimensions:Sequence[int], *out_sharding=None) -> Array

General broadcasting operation.

funcjax._src.lax.lax.broadcast_to_rank(x:ArrayLike, rank:int) -> Array

Adds leading dimensions of ``1`` to give ``x`` rank ``rank``.

funcjax._src.lax.lax.broadcasted_iota(dtype:DTypeLike, shape:Shape, dimension:int, *out_sharding=None) -> Array

Convenience wrapper around ``iota``.

funcjax._src.lax.lax.canonicalize_precision(precision:PrecisionLike) -> CanonicalPrecision

Turns an API precision specification into a pair of enumeration values.

funcjax._src.lax.lax.cbrt(x:ArrayLike, *accuracy:Tolerance | AccuracyMode | None=None) -> Array

Elementwise cube root: :math:`\sqrt[3]{x}`.

funcjax._src.lax.lax.ceil(x:ArrayLike) -> Array

Elementwise ceiling: :math:`\left\lceil x \right\rceil`.

funcjax._src.lax.lax.check_same_dtypes(name:str, *avals:ShapedArray) -> None

Check that dtypes agree, possibly ignoring float precision.

funcjax._src.lax.lax.clamp(min:ArrayLike, x:ArrayLike, max:ArrayLike) -> Array

Elementwise clamp.

funcjax._src.lax.lax.clz(x:ArrayLike) -> Array

Elementwise count-leading-zeros.

funcjax._src.lax.lax.collapse(operand:Array, start_dimension:int, stop_dimension:int | None=None) -> Array

Collapses dimensions of an array into a single dimension.

funcjax._src.lax.lax.complex(x:ArrayLike, y:ArrayLike) -> Array

Elementwise make complex number: :math:`x + jy`.

funcjax._src.lax.lax.concatenate(operands:Array | Sequence[ArrayLike], dimension:int) -> Array

Concatenates a sequence of arrays along `dimension`.

funcjax._src.lax.lax.conj(x:ArrayLike) -> Array

Elementwise complex conjugate function: :math:`\overline{x}`.

funcjax._src.lax.lax.convert_element_type(operand:ArrayLike, new_dtype:DTypeLike | dtypes.ExtendedDType) -> Array

Elementwise cast.

funcjax._src.lax.lax.cos(x:ArrayLike, *accuracy:Tolerance | AccuracyMode | None=None) -> Array

Elementwise cosine: :math:`\mathrm{cos}(x)`.

funcjax._src.lax.lax.cosh(x:ArrayLike) -> Array

Elementwise hyperbolic cosine: :math:`\mathrm{cosh}(x)`.

funcjax._src.lax.lax.div(x:ArrayLike, y:ArrayLike) -> Array

Elementwise division: :math:`x \over y`.

func

jax._src.lax.lax.dot(lhs:ArrayLike, rhs:ArrayLike, *dimension_numbers:DotDimensionNumbers | None=None, *precision:PrecisionLike=None, *preferred_element_type:DTypeLike | None=None, *out_sharding=None) -> Array

General dot product/contraction operator.

func

jax._src.lax.lax.dot_general(lhs:ArrayLike, rhs:ArrayLike, dimension_numbers:DotDimensionNumbers, precision:PrecisionLike=None, preferred_element_type:DTypeLike | None=None, *out_sharding=None) -> Array

Alias of :func:`jax.lax.dot`.

funcjax._src.lax.lax.eq(x:ArrayLike, y:ArrayLike) -> Array

Elementwise equals: :math:`x = y`.

funcjax._src.lax.lax.exp(x:ArrayLike, *accuracy:Tolerance | AccuracyMode | None=None) -> Array

Elementwise exponential: :math:`e^x`.

funcjax._src.lax.lax.exp2(x:ArrayLike, *accuracy:Tolerance | AccuracyMode | None=None) -> Array

Elementwise base-2 exponential: :math:`2^x`.

funcjax._src.lax.lax.expand_dims(array:ArrayLike, dimensions:Sequence[int]) -> Array

Insert any number of size 1 dimensions into an array.

funcjax._src.lax.lax.expm1(x:ArrayLike, *accuracy:Tolerance | AccuracyMode | None=None) -> Array

Elementwise :math:`e^{x} - 1`.

funcjax._src.lax.lax.floor(x:ArrayLike) -> Array

Elementwise floor: :math:`\left\lfloor x \right\rfloor`.

func

jax._src.lax.lax.full(shape:Shape, fill_value:ArrayLike, dtype:DTypeLike | None=None, *sharding:Sharding | None=None) -> Array

Returns an array of `shape` filled with `fill_value`.

funcjax._src.lax.lax.ge(x:ArrayLike, y:ArrayLike) -> Array

Elementwise greater-than-or-equals: :math:`x \geq y`.

funcjax._src.lax.lax.gt(x:ArrayLike, y:ArrayLike) -> Array

Elementwise greater-than: :math:`x > y`.

funcjax._src.lax.lax.imag(x:ArrayLike) -> Array

Elementwise extract imaginary part: :math:`\mathrm{Im}(x)`.

funcjax._src.lax.lax.integer_pow(x:ArrayLike, y:int) -> Array

Elementwise power: :math:`x^y`, where :math:`y` is a static integer.

funcjax._src.lax.lax.is_finite(x:ArrayLike) -> Array

Elementwise :math:`\mathrm{isfinite}`.

funcjax._src.lax.lax.le(x:ArrayLike, y:ArrayLike) -> Array

Elementwise less-than-or-equals: :math:`x \leq y`.

funcjax._src.lax.lax.log(x:ArrayLike, *accuracy:Tolerance | AccuracyMode | None=None) -> Array

Elementwise natural logarithm: :math:`\mathrm{log}(x)`.

funcjax._src.lax.lax.log1p(x:ArrayLike, *accuracy:Tolerance | AccuracyMode | None=None) -> Array

Elementwise :math:`\mathrm{log}(1 + x)`.

funcjax._src.lax.lax.lt(x:ArrayLike, y:ArrayLike) -> Array

Elementwise less-than: :math:`x < y`.

funcjax._src.lax.lax.max(x:ArrayLike, y:ArrayLike) -> Array

Elementwise maximum: :math:`\mathrm{max}(x, y)`.

funcjax._src.lax.lax.mul(x:ArrayLike, y:ArrayLike, *out_dtype:DTypeLike | None=None) -> Array

Elementwise multiplication: :math:`x \times y`.

funcjax._src.lax.lax.mulhi(x:ArrayLike, y:ArrayLike) -> Array

Elementwise multiply-high: high bits of :math:`x \times y`.

funcjax._src.lax.lax.ne(x:ArrayLike, y:ArrayLike) -> Array

Elementwise not-equals: :math:`x \neq y`.

funcjax._src.lax.lax.neg(x:ArrayLike) -> Array

Elementwise negation: :math:`-x`.

funcjax._src.lax.lax.pad(operand:ArrayLike, padding_value:ArrayLike, padding_config:Sequence[tuple[int, int, int]]) -> Array

Applies low, high, and/or interior padding to an array.

funcjax._src.lax.lax.population_count(x:ArrayLike) -> Array

Elementwise popcount, count the number of set bits in each element.

funcjax._src.lax.lax.pow(x:ArrayLike, y:ArrayLike) -> Array

Elementwise power: :math:`x^y`.

func

jax._src.lax.lax.ragged_dot(lhs:Array, rhs:Array, group_sizes:Array, precision:PrecisionLike=None, preferred_element_type:DTypeLike | None=None, group_offset:Array | None=None, out_sharding:NamedSharding | P | None=None) -> Array

Ragged matrix multiplication.

funcjax._src.lax.lax.real(x:ArrayLike) -> Array

Elementwise extract real part: :math:`\mathrm{Re}(x)`.

funcjax._src.lax.lax.reciprocal(x:ArrayLike) -> Array

Elementwise reciprocal: :math:`1 \over x`.

funcjax._src.lax.lax.reduce_and(operand:ArrayLike, axes:Sequence[int]) -> Array

Compute the bitwise AND of elements over one or more array axes.

funcjax._src.lax.lax.reduce_max(operand:ArrayLike, axes:Sequence[int], *out_sharding=None) -> Array

Compute the maximum of elements over one or more array axes.

funcjax._src.lax.lax.reduce_min(operand:ArrayLike, axes:Sequence[int], *out_sharding=None) -> Array

Compute the minimum of elements over one or more array axes.

funcjax._src.lax.lax.reduce_or(operand:ArrayLike, axes:Sequence[int]) -> Array

Compute the bitwise OR of elements over one or more array axes.

funcjax._src.lax.lax.reduce_prod(operand:ArrayLike, axes:Sequence[int]) -> Array

Compute the product of elements over one or more array axes.

funcjax._src.lax.lax.reduce_sum(operand:ArrayLike, axes:Sequence[int], *out_sharding=None) -> Array

Compute the sum of elements over one or more array axes.

funcjax._src.lax.lax.reduce_xor(operand:ArrayLike, axes:Sequence[int]) -> Array

Compute the bitwise XOR of elements over one or more array axes.

funcjax._src.lax.lax.rem(x:ArrayLike, y:ArrayLike) -> Array

Elementwise remainder: :math:`x \bmod y`.

funcjax._src.lax.lax.rng_uniform(a, b, shape)

Stateful PRNG generator.

funcjax._src.lax.lax.round(x:ArrayLike, rounding_method:RoundingMethod=RoundingMethod.AWAY_FROM_ZERO) -> Array

Elementwise round.

funcjax._src.lax.lax.rsqrt(x:ArrayLike, *accuracy:Tolerance | AccuracyMode | None=None) -> Array

Elementwise reciprocal square root: :math:`1 \over \sqrt{x}`.

funcjax._src.lax.lax.select(pred:ArrayLike, on_true:ArrayLike, on_false:ArrayLike) -> Array

Selects between two branches based on a boolean predicate.

funcjax._src.lax.lax.select_n(which:ArrayLike, *cases:ArrayLike) -> Array

Selects array values from multiple cases.

funcjax._src.lax.lax.shift_left(x:ArrayLike, y:ArrayLike) -> Array

Elementwise left shift: :math:`x \ll y`.

funcjax._src.lax.lax.shift_right_arithmetic(x:ArrayLike, y:ArrayLike) -> Array

Elementwise arithmetic right shift: :math:`x \gg y`.

funcjax._src.lax.lax.shift_right_logical(x:ArrayLike, y:ArrayLike) -> Array

Elementwise logical right shift: :math:`x \gg y`.

funcjax._src.lax.lax.sign(x:ArrayLike) -> Array

Elementwise sign.

funcjax._src.lax.lax.sin(x:ArrayLike, *accuracy:Tolerance | AccuracyMode | None=None) -> Array

Elementwise sine: :math:`\mathrm{sin}(x)`.

funcjax._src.lax.lax.sinh(x:ArrayLike) -> Array

Elementwise hyperbolic sine: :math:`\mathrm{sinh}(x)`.

funcjax._src.lax.lax.split(operand:ArrayLike, sizes:Sequence[DimSize], axis:int=0) -> Sequence[Array]

Splits an array along ``axis``.

funcjax._src.lax.lax.sqrt(x:ArrayLike, *accuracy:Tolerance | AccuracyMode | None=None) -> Array

Elementwise square root: :math:`\sqrt{x}`.

funcjax._src.lax.lax.square(x:ArrayLike) -> Array

Elementwise square: :math:`x^2`.

funcjax._src.lax.lax.squeeze(array:ArrayLike, dimensions:Sequence[int]) -> Array

Squeeze any number of size 1 dimensions from an array.

funcjax._src.lax.lax.stack(operands:Sequence[ArrayLike], axis:int=0) -> Array

Joins a sequence of arrays along a new axis.

funcjax._src.lax.lax.stage(x:ArrayLike) -> Array

Lifts a value into a trace.

funcjax._src.lax.lax.stop_gradient(x:T) -> T

Stops gradient computation.

funcjax._src.lax.lax.sub(x:ArrayLike, y:ArrayLike) -> Array

Elementwise subtraction: :math:`x - y`.

funcjax._src.lax.lax.tan(x:ArrayLike, *accuracy:Tolerance | AccuracyMode | None=None) -> Array

Elementwise tangent: :math:`\mathrm{tan}(x)`.

funcjax._src.lax.lax.tanh(x:ArrayLike, *accuracy:Tolerance | AccuracyMode | None=None) -> Array

Elementwise hyperbolic tangent: :math:`\mathrm{tanh}(x)`.

funcjax._src.lax.lax.tile(operand:ArrayLike, reps:Sequence[int]) -> Array

Tiles an array by repeating it along each dimension.

funcjax._src.lax.lax.unstack(x:ArrayLike, axis:int=0) -> tuple[Array, ...]

Unstacks an array along an axis.

classjax._src.lax.linalg.EigImplementation

Enum for eigendecomposition algorithm.

classjax._src.lax.linalg.EighImplementation

Implementation for symmetric/Hermitian eigendecomposition.

classjax._src.lax.linalg.SvdAlgorithm

Enum for SVD algorithm.

funcjax._src.lax.linalg.cholesky(x:Array, *symmetrize_input:bool=True) -> Array

Cholesky decomposition.

funcjax._src.lax.linalg.cholesky_update(r_matrix:ArrayLike, w_vector:ArrayLike) -> Array

Cholesky rank-1 update.

funcjax._src.lax.linalg.geqp3(a:ArrayLike, jpvt:ArrayLike, *use_magma:bool | None=None) -> tuple[Array, Array, Array]

Computes the column-pivoted QR decomposition of a matrix.

funcjax._src.lax.linalg.geqrf(a:ArrayLike) -> tuple[Array, Array]

Computes the QR decomposition of a matrix.

funcjax._src.lax.linalg.hessenberg(a:ArrayLike) -> tuple[Array, Array]

Reduces a square matrix to upper Hessenberg form.

funcjax._src.lax.linalg.householder_product(a:ArrayLike, taus:ArrayLike) -> Array

Product of elementary Householder reflectors.

funcjax._src.lax.linalg.lu(x:ArrayLike) -> tuple[Array, Array, Array]

LU decomposition with partial pivoting.

funcjax._src.lax.linalg.lu_solve(lu:ArrayLike, permutation:ArrayLike, b:ArrayLike, trans:int=0) -> Array

LU solve with broadcasting.

func

jax._src.lax.linalg.schur(x:ArrayLike, *compute_schur_vectors:bool=True, *sort_eig_vals:bool=False, *select_callable:Callable[..., Any] | None=None) -> tuple[Array, Array]

Schur decomposition.

func

jax._src.lax.linalg.symmetric_product(a_matrix:ArrayLike, c_matrix:ArrayLike, *alpha:float=1.0, *beta:float=0.0, *symmetrize_output:bool=False)

Symmetric product.

func

jax._src.lax.linalg.triangular_solve(a:ArrayLike, b:ArrayLike, *left_side:bool=False, *lower:bool=False, *transpose_a:bool=False, *conjugate_a:bool=False, *unit_diagonal:bool=False) -> Array

Triangular solve.

funcjax._src.lax.linalg.tridiagonal(a:ArrayLike, *lower:bool=True) -> tuple[Array, Array, Array, Array]

Reduces a symmetric/Hermitian matrix to tridiagonal form.

funcjax._src.lax.linalg.tridiagonal_solve(dl:Array, d:Array, du:Array, b:Array, *perturb_singular:bool=False) -> Array

Computes the solution of a tridiagonal linear system.

funcjax._src.lax.other.logaddexp(x1:ArrayLike, x2:ArrayLike) -> Array

Compute log(exp(x1) + exp(x2)) avoiding overflow.

funcjax._src.lax.other.logaddexp2(x1:ArrayLike, x2:ArrayLike) -> Array

Compute log2(exp2(x1) + exp2(x2)) avoiding overflow.

funcjax._src.lax.parallel.all_gather(x, axis_name, *axis_index_groups=None, *axis=0, *tiled=False, *to:str='varying')

Gather values of x across all replicas.

funcjax._src.lax.parallel.all_gather_invariant(x, axis_name, *axis:int=0, *tiled:bool=False)

Gather values of x across all replicas.

funcjax._src.lax.parallel.axis_index(axis_name:AxisName) -> Array

Return the index along the mapped axis ``axis_name``.

funcjax._src.lax.parallel.axis_size(axis_name:AxisName) -> int

Return the size of the mapped axis ``axis_name``.

func

jax._src.lax.slicing.gather(operand:ArrayLike, start_indices:ArrayLike, dimension_numbers:GatherDimensionNumbers, slice_sizes:Shape, *unique_indices:bool=False, *indices_are_sorted:bool=False, *mode:str | GatherScatterMode | None=None, *fill_value=None) -> Array

Gather operator.

func

jax._src.lax.slicing.scatter(operand:ArrayLike, scatter_indices:ArrayLike, updates:ArrayLike, dimension_numbers:ScatterDimensionNumbers, *indices_are_sorted:bool=False, *unique_indices:bool=False, *mode:str | GatherScatterMode | None=None) -> Array

Scatter-update operator.

func

jax._src.lax.slicing.scatter_add(operand:ArrayLike, scatter_indices:ArrayLike, updates:ArrayLike, dimension_numbers:ScatterDimensionNumbers, *indices_are_sorted:bool=False, *unique_indices:bool=False, *mode:str | GatherScatterMode | None=None) -> Array

Scatter-add operator.

func

jax._src.lax.slicing.scatter_max(operand:ArrayLike, scatter_indices:ArrayLike, updates:ArrayLike, dimension_numbers:ScatterDimensionNumbers, *indices_are_sorted:bool=False, *unique_indices:bool=False, *mode:str | GatherScatterMode | None=None) -> Array

Scatter-max operator.

func

jax._src.lax.slicing.scatter_min(operand:ArrayLike, scatter_indices:ArrayLike, updates:ArrayLike, dimension_numbers:ScatterDimensionNumbers, *indices_are_sorted:bool=False, *unique_indices:bool=False, *mode:str | GatherScatterMode | None=None) -> Array

Scatter-min operator.

func

jax._src.lax.slicing.scatter_mul(operand:ArrayLike, scatter_indices:ArrayLike, updates:ArrayLike, dimension_numbers:ScatterDimensionNumbers, *indices_are_sorted:bool=False, *unique_indices:bool=False, *mode:str | GatherScatterMode | None=None) -> Array

Scatter-multiply operator.

func

jax._src.lax.slicing.scatter_sub(operand:ArrayLike, scatter_indices:ArrayLike, updates:ArrayLike, dimension_numbers:ScatterDimensionNumbers, *indices_are_sorted:bool=False, *unique_indices:bool=False, *mode:str | GatherScatterMode | None=None) -> Array

Scatter-sub operator.

funcjax._src.lax.special.betainc(a:ArrayLike, b:ArrayLike, x:ArrayLike) -> Array

Elementwise regularized incomplete beta integral.

funcjax._src.lax.special.digamma(x:ArrayLike) -> Array

Elementwise digamma: :math:`\psi(x)`.

funcjax._src.lax.special.erf(x:ArrayLike) -> Array

Elementwise error function: :math:`\mathrm{erf}(x)`.

funcjax._src.lax.special.erf_inv(x:ArrayLike) -> Array

Elementwise inverse error function: :math:`\mathrm{erf}^{-1}(x)`.

funcjax._src.lax.special.igamma(a:ArrayLike, x:ArrayLike) -> Array

Elementwise regularized incomplete gamma function.

funcjax._src.lax.special.igamma_grad_a(a:ArrayLike, x:ArrayLike) -> Array

Elementwise derivative of the regularized incomplete gamma function.

funcjax._src.lax.special.igammac(a:ArrayLike, x:ArrayLike) -> Array

Elementwise complementary regularized incomplete gamma function.

funcjax._src.lax.special.lgamma(x:ArrayLike) -> Array

Elementwise log gamma: :math:`\mathrm{log}(\Gamma(x))`.

funcjax._src.lax.special.ndtr(x:Array) -> Array

Elementwise cumulative distribution function of the standard normal.

funcjax._src.lax.special.polygamma(m:ArrayLike, x:ArrayLike) -> Array

Elementwise polygamma: :math:`\psi^{(m)}(x)`.

funcjax._src.lax.special.random_gamma_grad(a:ArrayLike, x:ArrayLike, *dtype) -> Array

Elementwise derivative of samples from `Gamma(a, 1)`.

funcjax._src.lax.special.zeta(x:ArrayLike, q:ArrayLike) -> Array

Elementwise Hurwitz zeta function: :math:`\zeta(x, q)`

func

jax._src.lax.utils.index_dtype_for_axis_size(indices_dtype:DType, axis_size:DimSize, wrap_negative_indices:bool) -> DType

Upcast indices_dtype if necessary to avoid overflow.

funcjax._src.lax.utils.int_dtype_for_dim(d:DimSize, *signed:bool) -> DType

Returns a integer dtype large enough to contain indices in dimension d.

funcjax._src.lax.utils.int_dtype_for_shape(shape:Shape, *signed:bool) -> DType

Returns a integer dtype large enough to contain indices in `shape`.

funcjax._src.lax_reference.ragged_dot(lhs, rhs, group_sizes)

Reference ragged dot implementation.

func

jax._src.lazy_loader.attach(package_name:str, submodules:Sequence[str]) -> tuple[Callable[[str], Any], Callable[[], list[str]], list[str]]

Lazily loads submodules of a package.

classjax._src.linear_util.DebugInfo

Debugging info about a func, its arguments, and results.

methodjax._src.linear_util.DebugInfo.filter_arg_names(keep:Sequence[bool]) -> tuple[str, ...] | None

Keep only the arg_names for which `keep` is True.

methodjax._src.linear_util.DebugInfo.filter_result_paths(keep:Sequence[bool]) -> tuple[str, ...] | None

Keep only the result_paths for which `keep` is True.

methodjax._src.linear_util.DebugInfo.resolve_result_paths() -> DebugInfo

Return a debug info with resolved result paths.

methodjax._src.linear_util.DebugInfo.safe_arg_names(expected_count:int) -> tuple[str, ...]

Get the arg_names with a safety check.

methodjax._src.linear_util.DebugInfo.safe_result_paths(expected_count:int) -> tuple[str, ...]

Get the result paths with a safety check.

classjax._src.linear_util.WrappedFun

Represents a function `f` to which `transforms` are to be applied.

methodjax._src.linear_util.WrappedFun.call_wrapped(*args, **kwargs)

Calls the transformed function

methodjax._src.linear_util.WrappedFun.wrap(gen, gen_static_args, out_store:Store | EqualStore | None) -> WrappedFun

Add another transform and its store.

funcjax._src.linear_util.transformation2(gen, fun:WrappedFun, *gen_static_args) -> WrappedFun

Adds one more transformation to a WrappedFun.

funcjax._src.linear_util.wrap_init(f:Callable, params=None, *debug_info:DebugInfo) -> WrappedFun

Wraps function `f` as a `WrappedFun`, suitable for transformation.

classjax._src.literals.TypedNdArray

A TypedNdArray is a host-side array used by JAX during tracing.

classjax._src.lru_cache.LRUCache

Bounded cache with least-recently-used (LRU) eviction policy.

methodjax._src.lru_cache.LRUCache.get(key:str) -> bytes | None

Retrieves the cached value for the given key.

methodjax._src.lru_cache.LRUCache.put(key:str, value:bytes) -> None

Adds a new entry to the cache.

classjax._src.mesh.AbstractMesh

AbstractMesh contains only axis names and axis sizes.

classjax._src.mesh.use_abstract_mesh

Sets a abstract mesh in a thread-local context.

funcjax._src.monitoring.clear_event_listeners()

Clear event listeners.

funcjax._src.monitoring.get_event_duration_listeners() -> list[EventDurationListenerWithMetadata]

Get event duration listeners.

funcjax._src.monitoring.get_event_listeners() -> list[EventListenerWithMetadata]

Get event listeners.

funcjax._src.monitoring.get_event_time_span_listeners() -> list[EventTimeSpanListenerWithMetadata]

Get event time span listeners.

funcjax._src.monitoring.get_scalar_listeners() -> list[ScalarListenerWithMetadata]

Get scalar event listeners.

funcjax._src.monitoring.record_event(event:str, **kwargs:str | int) -> None

Record an event.

funcjax._src.monitoring.record_event_duration_secs(event:str, duration:float, **kwargs:str | int) -> None

Record an event duration in seconds (float).

funcjax._src.monitoring.record_event_time_span(event:str, start_time:float, end_time:float, **kwargs:str | int) -> None

Record an event start and end time in seconds (float).

funcjax._src.monitoring.record_scalar(event:str, value:float | int, **kwargs:str | int) -> None

Record a scalar summary value.

funcjax._src.monitoring.register_event_listener(callback:EventListenerWithMetadata) -> None

funcjax._src.monitoring.register_scalar_listener(callback:ScalarListenerWithMetadata) -> None

funcjax._src.monitoring.unregister_event_duration_listener(callback:EventDurationListenerWithMetadata) -> None

Unregister an event duration listener by callback.

funcjax._src.monitoring.unregister_event_listener(callback:EventListenerWithMetadata) -> None

Unregister an event listener by callback.

funcjax._src.monitoring.unregister_event_time_span_listener(callback:EventTimeSpanListenerWithMetadata) -> None

Unregister an event time span listener by callback.

funcjax._src.monitoring.unregister_scalar_listener(callback:ScalarListenerWithMetadata) -> None

Unregister a scalar event listener by callback.

classjax._src.named_sharding.NamedSharding

A :class:`NamedSharding` expresses sharding using named axes.

funcjax._src.nn.functions.celu(x:ArrayLike, alpha:ArrayLike=1.0) -> Array

Continuously-differentiable exponential linear unit activation.

funcjax._src.nn.functions.elu(x:ArrayLike, alpha:ArrayLike=1.0) -> Array

Exponential linear unit activation function.

funcjax._src.nn.functions.gelu(x:ArrayLike, approximate:bool=True) -> Array

Gaussian error linear unit activation function.

funcjax._src.nn.functions.glu(x:ArrayLike, axis:int=-1) -> Array

Gated linear unit activation function.

funcjax._src.nn.functions.hard_sigmoid(x:ArrayLike) -> Array

Hard Sigmoid activation function.

funcjax._src.nn.functions.hard_tanh(x:ArrayLike) -> Array

Hard :math:`\mathrm{tanh}` activation function.

funcjax._src.nn.functions.identity(x:ArrayLike) -> Array

Identity activation function.

funcjax._src.nn.functions.leaky_relu(x:ArrayLike, negative_slope:ArrayLike=0.01) -> Array

Leaky rectified linear unit activation function.

funcjax._src.nn.functions.log1mexp(x:ArrayLike) -> Array

Numerically stable calculation of :math:`\log(1 - \exp(-x))`.

funcjax._src.nn.functions.log_sigmoid(x:ArrayLike) -> Array

Log-sigmoid activation function.

funcjax._src.nn.functions.log_softmax(x:ArrayLike, axis:Axis=-1, where:ArrayLike | None=None) -> Array

Log-Softmax function.

funcjax._src.nn.functions.logmeanexp(x:ArrayLike, axis:Axis=None, where:ArrayLike | None=None, keepdims:bool=False) -> Array

Log mean exp.

funcjax._src.nn.functions.mish(x:ArrayLike) -> Array

Mish activation function.

func

jax._src.nn.functions.one_hot(x:Any, num_classes:int, *dtype:Any | None=None, *axis:int | AxisName=-1, *out_sharding:NamedSharding | P | None=None) -> Array

One-hot encodes the given indices.

funcjax._src.nn.functions.relu(x:ArrayLike) -> Array

Rectified linear unit activation function.

この情報について

掲載しているシグネチャは jax-ml/jax の公開ソースコードを Python の ast モジュールで静的解析し、引数名・デフォルト値・型注釈・戻り値型をそのまま抽出したものです。実装コードは保存していません。詳しくは仕組みの解説をご覧ください。

収録ライブラリ一覧（全 805 件）へ戻る