jaxlie.utils._utils

Module Contents

Functions

get_epsilon(dtype)	Helper for grabbing type-specific precision constants.
register_lie_group(*, matrix_dim, parameters_dim, ...)	Decorator for registering Lie group dataclasses.
broadcast_leading_axes(inputs)	Broadcast leading axes of arrays. Takes tuples of either:

Attributes

T
TupleOfBroadcastable

jaxlie.utils._utils.T

jaxlie.utils._utils.get_epsilon(dtype)

Helper for grabbing type-specific precision constants.

Parameters:

dtype (jax.numpy.dtype) – Datatype.

Returns:

Output float.

Return type:

float

jaxlie.utils._utils.register_lie_group(*, matrix_dim, parameters_dim, tangent_dim, space_dim)

Decorator for registering Lie group dataclasses.

Sets dimensionality class variables, and marks all methods for JIT compilation.

Parameters:

• matrix_dim (int) –

• parameters_dim (int) –

• tangent_dim (int) –

• space_dim (int) –

Return type:

Callable[[Type[T]], Type[T]]

jaxlie.utils._utils.TupleOfBroadcastable

jaxlie.utils._utils.broadcast_leading_axes(inputs)

Broadcast leading axes of arrays. Takes tuples of either:

• an array, which we assume has shape (*, D).

• a Lie group object.

Parameters:

inputs (TupleOfBroadcastable) –

Return type:

TupleOfBroadcastable