nifty.re.misc module

conditional_call(condition, fn, *args, **kwargs)

JAX JIT-safe call to fn if condition is True.

Warning:

The function fn may NOT dispatch ANY JAX code, including by comparing JAX objects! To safeguard against easy to miss fallacies, all JAX arrays are automatically converted to numpy arrays.

param condition:

If True, will call fn on the host

type condition:

boolean

param fn:

Function that will be called on the host if condition is True

type fn:

Callable

param args:

Positional arguments passed to fn

param kwargs:

Keyword arguments passed to fn

conditional_raise(condition: bool, exception)

JAX JIT-safe raise of the given Exception if condition is True.

Parameters:

• condition (bool) – If True, will raise exception on the host.

• exception (Exception) – Exception that will be raised if condition is True

deprecated(message: str)

doc_from(original)

hvp(f, primals, tangents)

interpolate(xmin=-7.0, xmax=7.0, N=14000) → Callable

Replaces a local nonlinearity such as jnp.exp with a linear interpolation

Interpolating functions speeds up code and increases numerical stability in some cases, but at a cost of precision and range.

Parameters:

• xmin (float) – Minimal interpolation value. Default: -7.

• xmax (float) – Maximal interpolation value. Default: 7.

• N (int) – Number of points used for the interpolation. Default: 14000

is_iterable_of_non_iterables(ls: Any) → bool

Indicates whether the input is one dimensional.

An object is considered one dimensional if it is an iterable of non-iterable items.

isiterable(candidate)

safeguard_arguments_against_accidental_calls_into_jax(func)

Safeguard against using JAX in callback, see !25861

split(mappable, keys)

Split a dictionary into one containing only the specified keys and one with all of the remaining ones.

wrap(call: Callable[[I], O], name: Hashable) → Callable[[Mapping[Hashable, I]], O]

wrap_left(call: Callable[[I], O], name: Hashable) → Callable[[I], Dict[Hashable, O]]