Repository: OpenMDAO/OpenMDAO

Issue:
Jax components fail if argument names are not valid python argument names.
### Description

When using a JaxExplicitComponent or JaxImplicitComponent, an error is raised when an input or output is not a valid python name.

```
             if args != compargs:
                 raise RuntimeError(f"{self.msginfo}: compute_primal method args {args} don't match "
                                    f"the expected args {compargs}.")
```

### Example

A new argument, arg_name, should be added to `add_input` and `add_output` that allows the user to specify the name to be used within `compute_primal`:

```python
comp.add_input('illegal:python:name', val=1.0, arg_name='foo')
```


### OpenMDAO Version

3.36.1-dev

### Relevant environment information

_No response_

Hints from issue comments:

Solution patch:
diff --git a/openmdao/components/explicit_func_comp.py b/openmdao/components/explicit_func_comp.py
index 119cfdabb..c4a1d8bde 100644
--- a/openmdao/components/explicit_func_comp.py
+++ b/openmdao/components/explicit_func_comp.py
@@ -119,7 +119,7 @@ class ExplicitFuncComp(ExplicitComponent):
                 self._dev_arrays_to_np_arrays(kwargs)
             self.add_output(name, **kwargs)
 
-    def _setup_jax(self, from_group=False):
+    def _setup_jax(self):
         # TODO: this is here to prevent the ExplicitComponent base class from trying to do its
         # own jax setup if derivs_method is 'jax'. We should probably refactor this...
         pass
diff --git a/openmdao/components/implicit_func_comp.py b/openmdao/components/implicit_func_comp.py
index 17862de10..37fb11509 100644
--- a/openmdao/components/implicit_func_comp.py
+++ b/openmdao/components/implicit_func_comp.py
@@ -135,7 +135,7 @@ class ImplicitFuncComp(ImplicitComponent):
             kwargs = _copy_with_ignore(meta, omf._allowed_add_output_args, ignore=('resid',))
             self.add_output(name, **kwargs)
 
-    def _setup_jax(self, from_group=False):
+    def _setup_jax(self):
         # TODO: this is here to prevent the ImplicitComponent base class from trying to do its
         # own jax setup if derivs_method is 'jax'. We should probably refactor this...
         pass
diff --git a/openmdao/components/jax_explicit_comp.py b/openmdao/components/jax_explicit_comp.py
index e7bf58f1e..c6d2ae72d 100644
--- a/openmdao/components/jax_explicit_comp.py
+++ b/openmdao/components/jax_explicit_comp.py
@@ -3,9 +3,12 @@ An ExplicitComponent that uses JAX for derivatives.
 """
 
 import sys
+from types import MethodType
+
 from openmdao.core.explicitcomponent import ExplicitComponent
-from openmdao.utils.jax_utils import jax
 from openmdao.utils.om_warnings import issue_warning
+from openmdao.utils.jax_utils import jax, jit, compute_partials as _jax_compute_partials, \
+    compute_jacvec_product as _jax_compute_jacvec_product, ReturnChecker, _jax_register_pytree_class
 
 
 class JaxExplicitComponent(ExplicitComponent):
@@ -34,3 +37,56 @@ class JaxExplicitComponent(ExplicitComponent):
             issue_warning(f"{self.msginfo}: JAX is not available, so '{fallback_derivs_method}' "
                           "will be used for derivatives.")
             self.options['derivs_method'] = fallback_derivs_method
+
+    def _setup_jax(self):
+        """
+        Set up the jax interface for this component.
+        """
+        if self.compute_primal is None:
+            raise RuntimeError(f"{self.msginfo}: compute_primal is not defined for this component.")
+
+        if self.matrix_free is True:
+            self.compute_jacvec_product = MethodType(_jax_compute_jacvec_product, self)
+        else:
+            self.compute_partials = MethodType(_jax_compute_partials, self)
+            self._has_compute_partials = True
+
+        # determine if the compute_primal method returns a tuple
+        self._compute_primal_returns_tuple = ReturnChecker(self.compute_primal).returns_tuple()
+
+        if self.options['use_jit']:
+            static_argnums = []
+            idx = len(self._var_rel_names['input']) + 1
+            static_argnums.extend(range(idx, idx + len(self._discrete_inputs)))
+            self.compute_primal = MethodType(jit(self.compute_primal.__func__,
+                                                 static_argnums=static_argnums), self)
+
+        _jax_register_pytree_class(self.__class__)
+
+    def _get_jac_func(self):
+        """
+        Return the jacobian function for this component.
+
+        In forward mode, jax.jacfwd is used, and in reverse mode, jax.jacrev is used.  The direction
+        is chosen automatically based on the sizes of the inputs and outputs.
+
+        Returns
+        -------
+        function
+            The jacobian function.
+        """
+        # TODO: modify this to use relevance and possibly compile multiple jac functions depending
+        # on DV/response so that we don't compute any derivatives that are always zero.
+        if self._jac_func_ is None:
+            fjax = jax.jacfwd if self.best_partial_deriv_direction() == 'fwd' else jax.jacrev
+            nstatic = len(self._discrete_inputs)
+            wrt_idxs = list(range(1, len(self._var_abs2meta['input']) + 1))
+            self._jac_func_ = MethodType(fjax(self.compute_primal.__func__, argnums=wrt_idxs), self)
+
+            if self.options['use_jit']:
+                static_argnums = tuple(range(1 + len(wrt_idxs), 1 + len(wrt_idxs) + nstatic))
+                self._jac_func_ = MethodType(jit(self._jac_func_.__func__,
+                                                 static_argnums=static_argnums),
+                                             self)
+
+        return self._jac_func_
diff --git a/openmdao/components/jax_implicit_comp.py b/openmdao/components/jax_implicit_comp.py
index fc752af4a..15af01c1e 100644
--- a/openmdao/components/jax_implicit_comp.py
+++ b/openmdao/components/jax_implicit_comp.py
@@ -3,9 +3,12 @@ An ImplicitComponent that uses JAX for derivatives.
 """
 
 import sys
+from types import MethodType
+
 from openmdao.core.implicitcomponent import ImplicitComponent
-from openmdao.utils.jax_utils import jax
 from openmdao.utils.om_warnings import issue_warning
+from openmdao.utils.jax_utils import jax, jit, \
+    linearize as _jax_linearize, apply_linear as _jax_apply_linear, _jax_register_pytree_class
 
 
 class JaxImplicitComponent(ImplicitComponent):
@@ -34,3 +37,40 @@ class JaxImplicitComponent(ImplicitComponent):
             issue_warning(f"{self.msginfo}: JAX is not available, so "
                           f"'{fallback_derivs_method}' will be used for derivatives.")
             self.options['derivs_method'] = fallback_derivs_method
+
+    def _setup_jax(self):
+        if self.compute_primal is None:
+            raise RuntimeError(f"{self.msginfo}: compute_primal is not defined for this component.")
+
+        if self.matrix_free is True:
+            self.apply_linear = MethodType(_jax_apply_linear, self)
+        else:
+            self.linearize = MethodType(_jax_linearize, self)
+            self._has_linearize = True
+
+        if self.options['use_jit']:
+            static_argnums = []
+            idx = len(self._var_rel_names['input']) + len(self._var_rel_names['output']) + 1
+            static_argnums.extend(range(idx, idx + len(self._discrete_inputs)))
+            self.compute_primal = MethodType(jit(self.compute_primal.__func__,
+                                                 static_argnums=static_argnums), self)
+
+        _jax_register_pytree_class(self.__class__)
+
+    def _get_jac_func(self):
+        # TODO: modify this to use relevance and possibly compile multiple jac functions depending
+        # on DV/response so that we don't compute any derivatives that are always zero.
+        if self._jac_func_ is None:
+            fjax = jax.jacfwd if self.best_partial_deriv_direction() == 'fwd' else jax.jacrev
+            wrt_idxs = list(range(1, len(self._var_abs2meta['input']) +
+                                  len(self._var_abs2meta['output']) + 1))
+            primal_func = self.compute_primal.__func__
+            self._jac_func_ = MethodType(fjax(primal_func, argnums=wrt_idxs), self)
+
+            if self.options['use_jit']:
+                static_argnums = tuple(range(1 + len(wrt_idxs), 1 + len(wrt_idxs) +
+                                             len(self._discrete_inputs)))
+                self._jac_func_ = MethodType(jit(self._jac_func_.__func__,
+                                                 static_argnums=static_argnums),
+                                             self)
+        return self._jac_func_
diff --git a/openmdao/core/component.py b/openmdao/core/component.py
index f48b61ebe..d8432ac1b 100644
--- a/openmdao/core/component.py
+++ b/openmdao/core/component.py
@@ -2,6 +2,7 @@
 
 import sys
 import types
+import inspect
 from collections import defaultdict
 from collections.abc import Iterable
 from itertools import product
@@ -27,7 +28,8 @@ from openmdao.utils.indexer import Indexer, indexer
 import openmdao.utils.coloring as coloring_mod
 from openmdao.utils.om_warnings import issue_warning, MPIWarning, DistributedComponentWarning, \
     DerivativesWarning, warn_deprecation, OMInvalidCheckDerivativesOptionsWarning
-from openmdao.utils.code_utils import is_lambda, LambdaPickleWrapper, get_partials_deps
+from openmdao.utils.code_utils import is_lambda, LambdaPickleWrapper, get_partials_deps, \
+    get_return_names
 from openmdao.approximation_schemes.complex_step import ComplexStep
 from openmdao.approximation_schemes.finite_difference import FiniteDifference
 
@@ -82,7 +84,7 @@ class Component(System):
     _var_rel_names : {'input': [str, ...], 'output': [str, ...]}
         List of relative names of owned variables existing on current proc.
         This is only needed while adding inputs and outputs. During setup, these are used to
-        determine the list of absolute names.
+        determine the list of absolute names. This includes only continuous variables.
     _static_var_rel_names : dict
         Static version of above - stores names of variables added outside of setup.
     _declared_partials_patterns : dict
@@ -97,6 +99,8 @@ class Component(System):
     _compute_primals_out_shape : tuple or None
         Cached (shape, istuple) of the output from compute_primal function.  If istuple is True,
         then shape is a tuple of shapes, otherwise it is a single shape.
+    _valid_name_map : dict
+        Mapping of declared input/output names to valid Python names.
     """
 
     def __init__(self, **kwargs):
@@ -116,6 +120,7 @@ class Component(System):
         self._no_check_partials = False
         self._has_distrib_outputs = False
         self._compute_primals_out_shape = None
+        self._valid_name_map = {}
 
     def _tree_flatten(self):
         """
@@ -331,6 +336,10 @@ class Component(System):
             if msg:
                 raise RuntimeError(msg)
 
+        if self.compute_primal is not None:
+            self._check_compute_primal_args()
+            self._check_compute_primal_returns()
+
         self._serial_idxs = None
         self._inconsistent_keys = set()
 
@@ -541,7 +550,7 @@ class Component(System):
 
     def add_input(self, name, val=1.0, shape=None, units=None, desc='', tags=None,
                   shape_by_conn=False, copy_shape=None, compute_shape=None,
-                  require_connection=False, distributed=None):
+                  require_connection=False, distributed=None, primal_name=None):
         """
         Add an input variable to the component.
 
@@ -575,6 +584,9 @@ class Component(System):
         distributed : bool
             If True, this variable is a distributed variable, so it can have different sizes/values
             across MPI processes.
+        primal_name : str or None
+            Valid python name to represent the variable in compute_primal if 'name' is not a valid
+            python name.
 
         Returns
         -------
@@ -635,6 +647,9 @@ class Component(System):
         if compute_shape is not None and is_lambda(compute_shape):
             compute_shape = LambdaPickleWrapper(compute_shape)
 
+        if primal_name is not None:
+            self._valid_name_map[name] = primal_name
+
         metadata = {
             'val': val,
             'shape': shape,
@@ -672,7 +687,7 @@ class Component(System):
 
         return metadata
 
-    def add_discrete_input(self, name, val, desc='', tags=None):
+    def add_discrete_input(self, name, val, desc='', tags=None, primal_name=None):
         """
         Add a discrete input variable to the component.
 
@@ -687,6 +702,9 @@ class Component(System):
         tags : str or list of strs
             User defined tags that can be used to filter what gets listed when calling
             list_inputs and list_outputs.
+        primal_name : str or None
+            Valid python name to represent the variable in compute_primal if 'name' is not a valid
+            python name.
 
         Returns
         -------
@@ -701,6 +719,9 @@ class Component(System):
         if tags is not None and not isinstance(tags, (str, list)):
             raise TypeError('%s: The tags argument should be a str or list' % self.msginfo)
 
+        if primal_name is not None:
+            self._valid_name_map[name] = primal_name
+
         metadata = {}
 
         metadata.update({
@@ -730,7 +751,8 @@ class Component(System):
 
     def add_output(self, name, val=1.0, shape=None, units=None, res_units=None, desc='',
                    lower=None, upper=None, ref=1.0, ref0=0.0, res_ref=None, tags=None,
-                   shape_by_conn=False, copy_shape=None, compute_shape=None, distributed=None):
+                   shape_by_conn=False, copy_shape=None, compute_shape=None, distributed=None,
+                   primal_name=None):
         """
         Add an output variable to the component.
 
@@ -784,6 +806,9 @@ class Component(System):
         distributed : bool
             If True, this variable is a distributed variable, so it can have different sizes/values
             across MPI processes.
+        primal_name : str or None
+            Valid python name to represent the variable in compute_primal if 'name' is not a valid
+            python name.
 
         Returns
         -------
@@ -891,6 +916,9 @@ class Component(System):
         if compute_shape is not None and is_lambda(compute_shape):
             compute_shape = LambdaPickleWrapper(compute_shape)
 
+        if primal_name is not None:
+            self._valid_name_map[name] = primal_name
+
         metadata = {
             'val': val,
             'shape': shape,
@@ -933,7 +961,7 @@ class Component(System):
 
         return metadata
 
-    def add_discrete_output(self, name, val, desc='', tags=None):
+    def add_discrete_output(self, name, val, desc='', tags=None, primal_name=None):
         """
         Add an output variable to the component.
 
@@ -948,6 +976,9 @@ class Component(System):
         tags : str or list of strs or set of strs
             User defined tags that can be used to filter what gets listed when calling
             list_inputs and list_outputs.
+        primal_name : str or None
+            Valid python name to represent the variable in compute_primal if 'name' is not a valid
+            python name.
 
         Returns
         -------
@@ -961,6 +992,9 @@ class Component(System):
         if tags is not None and not isinstance(tags, (str, set, list)):
             raise TypeError('%s: The tags argument should be a str, set, or list' % self.msginfo)
 
+        if primal_name is not None:
+            self._valid_name_map[name] = primal_name
+
         metadata = {}
 
         metadata.update({
@@ -2551,6 +2585,48 @@ class Component(System):
         # from Problem.check_partials and passable to assert_check_partials.
         return {self.pathname: partials_data}, worst
 
+    def _check_compute_primal_args(self):
+        """
+        Check that the compute_primal method args are in the correct order.
+        """
+        args = list(inspect.signature(self.compute_primal).parameters)
+        if args and args[0] == 'self':
+            args = args[1:]
+        compargs = self._get_compute_primal_argnames()
+        if args != compargs:
+            raise RuntimeError(f"{self.msginfo}: compute_primal method args {args} don't match "
+                               f"the args {compargs} mapped from this component's inputs. To "
+                               "map inputs to the compute_primal method, set the name used in "
+                               "compute_primal to the 'primal_name' arg when calling "
+                               "add_input/add_discrete_input. This is only necessary if the "
+                               "declared component input name is not a valid Python name.")
+
+    def _check_compute_primal_returns(self):
+        """
+        Check that the compute_primal method returns a tuple.
+        """
+        retnames = get_return_names(self.compute_primal)
+        if self._valid_name_map:
+            expected_names = [self._valid_name_map.get(n, n) for n in self._var_rel_names['output']]
+            expected_names.extend([self._valid_name_map.get(n, n) for n in self._discrete_outputs])
+        else:
+            expected_names = list(self._var_rel_names['output'])
+            expected_names.extend(self._discrete_outputs)
+
+        if len(retnames) != len(expected_names):
+            raise RuntimeError(f"{self.msginfo}: compute_primal method returns {len(retnames)} "
+                               f"values but expected {len(expected_names)}.")
+
+        for i, (expname, rname) in enumerate(zip(expected_names, retnames)):
+            if rname is not None and expname != rname:
+                raise RuntimeError(f"{self.msginfo}: compute_primal method returns {rname} "
+                                   f"for return value {i} but the name of the output that was "
+                                   f"mapped for this component is {expname}. To map outputs to "
+                                   "the compute_primal method, set the name used in compute_primal "
+                                   "to the 'primal_name' arg when calling "
+                                   "add_output/add_discrete_output. This is only necessary if "
+                                   "the declared component output name is not a valid Python name.")
+
 
 class _DictValues(object):
     """
diff --git a/openmdao/core/explicitcomponent.py b/openmdao/core/explicitcomponent.py
index ca8a23f01..b434efe16 100644
--- a/openmdao/core/explicitcomponent.py
+++ b/openmdao/core/explicitcomponent.py
@@ -1,9 +1,7 @@
 """Define the ExplicitComponent class."""
-import inspect
 
 import numpy as np
-from types import MethodType
-
+from itertools import chain
 
 from openmdao.jacobians.dictionary_jacobian import DictionaryJacobian
 from openmdao.core.component import Component
@@ -12,9 +10,6 @@ from openmdao.utils.class_util import overrides_method
 from openmdao.recorders.recording_iteration_stack import Recording
 from openmdao.core.constants import INT_DTYPE, _UNDEFINED
 from openmdao.utils.general_utils import is_undefined
-from openmdao.utils.jax_utils import jax, jit, ExplicitCompJaxify, \
-    compute_partials as _jax_compute_partials, \
-    compute_jacvec_product as _jax_compute_jacvec_product, ReturnChecker, _jax_register_pytree_class
 
 
 _tuplist = (tuple, list)
@@ -173,7 +168,8 @@ class ExplicitComponent(Component):
 
     def add_output(self, name, val=1.0, shape=None, units=None, res_units=None, desc='',
                    lower=None, upper=None, ref=1.0, ref0=0.0, res_ref=None, tags=None,
-                   shape_by_conn=False, copy_shape=None, compute_shape=None, distributed=None):
+                   shape_by_conn=False, copy_shape=None, compute_shape=None, distributed=None,
+                   primal_name=None):
         """
         Add an output variable to the component.
 
@@ -229,6 +225,9 @@ class ExplicitComponent(Component):
         distributed : bool
             If True, this variable is a distributed variable, so it can have different sizes/values
             across MPI processes.
+        primal_name : str or None
+            Valid python name to represent the variable in compute_primal if 'name' is not a valid
+            python name.
 
         Returns
         -------
@@ -244,7 +243,7 @@ class ExplicitComponent(Component):
                                   ref=ref, ref0=ref0, res_ref=res_ref,
                                   tags=tags, shape_by_conn=shape_by_conn,
                                   copy_shape=copy_shape, compute_shape=compute_shape,
-                                  distributed=distributed)
+                                  distributed=distributed, primal_name=primal_name)
 
     def _approx_subjac_keys_iter(self):
         is_output = self._outputs._contains_abs
@@ -620,85 +619,8 @@ class ExplicitComponent(Component):
         list
             List of argnames expected by the compute_primal method.
         """
-        argnames = []
-        argnames.extend(self._var_rel_names['input'])
-        if self._discrete_inputs:
-            argnames.extend(self._discrete_inputs)
-        return argnames
-
-    def _setup_jax(self, from_group=False):
-        """
-        Set up the jax interface for this component.
-
-        Parameters
-        ----------
-        from_group : bool
-            If True, this is being called from a Group setup.
-        """
-        if self.matrix_free is True:
-            self.compute_jacvec_product = MethodType(_jax_compute_jacvec_product, self)
+        if self._valid_name_map:
+            return [self._valid_name_map.get(n, n) for n in chain(self._var_rel_names['input'],
+                                                                  self._discrete_inputs)]
         else:
-            self.compute_partials = MethodType(_jax_compute_partials, self)
-            self._has_compute_partials = True
-
-        if self.compute_primal is None:
-            # convert the compute method to a compute_primal method
-            jaxifier = ExplicitCompJaxify(self, verbose=True)
-
-            if jaxifier.get_self_statics:
-                self.get_self_statics = MethodType(jaxifier.get_self_statics, self)
-            # replace existing compute method with base class method, so that compute_primal
-            # will be called.
-            self.compute = MethodType(ExplicitComponent.compute, self)
-
-            self.compute_primal = MethodType(jaxifier.compute_primal, self)
-            self._compute_primal_returns_tuple = True
-        else:
-            # check that compute_primal args are in the correct order
-            args = list(inspect.signature(self.compute_primal).parameters)
-            if args and args[0] == 'self':
-                args = args[1:]
-            compargs = self._get_compute_primal_argnames()
-            if args != compargs:
-                raise RuntimeError(f"{self.msginfo}: compute_primal method args {args} don't match "
-                                   f"the expected args {compargs}.")
-
-            # determine if the compute_primal method returns a tuple
-            self._compute_primal_returns_tuple = ReturnChecker(self.compute_primal).returns_tuple()
-
-        if not from_group and self.options['use_jit']:
-            static_argnums = []
-            idx = len(self._var_rel_names['input']) + 1
-            static_argnums.extend(range(idx, idx + len(self._discrete_inputs)))
-            self.compute_primal = MethodType(jit(self.compute_primal.__func__,
-                                                 static_argnums=static_argnums), self)
-
-        _jax_register_pytree_class(self.__class__)
-
-    def _get_jac_func(self):
-        """
-        Return the jacobian function for this component.
-
-        In forward mode, jax.jacfwd is used, and in reverse mode, jax.jacrev is used.  The direction
-        is chosen automatically based on the sizes of the inputs and outputs.
-
-        Returns
-        -------
-        function
-            The jacobian function.
-        """
-        # TODO: modify this to use relevance and possibly compile multiple jac functions depending
-        # on DV/response so that we don't compute any derivatives that are always zero.
-        if self._jac_func_ is None:
-            fjax = jax.jacfwd if self.best_partial_deriv_direction() == 'fwd' else jax.jacrev
-            nstatic = len(self._discrete_inputs)
-            wrt_idxs = list(range(1, len(self._var_abs2meta['input']) + 1))
-            self._jac_func_ = MethodType(fjax(self.compute_primal.__func__, argnums=wrt_idxs), self)
-
-            if self.options['use_jit']:
-                static_argnums = tuple(range(1 + len(wrt_idxs), 1 + len(wrt_idxs) + nstatic))
-                self._jac_func_ = MethodType(jit(self._jac_func_.__func__,
-                                                 static_argnums=static_argnums),
-                                             self)
-
-        return self._jac_func_
+            return list(chain(self._var_rel_names['input'], self._discrete_inputs))
diff --git a/openmdao/core/group.py b/openmdao/core/group.py
index aff11a5a7..2728c1d89 100644
--- a/openmdao/core/group.py
+++ b/openmdao/core/group.py
@@ -2374,7 +2374,7 @@ class Group(System):
         vnames = self._var_abs2prom[io] if local else self._var_allprocs_abs2prom[io]
         return set(vnames).difference(self._conn_global_abs_in2out)
 
-    def _setup_jax(self, from_group=False):
+    def _setup_jax(self):
         if jax is None:
             return
 
@@ -2382,7 +2382,7 @@ class Group(System):
         # jax anywhere below where it's active, then return.
         for subsys in self._subsystems_myproc:
             if isinstance(subsys, Group) or subsys.options['derivs_method'] == 'jax':
-                subsys._setup_jax(from_group)
+                subsys._setup_jax()
 
     def _setup_dynamic_shapes(self):
         """
diff --git a/openmdao/core/implicitcomponent.py b/openmdao/core/implicitcomponent.py
index 5c5845d8c..6dc691d33 100644
--- a/openmdao/core/implicitcomponent.py
+++ b/openmdao/core/implicitcomponent.py
@@ -1,9 +1,8 @@
 """Define the ImplicitComponent class."""
 
-import inspect
 from scipy.sparse import coo_matrix
+from itertools import chain
 import numpy as np
-from types import MethodType
 
 from openmdao.core.component import Component, _allowed_types
 from openmdao.core.constants import _UNDEFINED, _SetupStatus
@@ -16,8 +15,6 @@ from openmdao.utils.general_utils import format_as_float_or_array, _subjac_meta2
 from openmdao.utils.units import simplify_unit
 from openmdao.utils.rangemapper import RangeMapper
 from openmdao.utils.om_warnings import issue_warning
-from openmdao.utils.jax_utils import jax, jit, ImplicitCompJaxify, \
-    linearize as _jax_linearize, apply_linear as _jax_apply_linear, _jax_register_pytree_class
 
 
 _tuplist = (tuple, list)
@@ -939,69 +936,15 @@ class ImplicitComponent(Component):
             yield from discrete_inputs.values()
 
     def _get_compute_primal_argnames(self):
-        argnames = []
-        argnames.extend(name for name in self._var_rel_names['input']
-                        if name not in self._discrete_inputs)
-        argnames.extend(name for name in self._var_rel_names['output']
-                        if name not in self._discrete_outputs)
-        if self._discrete_inputs:
-            argnames.extend(self._discrete_inputs)
-        return argnames
-
-    def _setup_jax(self, from_group=False):
-        if self.matrix_free is True:
-            self.apply_linear = MethodType(_jax_apply_linear, self)
+        if self._valid_name_map:
+            argnames = [self._valid_name_map.get(n, n) for n in self._var_rel_names['input']]
+            argnames.extend(self._valid_name_map.get(n, n) for n in self._var_rel_names['output'])
+            if self._discrete_inputs:
+                argnames.extend(self._valid_name_map.get(n, n) for n in self._discrete_inputs)
+            return argnames
         else:
-            self.linearize = MethodType(_jax_linearize, self)
-            self._has_linearize = True
-
-        if self.compute_primal is None:
-            jaxifier = ImplicitCompJaxify(self, verbose=True)
-
-            if jaxifier.get_self_statics:
-                self.get_self_statics = MethodType(jaxifier.get_self_statics, self)
-
-            # replace existing apply_nonlinear method with base class method, so that compute_primal
-            # will be called.
-            self.apply_nonlinear = MethodType(ImplicitComponent.apply_nonlinear, self)
-
-            self.compute_primal = MethodType(jaxifier.compute_primal, self)
-        else:
-            # check that compute_primal args are in the correct order
-            args = list(inspect.signature(self.compute_primal).parameters)
-            if args and args[0] == 'self':
-                args = args[1:]
-            compargs = self._get_compute_primal_argnames()
-            if args != compargs:
-                raise RuntimeError(f"{self.msginfo}: compute_primal method args {args} don't match "
-                                   f"the expected args {compargs}.")
-
-        if not from_group and self.options['use_jit']:
-            static_argnums = []
-            idx = len(self._var_rel_names['input']) + len(self._var_rel_names['output']) + 1
-            static_argnums.extend(range(idx, idx + len(self._discrete_inputs)))
-            self.compute_primal = MethodType(jit(self.compute_primal.__func__,
-                                                 static_argnums=static_argnums), self)
-
-        _jax_register_pytree_class(self.__class__)
-
-    def _get_jac_func(self):
-        # TODO: modify this to use relevance and possibly compile multiple jac functions depending
-        # on DV/response so that we don't compute any derivatives that are always zero.
-        if self._jac_func_ is None:
-            fjax = jax.jacfwd if self.best_partial_deriv_direction() == 'fwd' else jax.jacrev
-            wrt_idxs = list(range(1, len(self._var_abs2meta['input']) +
-                                  len(self._var_abs2meta['output']) + 1))
-            primal_func = self.compute_primal.__func__
-            self._jac_func_ = MethodType(fjax(primal_func, argnums=wrt_idxs), self)
-
-            if self.options['use_jit']:
-                static_argnums = tuple(range(1 + len(wrt_idxs), 1 + len(wrt_idxs) +
-                                             len(self._discrete_inputs)))
-                self._jac_func_ = MethodType(jit(self._jac_func_.__func__,
-                                                 static_argnums=static_argnums),
-                                             self)
-        return self._jac_func_
+            return list(chain(self._var_rel_names['input'], self._var_rel_names['output'],
+                              self._discrete_inputs))
 
 
 def meta2range_iter(meta_dict, names=None, shp_name='shape'):
diff --git a/openmdao/core/indepvarcomp.py b/openmdao/core/indepvarcomp.py
index 3d21cf28e..3ec769831 100644
--- a/openmdao/core/indepvarcomp.py
+++ b/openmdao/core/indepvarcomp.py
@@ -227,17 +227,6 @@ class IndepVarComp(ExplicitComponent):
         # define this for IndepVarComp to avoid overhead of ExplicitComponent._linearize.
         pass
 
-    def compute_primal(self):
-        """
-        Compute the outputs.
-
-        Returns
-        -------
-        tuple
-            A tuple containing the output values.
-        """
-        return tuple(self._outputs.values())
-
     def _apply_nonlinear(self):
         """
         Compute residuals. The model is assumed to be in a scaled state.
diff --git a/openmdao/core/system.py b/openmdao/core/system.py
index 2a08b5ef2..4f827fdb4 100644
--- a/openmdao/core/system.py
+++ b/openmdao/core/system.py
@@ -6758,7 +6758,7 @@ class System(object, metaclass=SystemMetaclass):
         """
         return ()
 
-    def _setup_jax(self, from_group=False):
+    def _setup_jax(self):
         pass
 
     def _sys_tree_visitor(self, func, predicate=None, recurse=True, include_self=True,
diff --git a/openmdao/utils/code_utils.py b/openmdao/utils/code_utils.py
index a3af22905..c99189865 100644
--- a/openmdao/utils/code_utils.py
+++ b/openmdao/utils/code_utils.py
@@ -427,6 +427,7 @@ def get_return_names(func):
         def __init__(self, func):
             super().__init__()
             self._ret_infos = []
+            self.fstack = []
             self.visit(ast.parse(textwrap.dedent(inspect.getsource(func)), mode='exec'))
 
         def get_return_names(self):
@@ -462,6 +463,22 @@ def get_return_names(func):
             else:
                 self._ret_infos[-1].append(_get_return_name(node.value))
 
+        def visit_FunctionDef(self, node):
+            """
+            Visit a FunctionDef node.
+
+            Parameters
+            ----------
+            node : ASTnode
+                The function definition node being visited.
+            """
+            if self.fstack:
+                return  # skip nested functions
+            self.fstack.append(node)
+            for stmt in node.body:
+                self.visit(stmt)
+            self.fstack.pop()
+
     return _FuncRetNameCollector(func).get_return_names()
 
 
diff --git a/openmdao/utils/jax_utils.py b/openmdao/utils/jax_utils.py
index d194d8efb..eff9c9ea1 100644
--- a/openmdao/utils/jax_utils.py
+++ b/openmdao/utils/jax_utils.py
@@ -93,9 +93,7 @@ def register_jax_component(comp_class):
                                   f'\nCannot register class {comp_class} as a jax jit-compatible '
                                   f'component.')
 
-    jax.tree_util.register_pytree_node(comp_class,
-                                       comp_class._tree_flatten,
-                                       comp_class._tree_unflatten)
+    _jax_register_pytree_class(comp_class)
     return comp_class
 
 
@@ -333,7 +331,8 @@ class CompJaxifyBase(ast.NodeTransformer):
         # arg["name"] with name.
 
         # NOTE: this will only work if the subscript is a string constant. If the subscript is a
-        # variable or some other expression, then we don't modify it.
+        # variable or some other expression, then we don't modify it and the conversion will
+        # likely fail.
         if (isinstance(node.value, ast.Name) and node.value.id in self._orig_args and
                 isinstance(node.slice, ast.Constant) and isinstance(node.slice.value, str)):
             return ast.copy_location(ast.Name(id=_fixname(node.slice.value), ctx=node.ctx), node)
@@ -779,16 +778,12 @@ else:
         """
         Register a class with jax so that it can be used with jax.jit.
 
+        This can be called after instantiating the class if necessary.
+
         Parameters
         ----------
         cls : class
             The class to be registered.
-        name : str
-            The name of the class.
-        bases : tuple
-            The base classes of the class.
-        attrs : dict
-            The attributes of the class.
         """
         global _registered_classes
         if cls not in _registered_classes:


Test patch:
diff --git a/openmdao/core/tests/test_expl_comp.py b/openmdao/core/tests/test_expl_comp.py
index 779c10dad..aef6fa3eb 100644
--- a/openmdao/core/tests/test_expl_comp.py
+++ b/openmdao/core/tests/test_expl_comp.py
@@ -2,7 +2,7 @@
 
 from io import StringIO
 import unittest
-
+import itertools
 import numpy as np
 
 import openmdao.api as om
@@ -16,6 +16,12 @@ from openmdao.test_suite.components.rectangle import RectangleGroup, \
 from openmdao.utils.assert_utils import assert_near_equal
 from openmdao.utils.general_utils import printoptions, remove_whitespace
 from openmdao.utils.mpi import MPI
+from openmdao.utils.testing_utils import parameterized_name
+
+try:
+    from parameterized import parameterized
+except ImportError:
+    from openmdao.utils.assert_utils import SkipParameterized as parameterized
 
 
 class ExplCompTestCase(unittest.TestCase):
@@ -1154,5 +1160,192 @@ class TestMPIExplComp(unittest.TestCase):
                 if line and not line.startswith('-'):
                     self.assertEqual(remove_whitespace(text[i]), remove_whitespace(line))
 
+class DotProductMultDiscretePrimalBadNamesBase(om.ExplicitComponent):
+
+    def compute_primal(self, x, y, disc_in):
+        disc_out = -disc_in
+        if disc_in >= 0:
+            z = np.dot(x, y)
+            zz = y * 3.0
+        else:
+            z = -np.dot(x, y)
+            zz = y * 2.5
+
+        self._discrete_outputs.set_vals((disc_out,))
+        return (z, zz, disc_out)
+
+
+class DotProductMultDiscretePrimalBadNamesCorrect(DotProductMultDiscretePrimalBadNamesBase):
+    def setup(self):
+        super().setup()
+        self.add_input('my:x', primal_name='x', shape_by_conn=True)
+        self.add_input('my:y', primal_name='y', shape_by_conn=True)
+        self.add_discrete_input('my:disc_in', primal_name='disc_in', val=2)
+        self.add_output('my:z', primal_name='z', compute_shape=lambda shapes: (shapes['my:x'][0], shapes['my:y'][1]))
+        self.add_output('my:zz', primal_name='zz', copy_shape='my:y')
+        self.add_discrete_output('my:disc_out', primal_name='disc_out', val=3)
+        self.declare_partials(of=['my:z', 'my:zz'], wrt=['my:x', 'my:y'])
+
+
+class DotProductMultDiscretePrimalBadInsWrongOrder(DotProductMultDiscretePrimalBadNamesBase):
+    def setup(self):
+        super().setup()
+        self.add_input('my:x', primal_name='y', shape_by_conn=True)
+        self.add_input('my:y', primal_name='x', shape_by_conn=True)
+        self.add_discrete_input('my:disc_in', primal_name='disc_in', val=2)
+        self.add_output('z', compute_shape=lambda shapes: (shapes['my:x'][0], shapes['my:y'][1]))
+        self.add_output('zz', copy_shape='my:y')
+        self.add_discrete_output('disc_out', val=3)
+        self.declare_partials(of=['z', 'zz'], wrt=['my:x', 'my:y'])
+
+
+class DotProductMultDiscretePrimalBadInsNoMapping(DotProductMultDiscretePrimalBadNamesBase):
+    def setup(self):
+        super().setup()
+        self.add_input('my:x', shape_by_conn=True)
+        self.add_input('my:y', shape_by_conn=True)
+        self.add_discrete_input('my:disc_in', val=2)
+        self.add_output('z', compute_shape=lambda shapes: (shapes['my:x'][0], shapes['my:y'][1]))
+        self.add_output('zz', copy_shape='my:y')
+        self.add_discrete_output('disc_out', val=3)
+        self.declare_partials(of=['z', 'zz'], wrt=['my:x', 'my:y'])
+
+
+class DotProductMultDiscretePrimalBadOutsWrongOrder(DotProductMultDiscretePrimalBadNamesBase):
+    def setup(self):
+        super().setup()
+        self.add_input('x', shape_by_conn=True)
+        self.add_input('y', shape_by_conn=True)
+        self.add_discrete_input('disc_in', primal_name='disc_in', val=2)
+        self.add_output('my:z', primal_name='zz', compute_shape=lambda shapes: (shapes['x'][0], shapes['y'][1]))
+        self.add_output('my:zz', primal_name='z', copy_shape='my:y')
+        self.add_discrete_output('my:disc_out', primal_name='disc_out', val=3)
+        self.declare_partials(of=['my:z', 'my:zz'], wrt=['x', 'y'])
+
+
+class DotProductMultDiscretePrimalBadOutsNoMapping(DotProductMultDiscretePrimalBadNamesBase):
+    def setup(self):
+        super().setup()
+        self.add_input('x', shape_by_conn=True)
+        self.add_input('y', shape_by_conn=True)
+        self.add_discrete_input('disc_in', primal_name='disc_in', val=2)
+        self.add_output('my:z', compute_shape=lambda shapes: (shapes['x'][0], shapes['y'][1]))
+        self.add_output('my:zz', copy_shape='y')
+        self.add_discrete_output('my:disc_out', val=3)
+        self.declare_partials(of=['my:z', 'my:zz'], wrt=['x', 'y'])
+
+
+x_shape = (2, 3)
+y_shape = (3, 4)
+
+
+class TestMappedNames(unittest.TestCase):
+    @parameterized.expand(itertools.product(['fwd', 'rev'], [True, False]), name_func=parameterized_name)
+    def test_explicit_comp_with_discrete_primal_mapped_names(self, mode, matrix_free):
+        p = om.Problem()
+        ivc = p.model.add_subsystem('ivc', om.IndepVarComp('x', val=np.ones(x_shape)))
+        ivc.add_output('y', val=np.ones(y_shape))
+        ivc.add_discrete_output('disc_out', val=3)
+        comp = p.model.add_subsystem('comp', DotProductMultDiscretePrimalBadNamesCorrect())
+        comp.matrix_free = matrix_free
+
+        p.model.connect('ivc.x', 'comp.my:x')
+        p.model.connect('ivc.y', 'comp.my:y')
+        p.model.connect('ivc.disc_out', 'comp.my:disc_in')
+
+        p.setup(mode=mode)
+
+        x = np.arange(1,np.prod(x_shape)+1).reshape(x_shape) * 2.0
+        y = np.arange(1,np.prod(y_shape)+1).reshape(y_shape)* 3.0
+        p.set_val('ivc.x', x)
+        p.set_val('ivc.y', y)
+        p.final_setup()
+        p.run_model()
+
+        assert_near_equal(p.get_val('comp.my:z'), np.dot(x, y))
+        assert_near_equal(p.get_val('comp.my:zz'), y * 3.0)
+
+    def test_explicit_comp_with_discrete_primal_bad_inputs_wrong_order(self):
+        p = om.Problem()
+        ivc = p.model.add_subsystem('ivc', om.IndepVarComp('x', val=np.ones(x_shape)))
+        ivc.add_output('y', val=np.ones(y_shape))
+        ivc.add_discrete_output('disc_out', val=3)
+        p.model.add_subsystem('comp', DotProductMultDiscretePrimalBadInsWrongOrder())
+
+        p.model.connect('ivc.x', 'comp.my:x')
+        p.model.connect('ivc.y', 'comp.my:y')
+        p.model.connect('ivc.disc_out', 'comp.my:disc_in')
+
+        with self.assertRaises(RuntimeError) as ctx:
+            p.setup()
+
+        self.assertEqual(str(ctx.exception),
+                         "'comp' <class DotProductMultDiscretePrimalBadInsWrongOrder>: compute_primal method args ['x', 'y', 'disc_in'] "
+                         "don't match the args ['y', 'x', 'disc_in'] mapped from this component's inputs. To map inputs to the compute_primal "
+                         "method, set the name used in compute_primal to the 'primal_name' arg when calling add_input/add_discrete_input. "
+                         "This is only necessary if the declared component input name is not a valid Python name.")
+
+    def test_explicit_comp_with_discrete_primal_bad_outputs_wrong_order(self):
+        p = om.Problem()
+        ivc = p.model.add_subsystem('ivc', om.IndepVarComp('x', val=np.ones(x_shape)))
+        ivc.add_output('y', val=np.ones(y_shape))
+        ivc.add_discrete_output('disc_out', val=3)
+        p.model.add_subsystem('comp', DotProductMultDiscretePrimalBadOutsWrongOrder())
+
+        p.model.connect('ivc.x', 'comp.x')
+        p.model.connect('ivc.y', 'comp.y')
+        p.model.connect('ivc.disc_out', 'comp.disc_in')
+
+        with self.assertRaises(RuntimeError) as ctx:
+            p.setup()
+
+        self.assertEqual(str(ctx.exception),
+                         "'comp' <class DotProductMultDiscretePrimalBadOutsWrongOrder>: compute_primal method returns z for return value 0 "
+                         "but the name of the output that was mapped for this component is zz. To map outputs to the compute_primal method, "
+                         "set the name used in compute_primal to the 'primal_name' arg when calling add_output/add_discrete_output. "
+                         "This is only necessary if the declared component output name is not a valid Python name.")
+
+
+    def test_explicit_comp_with_discrete_primal_bad_inputs_no_mapping(self):
+        p = om.Problem()
+        ivc = p.model.add_subsystem('ivc', om.IndepVarComp('x', val=np.ones(x_shape)))
+        ivc.add_output('y', val=np.ones(y_shape))
+        ivc.add_discrete_output('disc_out', val=3)
+        p.model.add_subsystem('comp', DotProductMultDiscretePrimalBadInsNoMapping())
+
+        p.model.connect('ivc.x', 'comp.my:x')
+        p.model.connect('ivc.y', 'comp.my:y')
+        p.model.connect('ivc.disc_out', 'comp.my:disc_in')
+
+        with self.assertRaises(RuntimeError) as ctx:
+            p.setup()
+
+        self.assertEqual(str(ctx.exception),
+                         "'comp' <class DotProductMultDiscretePrimalBadInsNoMapping>: compute_primal method args ['x', 'y', 'disc_in'] "
+                         "don't match the args ['my:x', 'my:y', 'my:disc_in'] mapped from this component's inputs. To map inputs to the "
+                         "compute_primal method, set the name used in compute_primal to the 'primal_name' arg when calling add_input/add_discrete_input. "
+                         "This is only necessary if the declared component input name is not a valid Python name.")
+
+    def test_explicit_comp_with_discrete_primal_bad_outputs_no_mapping(self):
+        p = om.Problem()
+        ivc = p.model.add_subsystem('ivc', om.IndepVarComp('x', val=np.ones(x_shape)))
+        ivc.add_output('y', val=np.ones(y_shape))
+        ivc.add_discrete_output('disc_out', val=3)
+        p.model.add_subsystem('comp', DotProductMultDiscretePrimalBadOutsNoMapping())
+
+        p.model.connect('ivc.x', 'comp.x')
+        p.model.connect('ivc.y', 'comp.y')
+        p.model.connect('ivc.disc_out', 'comp.disc_in')
+
+        with self.assertRaises(RuntimeError) as ctx:
+            p.setup()
+
+        self.assertEqual(str(ctx.exception),
+                         "'comp' <class DotProductMultDiscretePrimalBadOutsNoMapping>: compute_primal method returns z for return value 0 "
+                         "but the name of the output that was mapped for this component is my:z. To map outputs to the compute_primal method, "
+                         "set the name used in compute_primal to the 'primal_name' arg when calling add_output/add_discrete_output. "
+                         "This is only necessary if the declared component output name is not a valid Python name.")
+
+
 if __name__ == '__main__':
     unittest.main()
diff --git a/openmdao/core/tests/test_impl_comp.py b/openmdao/core/tests/test_impl_comp.py
index 6b5a7e1c5..089960d52 100644
--- a/openmdao/core/tests/test_impl_comp.py
+++ b/openmdao/core/tests/test_impl_comp.py
@@ -1563,5 +1563,116 @@ class CacheLinSolutionTestCase(unittest.TestCase):
             p.driver._compute_totals(of=['C1.y'], wrt=['indeps.x'])
 
 
+class LinearSystemCompPrimal(om.ImplicitComponent):
+    def __init__(self, input_prefix, output_prefix, domap, **kwargs):
+        super().__init__(**kwargs)
+        self.Aname = input_prefix + 'A'
+        self.bname = input_prefix + 'b'
+        self.xname = output_prefix + 'x'
+        self.domap = domap
+
+    def initialize(self):
+        self.options.declare('size', default=1, types=int)
+
+    def setup(self):
+        size = self.options['size']
+
+        shape = (size, )
+
+        if ':' in self.Aname and self.domap:
+            self.add_input(self.Aname, primal_name=self.Aname.rpartition(':')[-1], val=np.eye(size))
+        else:
+            self.add_input(self.Aname, val=np.eye(size))
+        if ':' in self.bname and self.domap:
+            self.add_input(self.bname, primal_name=self.bname.rpartition(':')[-1], val=np.ones(shape))
+        else:
+            self.add_input(self.bname, val=np.ones(shape))
+        if ':' in self.xname and self.domap:
+            self.add_output(self.xname, primal_name=self.xname.rpartition(':')[-1], shape=shape)
+        else:
+            self.add_output(self.xname, shape=shape)
+
+    def setup_partials(self):
+        size = self.options['size']
+        mat_size = size * size
+        full_size = size
+
+        row_col = np.arange(full_size, dtype="int")
+
+        self.declare_partials(self.xname, self.bname, val=np.full(full_size, -1.0), rows=row_col, cols=row_col)
+
+        rows = np.repeat(np.arange(full_size), size)
+
+        cols = np.arange(mat_size)
+
+        self.declare_partials(self.xname, self.Aname, rows=rows, cols=cols)
+
+        cols = np.tile(np.arange(size), size)
+        cols += np.repeat(np.arange(1), mat_size) * size
+
+        self.declare_partials(of=self.xname, wrt=self.xname, rows=rows, cols=cols)
+
+        if self.matrix_free:
+            self.linear_solver = om.ScipyKrylov()
+        else:
+            self.linear_solver = om.DirectSolver()
+        self.nonlinear_solver = om.NewtonSolver(solve_subsystems=False)
+
+    def compute_primal(self, A, b, x):
+        return A.dot(x) - b
+
+
+class TestMappedNames(unittest.TestCase):
+    def test_implicit_comp_primal_mapped_names(self):
+        p = om.Problem()
+        p.model.add_subsystem('comp', LinearSystemCompPrimal('my:', 'my:', domap=True, size=3,
+                                                             derivs_method='cs'))
+
+        p.setup()
+
+        x = np.array([1, 2, -3])
+        A = np.array([[5.0, -3.0, 2.0], [1.0, 7.0, -4.0], [1.0, 0.0, 8.0]])
+        b = A.dot(x)
+
+        p.set_val('comp.my:A', A)
+        p.set_val('comp.my:b', b)
+        p.set_val('comp.my:x', x)
+        p.final_setup()
+        p.run_model()
+
+        assert_near_equal(p['comp.my:x'], x, .0001)
+
+    def test_implicit_comp_primal_bad_inputs_no_mapping(self):
+        p = om.Problem()
+        p.model.add_subsystem('comp', LinearSystemCompPrimal('my:', '', domap=False))
+
+        with self.assertRaises(RuntimeError) as ctx:
+            p.setup()
+
+        self.assertEqual(str(ctx.exception),
+                         "'comp' <class LinearSystemCompPrimal>: compute_primal method args ['A', 'b', 'x'] "
+                         "don't match the args ['my:A', 'my:b', 'x'] mapped from this component's inputs. To "
+                         "map inputs to the compute_primal method, set the name used in compute_primal to the "
+                         "'primal_name' arg when calling add_input/add_discrete_input. This is only necessary "
+                         "if the declared component input name is not a valid Python name.")
+
+
+
+    def test_implicit_comp_primal_bad_outputs_no_mapping(self):
+        p = om.Problem()
+        p.model.add_subsystem('comp', LinearSystemCompPrimal('', 'my:', domap=False))
+
+        with self.assertRaises(RuntimeError) as ctx:
+            p.setup()
+
+        self.assertEqual(str(ctx.exception),
+                         "'comp' <class LinearSystemCompPrimal>: compute_primal method args ['A', 'b', 'x'] "
+                         "don't match the args ['A', 'b', 'my:x'] mapped from this component's inputs. To map "
+                         "inputs to the compute_primal method, set the name used in compute_primal to the 'primal_name' "
+                         "arg when calling add_input/add_discrete_input. This is only necessary if the declared "
+                         "component input name is not a valid Python name.")
+
+
+
 if __name__ == '__main__':
     unittest.main()
diff --git a/openmdao/jax/tests/test_jax_explicit.py b/openmdao/jax/tests/test_jax_explicit.py
index 545b55170..6e2e67b90 100644
--- a/openmdao/jax/tests/test_jax_explicit.py
+++ b/openmdao/jax/tests/test_jax_explicit.py
@@ -85,51 +85,7 @@ class DotProdMultPrimalOption(om.JaxExplicitComponent):
         return z, zz
 
 
-
-class DotProdMultOption(om.ExplicitComponent):
-    def initialize(self):
-        self.options.declare('mult', default=2.5, desc='multiplier', types=(float,))
-
-    def setup(self):
-        self.add_input('x', shape_by_conn=True)
-        self.add_input('y', shape_by_conn=True)
-        self.add_output('z', compute_shape=lambda shapes: (shapes['x'][0], shapes['y'][1]))
-        self.add_output('zz', copy_shape='y')
-
-        self.declare_partials(of=['z', 'zz'], wrt=['x', 'y'])
-
-    def compute(self, inputs, outputs):
-        outputs['z'] = np.dot(inputs['x'], inputs['y'])
-        outputs['zz'] = inputs['y'] * self.options['mult']
-
-
-
-class DotProductMultDiscrete(om.ExplicitComponent):
-    def setup(self):
-        self.add_input('x', shape_by_conn=True)
-        self.add_input('y', shape_by_conn=True)
-        self.add_discrete_input('disc_in', val=2)
-        self.add_output('z', compute_shape=lambda shapes: (shapes['x'][0], shapes['y'][1]))
-        self.add_output('zz', copy_shape='y')
-        self.add_discrete_output('disc_out', val=3)
-
-        self.declare_partials(of=['z', 'zz'], wrt=['x', 'y'])
-
-    def compute(self, inputs, outputs, discrete_inputs, discrete_outputs):
-        discrete_outputs['disc_out'] = -discrete_inputs['disc_in']
-        if discrete_inputs['disc_in'] > 0:
-            outputs['z'] = np.dot(inputs['x'], inputs['y'])
-        else:
-            outputs['z'] = -np.dot(inputs['x'], inputs['y'])
-
-        if discrete_outputs['disc_out'] > 0:
-            outputs['zz'] = inputs['y'] * 2.5
-        else:
-            outputs['zz'] = inputs['y'] * 3.0
-
-
-
-class DotProductMultDiscretePrimal(om.ExplicitComponent):
+class DotProductMultDiscretePrimal(om.JaxExplicitComponent):
     def setup(self):
         self.add_input('x', shape_by_conn=True)
         self.add_input('y', shape_by_conn=True)
@@ -153,6 +109,7 @@ class DotProductMultDiscretePrimal(om.ExplicitComponent):
         self._discrete_outputs.set_vals((disc_out,))
         return (z, zz, disc_out)
 
+
 x_shape = (2, 3)
 y_shape = (3, 4)
 
@@ -160,7 +117,7 @@ y_shape = (3, 4)
 @unittest.skipIf(jax is None or sys.version_info < (3, 9), 'jax is not available or python < 3.9.')
 class TestJaxComp(unittest.TestCase):
 
-    @parameterized.expand(itertools.product(['fwd', 'rev'], [(True, 'jax'), (False, 'jax'), (False, 'fd')]), name_func=parameterized_name)
+    @parameterized.expand(itertools.product(['fwd', 'rev'], [(False, 'fd')]), name_func=parameterized_name)
     def test_jax_explicit_comp(self, mode, tup):
         matrix_free, derivs_method = tup
         xshape = x_shape
@@ -188,32 +145,6 @@ class TestJaxComp(unittest.TestCase):
         assert_check_totals(p.check_totals(of=['comp.z'], wrt=['comp.x', 'comp.y'], method='cs', show_only_incorrect=True))
         assert_check_partials(p.check_partials(method='cs', show_only_incorrect=True))
 
-    @parameterized.expand(itertools.product(['fwd', 'rev'], [True, False]), name_func=parameterized_name)
-    def test_jax_explicit_comp2(self, mode, matrix_free):
-        p = om.Problem()
-        ivc = p.model.add_subsystem('ivc', om.IndepVarComp('x', val=np.ones(x_shape)))
-        ivc.add_output('y', val=np.ones(y_shape))
-        comp = p.model.add_subsystem('comp', DotProdMult(derivs_method='jax'))
-        comp.matrix_free = matrix_free
-        p.model.connect('ivc.x', 'comp.x')
-        p.model.connect('ivc.y', 'comp.y')
-
-        p.setup(mode=mode)
-
-        x = np.arange(1,np.prod(x_shape)+1).reshape(x_shape) * 2.0
-        y = np.arange(1,np.prod(y_shape)+1).reshape(y_shape)* 3.0
-        p.set_val('ivc.x', x)
-        p.set_val('ivc.y', y)
-        p.final_setup()
-        p.run_model()
-
-        assert_near_equal(p.get_val('comp.z'), np.dot(x, y))
-        assert_near_equal(p.get_val('comp.zz'), y * 2.5)
-        assert_check_totals(p.check_totals(of=['comp.z','comp.zz'],
-                                           wrt=['comp.x', 'comp.y'], method='fd',
-                                           show_only_incorrect=True))
-        assert_check_partials(p.check_partials(show_only_incorrect=True))
-
     @parameterized.expand(itertools.product(['fwd', 'rev'], [True, False]), name_func=parameterized_name)
     def test_jax_explicit_comp2primal(self, mode, matrix_free):
         # this component defines its own compute_primal method
@@ -325,7 +256,7 @@ class TestJaxComp(unittest.TestCase):
         p = om.Problem()
         ivc = p.model.add_subsystem('ivc', om.IndepVarComp('x', val=np.ones(x_shape)))
         ivc.add_output('y', val=np.ones(y_shape))
-        comp = p.model.add_subsystem('comp', DotProdMultOption(mult=1.7, derivs_method='jax'))
+        comp = p.model.add_subsystem('comp', DotProdMultPrimalOption(mult=1.7))
         comp.matrix_free = matrix_free
 
         p.model.connect('ivc.x', 'comp.x')
@@ -341,7 +272,7 @@ class TestJaxComp(unittest.TestCase):
         p.run_model()
 
         assert_near_equal(p.get_val('comp.z'), np.dot(x, y))
-        assert_near_equal(p.get_val('comp.zz'), y * 1.7)
+        assert_near_equal(p.get_val('comp.zz'), y * 1.7 * 2.)
         assert_check_totals(p.check_totals(of=['comp.z','comp.zz'], wrt=['comp.x', 'comp.y'],
                                            method='fd', show_only_incorrect=True))
         assert_check_partials(p.check_partials(show_only_incorrect=True))
@@ -351,43 +282,7 @@ class TestJaxComp(unittest.TestCase):
         comp.options['mult'] = 1.9
         p.run_model()
         assert_near_equal(p.get_val('comp.z'), np.dot(x, y))
-        assert_near_equal(p.get_val('comp.zz'), y * 1.9)
-        assert_check_totals(p.check_totals(of=['comp.z','comp.zz'], wrt=['comp.x', 'comp.y'],
-                                           method='fd', show_only_incorrect=True))
-        assert_check_partials(p.check_partials(show_only_incorrect=True))
-
-    @parameterized.expand(itertools.product(['fwd', 'rev'], [True, False]), name_func=parameterized_name)
-    def test_jax_explicit_comp_with_discrete(self, mode, matrix_free):
-        p = om.Problem()
-        ivc = p.model.add_subsystem('ivc', om.IndepVarComp('x', val=np.ones(x_shape)))
-        ivc.add_output('y', val=np.ones(y_shape))
-        ivc.add_discrete_output('disc_out', val=3)
-        comp = p.model.add_subsystem('comp', DotProductMultDiscrete(derivs_method='jax'))
-        comp.matrix_free = matrix_free
-
-        p.model.connect('ivc.x', 'comp.x')
-        p.model.connect('ivc.y', 'comp.y')
-        p.model.connect('ivc.disc_out', 'comp.disc_in')
-
-        p.setup(mode=mode)
-
-        x = np.arange(1,np.prod(x_shape)+1).reshape(x_shape) * 2.0
-        y = np.arange(1,np.prod(y_shape)+1).reshape(y_shape)* 3.0
-        p.set_val('ivc.x', x)
-        p.set_val('ivc.y', y)
-        p.final_setup()
-        p.run_model()
-
-        assert_near_equal(p.get_val('comp.z'), np.dot(x, y))
-        assert_near_equal(p.get_val('comp.zz'), y * 3.0)
-        assert_check_totals(p.check_totals(of=['comp.z','comp.zz'], wrt=['comp.x', 'comp.y'],
-                                           method='fd', show_only_incorrect=True))
-        assert_check_partials(p.check_partials(show_only_incorrect=True))
-
-        p.set_val('ivc.disc_out', -2)
-        p.run_model()
-        assert_near_equal(p.get_val('comp.z'), -np.dot(x, y))
-        assert_near_equal(p.get_val('comp.zz'), y * 2.5)
+        assert_near_equal(p.get_val('comp.zz'), y * 1.9 * 2.)
         assert_check_totals(p.check_totals(of=['comp.z','comp.zz'], wrt=['comp.x', 'comp.y'],
                                            method='fd', show_only_incorrect=True))
         assert_check_partials(p.check_partials(show_only_incorrect=True))
@@ -398,7 +293,7 @@ class TestJaxComp(unittest.TestCase):
         ivc = p.model.add_subsystem('ivc', om.IndepVarComp('x', val=np.ones(x_shape)))
         ivc.add_output('y', val=np.ones(y_shape))
         ivc.add_discrete_output('disc_out', val=3)
-        comp = p.model.add_subsystem('comp', DotProductMultDiscretePrimal(derivs_method='jax'))
+        comp = p.model.add_subsystem('comp', DotProductMultDiscretePrimal())
         comp.matrix_free = matrix_free
 
         p.model.connect('ivc.x', 'comp.x')
@@ -431,7 +326,6 @@ class TestJaxComp(unittest.TestCase):
 
 if sys.version_info >= (3, 9):
 
-
     class CompRetValue(om.JaxExplicitComponent):
         def __init__(self, shape, nins=1, nouts=1, **kwargs):
             super().__init__(**kwargs)
diff --git a/openmdao/jax/tests/test_jax_implicit.py b/openmdao/jax/tests/test_jax_implicit.py
index 0a6dfff04..4f1fe0011 100644
--- a/openmdao/jax/tests/test_jax_implicit.py
+++ b/openmdao/jax/tests/test_jax_implicit.py
@@ -68,52 +68,6 @@ class JaxQuadraticCompPrimal(om.JaxImplicitComponent):
         return a * x ** 2 + b * x + c
 
 
-
-class JaxLinearSystemCompConverted(om.ImplicitComponent):
-
-    def initialize(self):
-        self.options.declare('size', default=1, types=int)
-
-    def setup(self):
-        size = self.options['size']
-
-        shape = (size, )
-
-        self.add_input("A", val=np.eye(size))
-        self.add_input("b", val=np.ones(shape))
-        self.add_output("x", shape=shape, val=.1)
-
-    def setup_partials(self):
-        size = self.options['size']
-        mat_size = size * size
-        full_size = size
-
-        row_col = np.arange(full_size, dtype="int")
-
-        self.declare_partials('x', 'b', val=np.full(full_size, -1.0), rows=row_col, cols=row_col)
-
-        rows = np.repeat(np.arange(full_size), size)
-
-        cols = np.arange(mat_size)
-
-        self.declare_partials('x', 'A', rows=rows, cols=cols)
-
-        cols = np.tile(np.arange(size), size)
-        cols += np.repeat(np.arange(1), mat_size) * size
-
-        self.declare_partials(of='x', wrt='x', rows=rows, cols=cols)
-
-        if self.matrix_free:
-            self.linear_solver = om.ScipyKrylov()
-        else:
-            self.linear_solver = om.DirectSolver()
-        self.nonlinear_solver = om.NewtonSolver(solve_subsystems=False)
-
-    def apply_nonlinear(self, inputs, outputs, residuals):
-        residuals['x'] = inputs['A'].dot(outputs['x']) - inputs['b']
-
-
-
 class JaxLinearSystemCompPrimal(om.JaxImplicitComponent):
 
     def initialize(self):
@@ -254,33 +208,6 @@ class JaxLinearSystemCompPrimalwDiscrete(om.JaxImplicitComponent):
 @unittest.skipIf(jax is None or sys.version_info < (3, 9), 'jax is not available or python < 3.9.')
 class TestJaxImplicitComp(unittest.TestCase):
 
-    @parameterized.expand(itertools.product(['fwd', 'rev'], ['matfree', '']), name_func=parameterized_name)
-    def test_quad_comp_converted(self, mode, matrix_free):
-        p = om.Problem()
-        # create an IVC manually so we can set the shapes.  Otherwise must set shape in the component
-        # itself.
-        ivc = p.model.add_subsystem('ivc', om.IndepVarComp('a'))
-        ivc.add_output('b')
-        ivc.add_output('c')
-        comp = p.model.add_subsystem('comp', QuadraticComp(derivs_method='jax'))
-        comp.matrix_free = bool(matrix_free)
-        p.model.connect('ivc.a', 'comp.a')
-        p.model.connect('ivc.b', 'comp.b')
-        p.model.connect('ivc.c', 'comp.c')
-        p.setup(mode=mode)
-
-        p.set_val('ivc.a', 1.0)
-        p.set_val('ivc.b', -4.0)
-        p.set_val('ivc.c', 3.0)
-
-        p.final_setup()
-        p.run_model()
-
-        assert_near_equal(p.get_val('comp.x'), 3.0)
-        assert_check_totals(p.check_totals(of=['comp.x'], wrt=['comp.a', 'comp.b', 'comp.c'], method='fd',
-                                           show_only_incorrect=True, abs_err_tol=3e-5, rel_err_tol=5e-6), atol=3e-5, rtol=5e-6)
-        assert_check_partials(p.check_partials(show_only_incorrect=True))
-
     @parameterized.expand(itertools.product(['fwd', 'rev'], ['matfree', ''], [(), (3,), (2,3)]), name_func=parameterized_name)
     def test_jax_quad_comp(self, mode, matrix_free, shape):
         p = om.Problem()
@@ -308,36 +235,6 @@ class TestJaxImplicitComp(unittest.TestCase):
                                            show_only_incorrect=True, abs_err_tol=3.5e-5, rel_err_tol=5e-6), atol=3.5e-5, rtol=5e-6)
         assert_check_partials(p.check_partials(show_only_incorrect=True, abs_err_tol=3e-6, rel_err_tol=3e-6), atol=2.5e-6)
 
-    @parameterized.expand(itertools.product(['fwd', 'rev'], ['matfree', '']), name_func=parameterized_name)
-    def test_lin_system_converted(self, mode, matrix_free):
-        A = np.array([[1., 1., 1.], [1., 2., 3.], [0., 1., 3.]])
-        b = np.array([1, 2, -3])
-
-        prob = om.Problem()
-
-        ivc = prob.model.add_subsystem('ivc', om.IndepVarComp())
-        ivc.add_output('A', A)
-        ivc.add_output('b', b)
-
-        lingrp = prob.model.add_subsystem('lingrp', om.Group())
-        lin = lingrp.add_subsystem('lin', JaxLinearSystemCompConverted(size=3, derivs_method='jax'))
-        lin.matrix_free = bool(matrix_free)
-
-        prob.model.connect('ivc.A', 'lingrp.lin.A')
-        prob.model.connect('ivc.b', 'lingrp.lin.b')
-
-        prob.setup(mode=mode)
-        prob.set_solver_print(level=0)
-
-        prob.run_model()
-
-        assert_near_equal(prob['lingrp.lin.x'], np.array([-4, 9, -4]), tolerance=1e-14)
-
-        assert_check_totals(prob.check_totals(of=['lingrp.lin.x'], wrt=['ivc.b', 'ivc.A'],
-                                              abs_err_tol=2e-4, rel_err_tol=3e-6, show_only_incorrect=True),
-                            atol=2e-4, rtol=3e-6)
-        assert_check_partials(prob.check_partials(show_only_incorrect=True), rtol=1e-5)
-
     @parameterized.expand(itertools.product(['fwd', 'rev'], ['matfree', '']), name_func=parameterized_name)
     def test_jax_lin_system_primal(self, mode, matrix_free):
         A = np.array([[1., 1., 1.], [1., 2., 3.], [0., 1., 3.]])