jax API reference

400 public APIs from jax (jax-ml/jax) — 60 classes, 303 functions, 37 methods. Signatures extracted by static analysis of the actual source.

Repository: jax-ml/jax

Kind	Count
Classes	60
Functions	303
Methods	37

API list

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.

About this data

These signatures were extracted from the public source of jax-ml/jax using Python's ast module. Argument names, default values, type annotations and return types are taken verbatim from the code. Implementation bodies are never stored. See how it works for details.

Back to all 805 libraries