jax_explicit_comp.py

An ExplicitComponent that uses JAX for derivatives.

class openmdao.components.jax_explicit_comp.JaxExplicitComponent(matrix_free=False, fallback_derivs_method='fd', **kwargs)

Bases: ExplicitComponent

Base class for explicit components when using JAX for derivatives.

Parameters:

matrix_freebool

If True, this component will compute derivatives using matrix vector products.

fallback_derivs_methodstr

The method to use if JAX is not available. Default is ‘fd’.

**kwargsdict

Additional arguments to be passed to the base class.

Attributes:

_tangentsdict

The tangents for the inputs and outputs.

_do_sparsitybool

If True, compute the sparsity.

_sparsitycoo_matrix or None

The sparsity of the Jacobian.

_jac_func_function or None

The function that computes the jacobian.

_jac_colored_function or None

The function that computes the colored jacobian.

_static_hashtuple

The hash of the static values.

_fwd_jac_func_function or None

Cached, jitted jvp function used by matrix_free ‘fwd’ mode (see _compute_jacvec_product). Primal values are passed in as call arguments rather than captured at trace time. Unlike _vjp_fun below this is valid for the life of the Problem and is only rebuilt when discrete inputs or static config change.

_fwd_static_hashtuple or None

The (discrete inputs, get_self_statics()) key _fwd_jac_func_ was built from.

_orig_compute_primalfunction

The original compute_primal method.

_ret_tuple_compute_primalfunction

The compute_primal method that returns a tuple.

_output_shapesdict

A dict of output shapes used when shapes are computed dynamically.

_do_shape_checkbool

If True, check the declared output shapes vs. the shapes of the outputs returned from compute_primal.

_deriv_output_idxstuple of int or None

Indices, into the full list of continuous outputs, of the outputs that have at least one dependent partial declared. None if all outputs need derivatives (the common case). Outputs can be excluded by declaring their partials as dependent=False (e.g. self.declare_partials('my_output', '*', dependent=False)), which allows the jax derivative computation to skip them entirely.

_deriv_output_namestuple of str or None

The relative names corresponding to _deriv_output_idxs. None if _deriv_output_idxs is None.

Methods

abs_meta_iter(iotype[, local, cont, discrete])	Iterate over absolute variable names and their metadata for this System.
add_constraint(name[, lower, upper, equals, ...])	Add a constraint variable to this system.
add_design_var(name[, lower, upper, ref, ...])	Add a design variable to this system.
add_discrete_input(name, val[, desc, tags, ...])	Add a discrete input variable to the component.
add_discrete_output(name, val[, desc, tags, ...])	Add an output variable to the component.
add_input(name, **kwargs)	Add an input to the component.
add_objective(name[, ref, ref0, index, ...])	Add a response variable to this system.
add_output(name, **kwargs)	Add an output to the component.
add_recorder(recorder[, recurse])	Add a recorder to the system.
add_response(name, type_[, lower, upper, ...])	Add a response variable to this system.
best_partial_deriv_direction()	Return the best direction for partial deriv calculations based on input and output sizes.
check_config(logger)	Perform optional error checks.
check_partials([out_stream, compact_print, ...])	Check partial derivatives comprehensively for this component.
check_sparsity([method, max_nz, out_stream])	Check the sparsity of the computed jacobian against the declared sparsity.
cleanup()	Clean up resources prior to exit.
comm_info_iter()	Yield comm size for this system and all subsystems.
compute(inputs, outputs[, discrete_inputs, ...])	Compute outputs given inputs.
compute_fd_jac(jac[, method])	Force the use of finite difference to compute a jacobian.
compute_fd_sparsity([method, num_full_jacs, ...])	Use finite difference to compute a sparsity matrix.
compute_jacvec_product(inputs, d_inputs, ...)	Compute jac-vector product.
compute_partials(inputs, partials[, ...])	Compute sub-jacobian parts.
compute_sparsity([direction, num_iters, ...])	Get the sparsity of the Jacobian.
convert2units(name, val, units)	Convert the given value to the specified units.
convert_from_units(name, val, units)	Convert the given value from the specified units to those of the named variable.
convert_units(name, val, units_from, units_to)	Wrap the utility convert_units and give a good error message.
declare_coloring(**kwargs)	Declare coloring for this component.
declare_partials(of, wrt[, dependent, rows, ...])	Declare information about this component's subjacobians.
dist_size_iter(io, top_comm)	Yield names and distributed ranges of all local and remote variables in this system.
get_coloring_fname(mode)	Return the full pathname to a coloring file.
get_conn_graph()	Return the model connection graph.
get_constraints([recurse, get_sizes, ...])	Get the Constraint settings from this system.
get_declare_partials_calls([sparsity])	Return a string containing declare_partials() calls based on the subjac sparsity.
get_design_vars([recurse, get_sizes, ...])	Get the DesignVariable settings from this system.
get_io_metadata([iotypes, metadata_keys, ...])	Retrieve metadata for a filtered list of variables.
get_linear_vectors()	Return the linear inputs, outputs, and residuals vectors.
get_nonlinear_vectors()	Return the inputs, outputs, and residuals vectors.
get_objectives([recurse, get_sizes, ...])	Get the Objective settings from this system.
get_outputs_dir(*subdirs[, mkdir])	Get the path under which all output files of this system are to be placed.
get_promotions([inprom, outprom])	Return all promotions for the given promoted variable(s).
get_reports_dir()	Get the path to the directory where the report files should go.
get_responses([recurse, get_sizes, use_prom_ivc])	Get the response variable settings from this system.
get_self_statics()	Override this in derived classes if compute_primal references static values.
get_source(name)	Return the source variable connected to the given named variable.
get_val(name[, units, indices, get_remote, ...])	Get an output/input/residual variable.
get_var_dup_info(name, io)	Return information about how the given variable is duplicated across MPI processes.
get_var_sizes(name, io)	Return the sizes of the given variable on all procs.
has_vectors()	Check if the system vectors have been initialized.
initialize()	Perform any one-time initialization run at instantiation.
is_explicit([is_comp])	Return True if this is an explicit component.
list_inputs([val, prom_name, units, shape, ...])	Write a list of input names and other optional information to a specified stream.
list_options([include_default, ...])	Write a list of output names and other optional information to a specified stream.
list_outputs([explicit, implicit, val, ...])	Write a list of output names and other optional information to a specified stream.
list_vars([val, prom_name, residuals, ...])	Write a list of inputs and outputs sorted by component in execution order.
load_case(case)	Pull all input and output variables from a Case into this System.
load_model_options()	Load the relevant model options from Problem._metadata['model_options'].
override_method(name, method)	Dynamically add a method to this component instance.
record_iteration()	Record an iteration of the current System.
run_apply_linear(mode[, scope_out, scope_in])	Compute jac-vec product.
run_apply_nonlinear()	Compute residuals.
run_linearize([sub_do_ln])	Compute jacobian / factorization.
run_solve_linear(mode)	Apply inverse jac product.
run_solve_nonlinear()	Compute outputs.
run_validation()	Run validate method on all systems below this system.
set_check_partial_options(wrt[, method, ...])	Set options that will be used for checking partial derivatives.
set_constraint_options(name[, ref, ref0, ...])	Set options for constraints in the model.
set_design_var_options(name[, lower, upper, ...])	Set options for design vars in the model.
set_objective_options(name[, ref, ref0, ...])	Set options for objectives in the model.
set_output_solver_options(name[, lower, ...])	Set solver output options.
set_solver_print([level, depth, type_, ...])	Control printing for solvers and subsolvers in the model.
set_val(name, val[, units, indices])	Set an input or output variable.
setup()	Declare inputs and outputs.
setup_partials()	Declare partials.
sparsity_matches_fd([direction, outstream])	Compare the sparsity computed by this system vs.
subjac_sparsity_iter(sparsity[, wrt_matches])	Iterate over sparsity for each subjac in the jacobian.
system_iter([include_self, recurse, typ, ...])	Yield a generator of local subsystems of this system.
total_local_size(io)	Return the total local size of the given variable.
use_fixed_coloring([coloring, recurse])	Use a precomputed coloring for this System.
uses_approx()	Return True if the system uses approximations to compute derivatives.
validate(inputs, outputs[, discrete_inputs, ...])	Check any final input / output values after a run.

__init__(matrix_free=False, fallback_derivs_method='fd', **kwargs)

Store some bound methods so we can detect runtime overrides.

add_input(name, **kwargs)

Add an input to the component.

This overrides the base class method to update the kwargs to use dynamic shaping by default.

Parameters:

namestr

The name of the input.

**kwargsdict

The kwargs to pass to the base class method.

add_output(name, **kwargs)

Add an output to the component.

This overrides the base class method to update the kwargs to use dynamic shaping by default.

Parameters:

namestr

The name of the output.

**kwargsdict

The kwargs to pass to the base class method.

compute(inputs, outputs, discrete_inputs=None, discrete_outputs=None)

Compute outputs given inputs. The model is assumed to be in an unscaled state.

An inherited component may choose to either override this function or to define a compute_primal function.

Parameters:

inputsVector

Unscaled, dimensional input variables read via inputs[key].

outputsVector

Unscaled, dimensional output variables read via outputs[key].

discrete_inputsdict-like or None

If not None, dict-like object containing discrete input values.

discrete_outputsdict-like or None

If not None, dict-like object containing discrete output values.

compute_sparsity(direction=None, num_iters=1, perturb_size=1e-09)

Get the sparsity of the Jacobian.

Parameters:

directionstr

The direction to compute the sparsity for.

num_itersint

The number of times to run the perturbation iteration.

perturb_sizefloat

The size of the perturbation to use.

Returns:

coo_matrix

The sparsity of the Jacobian.

declare_coloring(**kwargs)

Declare coloring for this component.

The ‘method’ argument is set to ‘jax’ and passed to the base class.

Parameters:

**kwargsdict

Additional arguments to be passed to the base class.