Auto jax discipline

auto_jax_discipline

A discipline interfacing a JAX function.

Classes

AutoJAXDiscipline

AutoJAXDiscipline(
    function: Callable[
        [NumberLike, ..., Any, ...], tuple[NumberLike]
    ],
    static_args: Mapping[str, Any] = READ_ONLY_EMPTY_DICT,
    differentiation_method: DifferentiationMethod = AUTO,
    differentiate_at_execution: bool = False,
    name: str | None = None,
)

Bases: JAXDiscipline

Automatically wrap a JAX function into a discipline.

Initialize the JAXDiscipline.

Parameters:

• function (Callable[[NumberLike, ..., Any, ...], tuple[NumberLike]]) –

  The JAX function.

• static_args (Mapping[str, Any], default: READ_ONLY_EMPTY_DICT ) –

  The names and values of the static arguments of the JAX function. These arguments are constant at discipline execution. The non-numeric arguments can also be included.

• differentiation_method (DifferentiationMethod, default: AUTO ) –

  The method to compute the Jacobian.

• differentiate_at_execution (bool, default: False ) –

  Whether to compute the Jacobian when executing the discipline.

• name (str | None, default: None ) –

  The name of the discipline. If empty, use the name of the class.

Source code in src/gemseo_jax/auto_jax_discipline.py

def __init__(
    self,
    function: Callable[[NumberLike, ..., Any, ...], tuple[NumberLike]],
    static_args: Mapping[str, Any] = READ_ONLY_EMPTY_DICT,
    differentiation_method: JAXDiscipline.DifferentiationMethod = JAXDiscipline.DifferentiationMethod.AUTO,  # noqa: E501
    differentiate_at_execution: bool = False,
    name: str | None = None,
) -> None:
    """
    Args:
        function: The JAX function.
        static_args: The names and values of the static arguments of the JAX
            function. These arguments are constant at discipline execution.
            The non-numeric arguments can also be included.
    """  # noqa: D205, D212, D415, D417
    arg_names, output_names, default_args = self.__parse_function(function)
    self.__static_args = {
        input_name: input_value
        for input_name, input_value in static_args.items()
        if input_name in arg_names
    }

    input_names = list(set(arg_names) - set(self.__static_args))
    default_inputs = {
        input_name: default_value
        for input_name, default_value in default_args.items()
        if input_name not in self.__static_args
    }
    self.__original_function = function

    super().__init__(
        function=self.__function,
        input_names=input_names,
        output_names=output_names,
        default_inputs=default_inputs,
        differentiation_method=differentiation_method,
        differentiate_at_execution=differentiate_at_execution,
        name=name,
    )

Attributes

jax_out_func instance-attribute

jax_out_func: Callable[[DataType], DataType] = function

The JAX function to compute the outputs from the inputs.

Classes

DifferentiationMethod

Bases: StrEnum

The method to compute the Jacobian.

Functions

add_differentiated_inputs

add_differentiated_inputs(
    input_names: Iterable[str] = (),
) -> None

Add the inputs with respect to which to differentiate the outputs.

The inputs that do not represent continuous numbers are filtered out.

Parameters:

• input_names (Iterable[str], default: () ) –

  The input variables with respect to which to differentiate the outputs. If empty, use all the inputs.

Raises:

• ValueError –

  When an input name is not the name of a discipline input.

Notes

The Jacobian is also filtered to view non-differentiated static.

Source code in src/gemseo_jax/jax_discipline.py

def add_differentiated_inputs(
    self,
    input_names: Iterable[str] = (),
) -> None:
    """
    Notes:
        The Jacobian is also filtered to view non-differentiated static.
    """  # noqa: D205, D212, D415
    old_differentiated_inputs = self._differentiated_input_names.copy()
    super().add_differentiated_inputs(input_names=input_names)
    refilter = any(
        input_name not in old_differentiated_inputs
        for input_name in self._differentiated_input_names
    )
    if refilter:
        self._filter_jacobian()

add_differentiated_outputs

add_differentiated_outputs(
    output_names: Iterable[str] = (),
) -> None

Add the outputs to be differentiated.

The outputs that do not represent continuous numbers are filtered out.

Parameters:

• output_names (Iterable[str], default: () ) –

  The outputs to be differentiated. If empty, use all the outputs.

Raises:

• ValueError –

  When an output name is not the name of a discipline output.

Notes

The Jacobian is also filtered to view non-differentiated static.

Source code in src/gemseo_jax/jax_discipline.py

def add_differentiated_outputs(
    self,
    output_names: Iterable[str] = (),
) -> None:
    """
    Notes:
        The Jacobian is also filtered to view non-differentiated static.
    """  # noqa: D205, D212, D415
    old_differentiated_outputs = self._differentiated_output_names.copy()
    super().add_differentiated_outputs(output_names=output_names)
    refilter = any(
        output_name not in old_differentiated_outputs
        for output_name in self._differentiated_output_names
    )
    if refilter:
        self._filter_jacobian()

compile_jit

compile_jit(pre_run: bool = True) -> None

Apply jit compilation over function and jacobian.

Parameters:

• pre_run (bool, default: True ) –

  Whether to call jitted callables once to trigger compilation and log times.

Source code in src/gemseo_jax/jax_discipline.py

def compile_jit(
    self,
    pre_run: bool = True,
) -> None:
    """Apply jit compilation over function and jacobian.

    Args:
        pre_run: Whether to call jitted callables once to trigger compilation and
            log times.
    """
    self.jax_out_func = jit(self.jax_out_func)
    self.__jax_jac_func = jit(self.__jax_jac_func)
    if pre_run:
        self._jit_pre_run()