diff --git a/tests/include/J2FlowDirection.h b/include/neml2/models/solid_mechanics/AssociativeJ2FlowDirection.h
similarity index 90%
rename from tests/include/J2FlowDirection.h
rename to include/neml2/models/solid_mechanics/AssociativeJ2FlowDirection.h
index 15a8a590da..bd90bdbdc5 100644
--- a/tests/include/J2FlowDirection.h
+++ b/include/neml2/models/solid_mechanics/AssociativeJ2FlowDirection.h
@@ -28,13 +28,13 @@
 
 namespace neml2
 {
-/// The plastic flow direction assuming J2 flow.
-class J2FlowDirection : public Model
+/// The plastic flow direction assuming an associative J2 flow.
+class AssociativeJ2FlowDirection : public Model
 {
 public:
   static OptionSet expected_options();
 
-  J2FlowDirection(const OptionSet & options);
+  AssociativeJ2FlowDirection(const OptionSet & options);
 
 protected:
   virtual void set_value(bool, bool, bool) override;
diff --git a/include/neml2/models/solid_mechanics/LinearIsotropicElasticJ2TrialStressUpdate.h b/include/neml2/models/solid_mechanics/LinearIsotropicElasticJ2TrialStressUpdate.h
new file mode 100644
index 0000000000..08fb95f96d
--- /dev/null
+++ b/include/neml2/models/solid_mechanics/LinearIsotropicElasticJ2TrialStressUpdate.h
@@ -0,0 +1,67 @@
+// Copyright 2024, UChicago Argonne, LLC
+// All Rights Reserved
+// Software Name: NEML2 -- the New Engineering material Model Library, version 2
+// By: Argonne National Laboratory
+// OPEN SOURCE LICENSE (MIT)
+//
+// Permission is hereby granted, free of charge, to any person obtaining a copy
+// of this software and associated documentation files (the "Software"), to deal
+// in the Software without restriction, including without limitation the rights
+// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+// copies of the Software, and to permit persons to whom the Software is
+// furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in
+// all copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+// THE SOFTWARE.
+
+#pragma once
+
+#include "neml2/models/Model.h"
+
+namespace neml2
+{
+/**
+ * @brief Update the trial stress under the assumptions of J2 plasticity and isotropic linear
+ * elasticity
+ *
+ * This allows the construction of fully scalar return mapping models for
+ * isotropic materials.
+ */
+class LinearIsotropicElasticJ2TrialStressUpdate : public Model
+{
+public:
+  static OptionSet expected_options();
+
+  LinearIsotropicElasticJ2TrialStressUpdate(const OptionSet & options);
+
+protected:
+  /// compute updated trial stress
+  virtual void set_value(bool out, bool dout_din, bool d2out_din2) override;
+
+  /// input trial stress (i.e., assuming a purely elastic step)
+  const Variable<Scalar> & _elastic_trial_stress;
+
+  /// input inelastic strain
+  const Variable<Scalar> & _inelastic_strain;
+
+  /// input old inelastic strain
+  const Variable<Scalar> & _inelastic_strain_old;
+
+  /// output (updated) trial stress
+  Variable<Scalar> & _updated_trial_stress;
+
+  /// Young's modulus
+  const Scalar & _E;
+
+  /// Poisson's ratio
+  const Scalar & _nu;
+};
+} // namespace neml2
diff --git a/tests/src/J2FlowDirection.cxx b/src/neml2/models/solid_mechanics/AssociativeJ2FlowDirection.cxx
similarity index 72%
rename from tests/src/J2FlowDirection.cxx
rename to src/neml2/models/solid_mechanics/AssociativeJ2FlowDirection.cxx
index e1736b8c74..b9bb13f1ed 100644
--- a/tests/src/J2FlowDirection.cxx
+++ b/src/neml2/models/solid_mechanics/AssociativeJ2FlowDirection.cxx
@@ -22,23 +22,29 @@
 // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 // THE SOFTWARE.
 
-#include "J2FlowDirection.h"
+#include "neml2/models/solid_mechanics/AssociativeJ2FlowDirection.h"
 #include "neml2/tensors/SSR4.h"
 
 namespace neml2
 {
-register_NEML2_object(J2FlowDirection);
+register_NEML2_object(AssociativeJ2FlowDirection);
 
 OptionSet
-J2FlowDirection::expected_options()
+AssociativeJ2FlowDirection::expected_options()
 {
   auto options = Model::expected_options();
+  options.doc() = "The plastic flow direction assuming an associative J2 flow.";
+
   options.set<VariableName>("mandel_stress") = VariableName("state", "M");
+  options.set("mandel_stress").doc() = "Mandel stress";
+
   options.set<VariableName>("flow_direction") = VariableName("state", "NM");
+  options.set("flow_direction").doc() = "Flow direction";
+
   return options;
 }
 
-J2FlowDirection::J2FlowDirection(const OptionSet & options)
+AssociativeJ2FlowDirection::AssociativeJ2FlowDirection(const OptionSet & options)
   : Model(options),
     _M(declare_input_variable<SR2>("mandel_stress")),
     _N(declare_output_variable<SR2>("flow_direction"))
@@ -46,7 +52,7 @@ J2FlowDirection::J2FlowDirection(const OptionSet & options)
 }
 
 void
-J2FlowDirection::set_value(bool out, bool dout_din, bool d2out_din2)
+AssociativeJ2FlowDirection::set_value(bool out, bool dout_din, bool d2out_din2)
 {
   neml_assert_dbg(!d2out_din2, "Second derivatives not implemented");
 
@@ -60,10 +66,12 @@ J2FlowDirection::set_value(bool out, bool dout_din, bool d2out_din2)
   }
 
   if (dout_din)
-  {
-    auto I = SSR4::identity_sym(options());
-    auto J = SSR4::identity_dev(options());
-    _N.d(_M) = 3.0 / 2.0 * (I - 2.0 / 3.0 * dvm_dM.outer(dvm_dM)) * J / vm;
-  }
+    if (_M.is_dependent())
+    {
+      auto I = SSR4::identity_sym(options());
+      auto J = SSR4::identity_dev(options());
+
+      _N.d(_M) = 3.0 / 2.0 * (I - 2.0 / 3.0 * dvm_dM.outer(dvm_dM)) * J / vm;
+    }
 }
 } // namespace neml2
diff --git a/src/neml2/models/solid_mechanics/LinearIsotropicElasticJ2TrialStressUpdate.cxx b/src/neml2/models/solid_mechanics/LinearIsotropicElasticJ2TrialStressUpdate.cxx
new file mode 100644
index 0000000000..f35a413ffc
--- /dev/null
+++ b/src/neml2/models/solid_mechanics/LinearIsotropicElasticJ2TrialStressUpdate.cxx
@@ -0,0 +1,97 @@
+// Copyright 2024, UChicago Argonne, LLC
+// All Rights Reserved
+// Software Name: NEML2 -- the New Engineering material Model Library, version 2
+// By: Argonne National Laboratory
+// OPEN SOURCE LICENSE (MIT)
+//
+// Permission is hereby granted, free of charge, to any person obtaining a copy
+// of this software and associated documentation files (the "Software"), to deal
+// in the Software without restriction, including without limitation the rights
+// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+// copies of the Software, and to permit persons to whom the Software is
+// furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in
+// all copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+// THE SOFTWARE.
+
+#include "neml2/models/solid_mechanics/LinearIsotropicElasticJ2TrialStressUpdate.h"
+
+namespace neml2
+{
+register_NEML2_object(LinearIsotropicElasticJ2TrialStressUpdate);
+
+OptionSet
+LinearIsotropicElasticJ2TrialStressUpdate::expected_options()
+{
+  OptionSet options = Model::expected_options();
+
+  options.doc() = "Update the trial stress under the assumptions of J2 plasticity and isotropic "
+                  "linear elasticity";
+
+  options.set_input("elastic_trial_stress") = VariableName("forces", "s");
+  options.set("elastic_trial_stress").doc() = "Initial trial stress assuming a purely elastic step";
+
+  options.set_input("equivalent_plastic_strain") = VariableName("state", "ep");
+  options.set("equivalent_plastic_strain").doc() =
+      "Current guess for the equivalent plastic strain";
+
+  options.set_output("updated_trial_stress") = VariableName("state", "s");
+  options.set("updated_trial_stress").doc() =
+      "Trial stress corrected for the current increment of plastic deformation";
+
+  options.set_parameter<CrossRef<Scalar>>("youngs_modulus");
+  options.set("youngs_modulus").doc() = "Young's modulus";
+
+  options.set_parameter<CrossRef<Scalar>>("poisson_ratio");
+  options.set("poisson_ratio").doc() = "Poisson's ratio";
+
+  return options;
+}
+
+LinearIsotropicElasticJ2TrialStressUpdate::LinearIsotropicElasticJ2TrialStressUpdate(
+    const OptionSet & options)
+  : Model(options),
+    _elastic_trial_stress(declare_input_variable<Scalar>("elastic_trial_stress")),
+    _inelastic_strain(declare_input_variable<Scalar>("equivalent_plastic_strain")),
+    _inelastic_strain_old(declare_input_variable<Scalar>(_inelastic_strain.name().old())),
+    _updated_trial_stress(declare_output_variable<Scalar>("updated_trial_stress")),
+    _E(declare_parameter<Scalar>("E", "youngs_modulus", /*allow_nonlinear=*/false)),
+    _nu(declare_parameter<Scalar>("nu", "poisson_ratio", /*allow_nonlinear=*/false))
+{
+}
+
+void
+LinearIsotropicElasticJ2TrialStressUpdate::set_value(bool out, bool dout_din, bool d2out_din2)
+{
+  const auto three_shear = 3.0 * _E / (2.0 * (1.0 + _nu));
+
+  if (out)
+    _updated_trial_stress =
+        _elastic_trial_stress - three_shear * (_inelastic_strain - _inelastic_strain_old);
+
+  if (dout_din)
+  {
+    if (_elastic_trial_stress.is_dependent())
+      _updated_trial_stress.d(_elastic_trial_stress) = Scalar::identity_map(options());
+
+    if (_inelastic_strain.is_dependent())
+      _updated_trial_stress.d(_inelastic_strain) = -three_shear;
+
+    if (_inelastic_strain_old.is_dependent())
+      _updated_trial_stress.d(_inelastic_strain_old) = three_shear;
+  }
+
+  if (d2out_din2)
+  {
+    // zero
+  }
+}
+} // namespace neml2
diff --git a/tests/regression/solid_mechanics/viscoplasticity/misc/polynomial/model.i b/tests/regression/solid_mechanics/viscoplasticity/misc/polynomial/model.i
index ed95a6a04f..1568ca0849 100644
--- a/tests/regression/solid_mechanics/viscoplasticity/misc/polynomial/model.i
+++ b/tests/regression/solid_mechanics/viscoplasticity/misc/polynomial/model.i
@@ -145,7 +145,7 @@
     stress = 'forces/S'
   []
   [trial_flow_direction]
-    type = J2FlowDirection
+    type = AssociativeJ2FlowDirection
     mandel_stress = 'forces/S'
     flow_direction = 'forces/N'
   []
diff --git a/tests/regression/solid_mechanics/viscoplasticity/misc/torch_script/model.i b/tests/regression/solid_mechanics/viscoplasticity/misc/torch_script/model.i
index 4fe90de024..549964b35e 100644
--- a/tests/regression/solid_mechanics/viscoplasticity/misc/torch_script/model.i
+++ b/tests/regression/solid_mechanics/viscoplasticity/misc/torch_script/model.i
@@ -103,7 +103,7 @@ nbatch = 20
     stress = 'forces/S'
   []
   [trial_flow_direction]
-    type = J2FlowDirection
+    type = AssociativeJ2FlowDirection
     mandel_stress = 'forces/S'
     flow_direction = 'forces/N'
   []
diff --git a/tests/regression/solid_mechanics/viscoplasticity/misc/torch_script/model_scalar.i b/tests/regression/solid_mechanics/viscoplasticity/misc/torch_script/model_scalar.i
new file mode 100644
index 0000000000..a7e9ce9f03
--- /dev/null
+++ b/tests/regression/solid_mechanics/viscoplasticity/misc/torch_script/model_scalar.i
@@ -0,0 +1,214 @@
+ntime = 100
+nbatch = 20
+
+[Tensors]
+  [end_time]
+    type = LogspaceScalar
+    start = 0
+    end = 1
+    nstep = ${nbatch}
+  []
+  [times]
+    type = LinspaceScalar
+    start = 0
+    end = end_time
+    nstep = ${ntime}
+  []
+  [start_temperature]
+    type = LinspaceScalar
+    start = 300
+    end = 500
+    nstep = ${nbatch}
+  []
+  [end_temperature]
+    type = LinspaceScalar
+    start = 1800
+    end = 1200
+    nstep = ${nbatch}
+  []
+  [temperatures]
+    type = LinspaceScalar
+    start = start_temperature
+    end = end_temperature
+    nstep = ${ntime}
+  []
+  [exx]
+    type = FullScalar
+    batch_shape = '(${nbatch})'
+    value = 0.1
+  []
+  [eyy]
+    type = FullScalar
+    batch_shape = '(${nbatch})'
+    value = -0.05
+  []
+  [ezz]
+    type = FullScalar
+    batch_shape = '(${nbatch})'
+    value = -0.05
+  []
+  [max_strain]
+    type = FillSR2
+    values = 'exx eyy ezz'
+  []
+  [strains]
+    type = LinspaceSR2
+    start = 0
+    end = max_strain
+    nstep = ${ntime}
+  []
+[]
+
+[Drivers]
+  [driver]
+    type = SolidMechanicsDriver
+    model = 'model'
+    times = 'times'
+    prescribed_strains = 'strains'
+    prescribed_temperatures = 'temperatures'
+    predictor = LINEAR_EXTRAPOLATION
+    save_as = 'result.pt'
+    enable_AD = true
+  []
+  [regression]
+    type = TransientRegression
+    driver = 'driver'
+    reference = 'gold/result.pt'
+  []
+[]
+
+[Solvers]
+  [newton]
+    type = Newton
+    abs_tol = 1e-8
+    rel_tol = 1e-6
+  []
+[]
+
+[Models]
+  #####################################################################################
+  # Compute the invariant plastic flow direction since we are doing J2 radial return
+  #####################################################################################
+  [trial_elastic_strain]
+    type = SR2LinearCombination
+    to_var = 'forces/Ee'
+    from_var = 'forces/E old_state/Ep'
+    coefficients = '1 -1'
+  []
+  [trial_cauchy_stress]
+    type = LinearIsotropicElasticity
+    youngs_modulus = 1e5
+    poisson_ratio = 0.3
+    strain = 'forces/Ee'
+    stress = 'forces/S'
+  []
+  [trial_flow_direction]
+    type = AssociativeJ2FlowDirection
+    mandel_stress = 'forces/S'
+    flow_direction = 'forces/N'
+  []
+  [vonmises]
+    type = SR2Invariant
+    invariant_type = 'VONMISES'
+    tensor = 'forces/S'
+    invariant = 'forces/s'
+  []
+  [trial_state]
+    type = ComposedModel
+    models = 'trial_elastic_strain trial_cauchy_stress trial_flow_direction vonmises'
+  []
+
+  #####################################################################################
+  # Stress update
+  #####################################################################################
+  [ep_rate]
+    type = ScalarVariableRate
+    variable = 'state/ep'
+    rate = 'state/ep_rate'
+  []
+  [plastic_strain_rate]
+    type = AssociativePlasticFlow
+    flow_direction = 'forces/N'
+    flow_rate = 'state/ep_rate'
+    plastic_strain_rate = 'state/Ep_rate'
+  []
+  [plastic_strain]
+    type = SR2ForwardEulerTimeIntegration
+    variable = 'state/Ep'
+  []
+  [plastic_update]
+    type = ComposedModel
+    models = 'ep_rate plastic_strain_rate plastic_strain'
+  []
+  [elastic_strain]
+    type = SR2LinearCombination
+    from_var = 'forces/E state/Ep'
+    to_var = 'state/Ee'
+    coefficients = '1 -1'
+  []
+  [cauchy_stress]
+    type = LinearIsotropicElasticity
+    youngs_modulus = 1e5
+    poisson_ratio = 0.3
+    strain = 'state/Ee'
+    stress = 'state/S'
+  []
+  [stress_update]
+    type = ComposedModel
+    models = 'elastic_strain cauchy_stress'
+  []
+
+  #####################################################################################
+  # Compute the rates of equivalent plastic strain and internal variables
+  #####################################################################################
+  [trial_stress_update]
+    type = LinearIsotropicElasticJ2TrialStressUpdate
+    youngs_modulus = 1e5
+    poisson_ratio = 0.3
+    elastic_trial_stress = 'forces/s'
+    inelastic_strain = 'state/ep'
+    updated_trial_stress = 'state/s'
+  []
+  [rom]
+    type = TorchScriptFlowRate
+    von_mises_stress = 'state/s'
+    temperature = 'forces/T'
+    internal_state_1 = 'state/G'
+    internal_state_2 = 'state/C'
+    equivalent_plastic_strain_rate = 'state/ep_rate'
+    internal_state_1_rate = 'state/G_rate'
+    internal_state_2_rate = 'state/C_rate'
+    torch_script = 'gold/surrogate.pt'
+  []
+  [integrate_ep]
+    type = ScalarBackwardEulerTimeIntegration
+    variable = 'state/ep'
+  []
+  [integrate_G]
+    type = ScalarBackwardEulerTimeIntegration
+    variable = 'state/G'
+  []
+  [integrate_C]
+    type = ScalarBackwardEulerTimeIntegration
+    variable = 'state/C'
+  []
+  [rate]
+    type = ComposedModel
+    models = "trial_stress_update rom
+              integrate_ep integrate_G integrate_C"
+  []
+  [radial_return]
+    type = ImplicitUpdate
+    implicit_model = 'rate'
+    solver = 'newton'
+  []
+
+  #####################################################################################
+  # Put the models together
+  #####################################################################################
+  [model]
+    type = ComposedModel
+    models = 'trial_state radial_return plastic_update stress_update trial_stress_update rom'
+    additional_outputs = 'state/s state/ep_rate state/G_rate state/C_rate state/ep state/G state/C state/S state/Ep'
+  []
+[]
diff --git a/tests/unit/models/solid_mechanics/AssociativeJ2FlowDirection.i b/tests/unit/models/solid_mechanics/AssociativeJ2FlowDirection.i
new file mode 100644
index 0000000000..0c4a5050f0
--- /dev/null
+++ b/tests/unit/models/solid_mechanics/AssociativeJ2FlowDirection.i
@@ -0,0 +1,31 @@
+[Tensors]
+  [M]
+    type = FillSR2
+    values = '1 2 3 4 5 6'
+  []
+  [NM]
+    type = FillSR2
+    values = '-0.09805807 0 0.09805807 0.39223227 0.49029034 0.58834841'
+  []
+[]
+
+[Drivers]
+  [unit]
+    type = ModelUnitTest
+    model = 'model'
+    batch_shape = '(10)'
+    input_symr2_names = 'state/M'
+    input_symr2_values = 'M'
+    output_symr2_names = 'state/NM'
+    output_symr2_values = 'NM'
+    derivatives_abs_tol = 1e-6
+  []
+[]
+
+[Models]
+  [model]
+    type = AssociativeJ2FlowDirection
+    mandel_stress = 'state/M'
+    flow_direction = 'state/NM'
+  []
+[]
diff --git a/tests/unit/models/solid_mechanics/LinearIsotropicElasticJ2TrialStressUpdate.i b/tests/unit/models/solid_mechanics/LinearIsotropicElasticJ2TrialStressUpdate.i
new file mode 100644
index 0000000000..35e9ae88da
--- /dev/null
+++ b/tests/unit/models/solid_mechanics/LinearIsotropicElasticJ2TrialStressUpdate.i
@@ -0,0 +1,20 @@
+[Drivers]
+  [unit]
+    type = ModelUnitTest
+    model = 'model'
+    batch_shape = '(10)'
+    input_scalar_names = 'forces/s state/ep old_state/ep'
+    input_scalar_values = '1.5 0.1 0.05'
+    output_scalar_names = 'state/s'
+    output_scalar_values = '1.38461538462'
+    derivatives_abs_tol = 1e-6
+  []
+[]
+
+[Models]
+  [model]
+    type = LinearIsotropicElasticJ2TrialStressUpdate
+    youngs_modulus = 2
+    poisson_ratio = 0.3
+  []
+[]